Fast iteration over STL container in cython

Question

I store some data in an unordered_map in cython and would also like to expose the data via the python __iter__ interface. But iteration is significantly slower than the standard dict implementation in python. Can this be sped up?

The data structure looks something like this.

# distutils: language = c++

from cython.operator cimport dereference, preincrement
from libcpp.unordered_map cimport unordered_map
from libcpp.unordered_set cimport unordered_set

ctypedef unordered_map[int, unordered_set[int]] container_t


cdef class Base:
    """
    Base class holding the data.
    """
    cdef container_t container

    def __init__(self, int n):
        # Populate the container ([] creates an entry).
        for i in range(n):
            self.container[i]

    def iterate(self):
        """
        Iterate in C++ without python interaction as a baseline.
        """
        it = self.container.begin()
        while it != self.container.end():
            preincrement(it)

I've tried two different implementations of __iter__ to support python iteration: one with yield and the other using a dedicated iterator (see below).

cdef class A(Base):
    """
    Simple implementation using `yield`.
    """
    def __iter__(self):
        it = self.container.begin()
        while it != self.container.end():
            yield dereference(it).first
            preincrement(it)
    
    
cdef class Iterator:
    """
    Dedicated iterator.
    """
    cdef Base container
    # Using container_t.iterator doesn't seem to work.
    cdef unordered_map[int, unordered_set[int]].iterator it
    
    def __init__(self, Base container):
        self.container = container
        self.it = container.container.begin()
    
    def __next__(self):
        if self.it == self.container.container.end():
            raise StopIteration
        value = dereference(self.it).first
        preincrement(self.it)
        return value
    
    
cdef class B(Base):
    """
    Implementation using a dedicated iterator.
    """
    def __iter__(self):
        # Updated interface following @DavidW's comment.
        return Iterator(self)

So let's compare the runtimes.

# Number of elements in the container.
n = 1_000_000

%%timeit x = {i: set() for i in range(n)}
# Baseline python implementation.
list(x)
# 5.64 ms ± 14 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit x = A(n)
list(x)
# 19.2 ms ± 141 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit x = B(n)
list(x)
14.6 ms ± 61.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Using a dedicated iterator is faster than using yield but still about 3x slower than the standard dict implementation. Iteration over the STL container is faster, however.

%%timeit x = Base(n)
x.iterate()
# 1.75 ms ± 34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Is there a way to get comparable performance with dict (or better)? Or is this just a fundamental limitation due to hopping back and forth between python ints and C++ ints? For reference, the full code is here.

Fast iteration over STL container in cython

0 Answers