CODEWITHSUNDEEP
pythongenerator
blueFast
blueFast

Reputation: 44381

Length of a finite generator

I have these two implementations to compute the length of a finite generator, while keeping the data for further processing:

def count_generator1(generator):
    '''- build a list with the generator data
       - get the length of the data
       - return both the length and the original data (in a list)
       WARNING: the memory use is unbounded, and infinite generators will block this'''
    l = list(generator)
    return len(l), l

def count_generator2(generator):
    '''- get two generators from the original generator
       - get the length of the data from one of them
       - return both the length and the original data, as returned by tee
       WARNING: tee can use up an unbounded amount of memory, and infinite generators will block this'''
    for_length, saved  = itertools.tee(generator, 2)
    return sum(1 for _ in for_length), saved

Both have drawbacks, both do the job. Could somebody comment on them, or even offer a better alternative?

Upvotes: 10

Views: 3609

Answers (3)

Vajk Hermecz
Vajk Hermecz

Reputation: 5712

If you don't need the length of the iterator prior to processing the data, you could use a helper method with a future to add in counting into the processing of your iterator/stream:

import asyncio
def ilen(iter):
    """
    Get future with length of iterator
    The future will hold the length once the iteartor is exhausted
    @returns: <iter, cnt-future>
    """
    def ilen_inner(iter, future):
        cnt = 0
        for row in iter:
            cnt += 1
            yield row
        future.set_result(cnt)
    cnt_future = asyncio.Future()
    return ilen_inner(iter, cnt_future), cnt_future

Usage would be:

data = db_connection.execute(query)
data, cnt = ilen(data)
solve_world_hunger(data)
print(f"Processed {cnt.result()} items")

Upvotes: 0

John Crawford
John Crawford

Reputation: 2344

I ran Windows 64-bit Python 3.4.3 timeit on a few approaches I could think of:

>>> from timeit import timeit
>>> from textwrap import dedent as d
>>> timeit(
...     d("""
...     count = -1
...     for _ in s:
...         count += 1
...     count += 1
...     """),
...     "s = range(1000)",
... )
50.70772041983173
>>> timeit(
...     d("""
...     count = -1
...     for count, _ in enumerate(s):
...         pass
...     count += 1
...     """),
...     "s = range(1000)",
... )
42.636973504498656
>>> timeit(
...     d("""
...     count, _ = reduce(f, enumerate(range(1000)), (-1, -1))
...     count += 1
...     """),
...     d("""
...     from functools import reduce
...     def f(_, count):
...         return count
...     s = range(1000)
...     """),
... )
121.15513102540672
>>> timeit("count = sum(1 for _ in s)", "s = range(1000)")
58.179126025925825
>>> timeit("count = len(tuple(s))", "s = range(1000)")
19.777029680237774
>>> timeit("count = len(list(s))", "s = range(1000)")
18.145157531932
>>> timeit("count = len(list(1 for _ in s))", "s = range(1000)")
57.41422175998332

Shockingly, the fastest approach was to use a list (not even a tuple) to exhaust the iterator and get the length from there:

>>> timeit("count = len(list(s))", "s = range(1000)")
18.145157531932

Of course, this risks memory issues. The best low-memory alternative was to use enumerate on a NOOP for-loop:

>>> timeit(
...     d("""
...     count = -1
...     for count, _ in enumerate(s):
...         pass
...     count += 1
...     """),
...     "s = range(1000)",
... )
42.636973504498656

Cheers!

Upvotes: 4

Gareth Latty
Gareth Latty

Reputation: 89017

If you have to do this, the first method is much better - as you consume all the values, itertools.tee() will have to store all the values anyway, meaning a list will be more efficient.

To quote from the docs:

This itertool may require significant auxiliary storage (depending on how much temporary data needs to be stored). In general, if one iterator uses most or all of the data before another iterator starts, it is faster to use list() instead of tee().

Upvotes: 13

Related Questions