Split a generator into chunks without pre-walking it
PythonGeneratorPython Problem Overview
(This question is related to this one and this one, but those are pre-walking the generator, which is exactly what I want to avoid)
I would like to split a generator in chunks. The requirements are:
- do not pad the chunks: if the number of remaining elements is less than the chunk size, the last chunk must be smaller.
- do not walk the generator beforehand: computing the elements is expensive, and it must only be done by the consuming function, not by the chunker
- which means, of course: do not accumulate in memory (no lists)
I have tried the following code:
def head(iterable, max=10):
for cnt, el in enumerate(iterable):
yield el
if cnt >= max:
break
def chunks(iterable, size=10):
i = iter(iterable)
while True:
yield head(i, size)
# Sample generator: the real data is much more complex, and expensive to compute
els = xrange(7)
for n, chunk in enumerate(chunks(els, 3)):
for el in chunk:
print 'Chunk %3d, value %d' % (n, el)
And this somehow works:
Chunk 0, value 0
Chunk 0, value 1
Chunk 0, value 2
Chunk 1, value 3
Chunk 1, value 4
Chunk 1, value 5
Chunk 2, value 6
^CTraceback (most recent call last):
File "xxxx.py", line 15, in <module>
for el in chunk:
File "xxxx.py", line 2, in head
for cnt, el in enumerate(iterable):
KeyboardInterrupt
Buuuut ... it never stops (I have to press ^C) because of the while True. I would like to stop that loop whenever the generator has been consumed, but I do not know how to detect that situation. I have tried raising an Exception:
class NoMoreData(Exception):
pass
def head(iterable, max=10):
for cnt, el in enumerate(iterable):
yield el
if cnt >= max:
break
if cnt == 0 : raise NoMoreData()
def chunks(iterable, size=10):
i = iter(iterable)
while True:
try:
yield head(i, size)
except NoMoreData:
break
# Sample generator: the real data is much more complex, and expensive to compute
els = xrange(7)
for n, chunk in enumerate(chunks(els, 2)):
for el in chunk:
print 'Chunk %3d, value %d' % (n, el)
But then the exception is only raised in the context of the consumer, which is not what I want (I want to keep the consumer code clean)
Chunk 0, value 0
Chunk 0, value 1
Chunk 0, value 2
Chunk 1, value 3
Chunk 1, value 4
Chunk 1, value 5
Chunk 2, value 6
Traceback (most recent call last):
File "xxxx.py", line 22, in <module>
for el in chunk:
File "xxxx.py", line 9, in head
if cnt == 0 : raise NoMoreData
__main__.NoMoreData()
How can I detect that the generator is exhausted in the chunks function, without walking it?
Python Solutions
Solution 1 - Python
One way would be to peek at the first element, if any, and then create and return the actual generator.
def head(iterable, max=10):
first = next(iterable) # raise exception when depleted
def head_inner():
yield first # yield the extracted first element
for cnt, el in enumerate(iterable):
yield el
if cnt + 1 >= max: # cnt + 1 to include first
break
return head_inner()
Just use this in your chunk generator and catch the StopIteration exception like you did with your custom exception.
Update: Here's another version, using itertools.islice to replace most of the head function, and a for loop. This simple for loop in fact does exactly the same thing as that unwieldy while-try-next-except-break construct in the original code, so the result is much more readable.
def chunks(iterable, size=10):
iterator = iter(iterable)
for first in iterator: # stops when iterator is depleted
def chunk(): # construct generator for next chunk
yield first # yield element from for loop
for more in islice(iterator, size - 1):
yield more # yield more elements from the iterator
yield chunk() # in outer generator, yield next chunk
And we can get even shorter than that, using itertools.chain to replace the inner generator:
def chunks(iterable, size=10):
iterator = iter(iterable)
for first in iterator:
yield chain([first], islice(iterator, size - 1))
Solution 2 - Python
Another way to create groups/chunks and not prewalk the generator is using itertools.groupby on a key function that uses an itertools.count object. Since the count object is independent of the iterable, the chunks can be easily generated without any knowledge of what the iterable holds.
Every iteration of groupby calls the next method of the count object and generates a group/chunk key (followed by items in the chunk) by doing an integer division of the current count value by the size of the chunk.
from itertools import groupby, count
def chunks(iterable, size=10):
c = count()
for _, g in groupby(iterable, lambda _: next(c)//size):
yield g
Each group/chunk g yielded by the generator function is an iterator. However, since groupby uses a shared iterator for all groups, the group iterators cannot be stored in a list or any container, each group iterator should be consumed before the next.
Solution 3 - Python
Even faster solution (new as of 2021-12-02) for smaller n:
When the chunk size is typically small, the fastest solution is this one, adapted from rhettg's answer:
from itertools import takewhile, zip_longest
def chunker(n, iterable):
'''chunker(3, 'ABCDEFG') --> ('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)'''
fillvalue = object() # Anonymous sentinel object that can't possibly appear in input
args = (iter(iterable),) * n
for x in zip_longest(*args, fillvalue=fillvalue):
if x[-1] is fillvalue:
# takewhile optimizes a bit for when n is large and the final
# group is small; at the cost of a little performance, you can
# avoid the takewhile import and simplify to:
# yield tuple(v for v in x if v is not fillvalue)
yield tuple(takewhile(lambda v: v is not fillvalue, x))
else:
yield x
Old answer (still fast, but loses to above by a little in basically all cases, and by a roughly factor of 2x in common cases):
Fastest possible solution I could come up with, thanks to (in CPython) using purely C-level builtins. By doing so, no Python byte code is needed to produce each chunk (unless the underlying generator is implemented in Python) which has a huge performance benefit. It does walk each chunk before returning it, but it doesn't do any pre-walking beyond the chunk it's about to return:
# Py2 only to get generator based map
from future_builtins import map
from itertools import islice, repeat, starmap, takewhile
# operator.truth is *significantly* faster than bool for the case of
# exactly one positional argument prior to 3.10; in 3.10+, you can
# just use bool (which is trivially faster than truth)
from operator import truth
def chunker(n, iterable): # n is size of each chunk; last chunk may be smaller
return takewhile(truth, map(tuple, starmap(islice, repeat((iter(iterable), n)))))
Since that's a bit dense, the spread out version for illustration:
def chunker(n, iterable):
iterable = iter(iterable)
while True:
x = tuple(islice(iterable, n))
if not x:
return
yield x
Wrapping a call to chunker in enumerate would let you number the chunks if it's needed.
Solution 4 - Python
more-itertools has provided chunked and ichunked that can achieve the goal, it is mentioned on the Python 3 itertools document page.
Solution 5 - Python
How about using itertools.islice:
import itertools
els = iter(xrange(7))
print list(itertools.islice(els, 2))
print list(itertools.islice(els, 2))
print list(itertools.islice(els, 2))
print list(itertools.islice(els, 2))
Which gives:
[0, 1]
[2, 3]
[4, 5]
[6]
Solution 6 - Python
Started to realize the usefulness of this scenario when crafting a solution to DB insertion of 500k+ rows at higher speed.
A generator processes the data from source and "yield"s it line by line; and then another generator groups the output in chunks and "yield "s it chunk by chunk. The second generator is only aware of the chunk size and nothing more.
Below is a sample to highlight the concept:
#!/usr/bin/python
def firstn_gen(n):
num = 0
while num < n:
yield num
num += 1
def chunk_gen(some_gen, chunk_size=7):
res_chunk = []
for count, item in enumerate(some_gen, 1):
res_chunk.append(item)
if count % chunk_size == 0:
yield res_chunk
res_chunk[:] = []
else:
yield res_chunk
if __name__ == '__main__':
for a_chunk in chunk_gen(firstn_gen(33)):
print(a_chunk)
Tested in Python 2.7.12:
[0, 1, 2, 3, 4, 5, 6]
[7, 8, 9, 10, 11, 12, 13]
[14, 15, 16, 17, 18, 19, 20]
[21, 22, 23, 24, 25, 26, 27]
[28, 29, 30, 31, 32]
Solution 7 - Python
I had this same issue, but found a simpler solution than those mentioned here:
def chunker(iterable, chunk_size):
els = iter(iterable)
while True:
next_el = next(els)
yield chain([next_el], islice(els, chunk_size - 1))
for i, chunk in enumerate(chunker(range(11), 2)):
for el in chunk:
print(i, el)
# Prints the following:
0 0
0 1
1 2
1 3
2 4
2 5
3 6
3 7
4 8
4 9
5 10
Solution 8 - Python
from itertools import islice
def chunk(it, n):
'''
# returns chunks of n elements each
>>> list(chunk(range(10), 3))
[
[0, 1, 2, ],
[3, 4, 5, ],
[6, 7, 8, ],
[9, ]
]
>>> list(chunk(list(range(10)), 3))
[
[0, 1, 2, ],
[3, 4, 5, ],
[6, 7, 8, ],
[9, ]
]
'''
def _w(g):
return lambda: tuple(islice(g, n))
return iter(_w(iter(it)), ())
Solution 9 - Python
Inspired by Moses Koledoye's answer, I tried to make a solution that uses itertools.groupby but doesn't require a divide at each step.
The following function can be used as a key to groupby, and it simply returns a boolean, which flips after a pre-defined number of calls.
def chunks(chunksize=3):
def flag_gen():
flag = False
while True:
for num in range(chunksize):
yield flag
flag = not flag
flag_iter = flag_gen()
def flag_func(*args, **kwargs):
return next(flag_iter)
return flag_func
Which can be used like this:
from itertools import groupby
my_long_generator = iter("abcdefghijklmnopqrstuvwxyz")
chunked_generator = groupby(my_long_generator, key=chunks(chunksize=5))
for flag, chunk in chunked_generator:
print("Flag is {f}".format(f=flag), list(chunk))
Output:
Flag is False ['a', 'b', 'c', 'd', 'e']
Flag is True ['f', 'g', 'h', 'i', 'j']
Flag is False ['k', 'l', 'm', 'n', 'o']
Flag is True ['p', 'q', 'r', 's', 't']
Flag is False ['u', 'v', 'w', 'x', 'y']
Flag is True ['z']
I've made a fiddle demonstrating this code.
Solution 10 - Python
You've said you don't wish to store things in memory, so does this mean that you can't build an intermediate list for the current chunk?
Why not traverse the generator and insert a sentinel value between chunks? The consumer (or a suitable wrapper) could ignore the sentinel:
class Sentinel(object):
pass
def chunk(els, size):
for i, el in enumerate(els):
yield el
if i > 0 and i % size == 0:
yield Sentinel
Solution 11 - Python
EDIT other solution with a generator of generators
You should not do a while True in your iterator, but simply iterate through it and update the chunk number at each iteration :
def chunk(it, maxv):
n = 0
for i in it:
yield n // mavx, i
n += 1
If you want a generator of generators, you can have :
def chunk(a, maxv):
def inner(it, maxv, l):
l[0] = False
for i in range(maxv):
yield next(it)
l[0] = True
raise StopIteration
it = iter(a)
l = [True]
while l[0] == True:
yield inner(it, maxv, l)
raise StopIteration
with a being an iterable.
Tests : on python 2.7 and 3.4:
for i in chunk(range(7), 3):
print 'CHUNK'
for a in i:
print a
gives :
CHUNK
0
1
2
CHUNK
3
4
5
CHUNK
6
And on 2.7 :
for i in chunk(xrange(7), 3):
print 'CHUNK'
for a in i:
print a
gives same result.
But BEWARE : list(chunk(range(7)) blocks on 2.7 and 3.4