Batch a stream lazily

Embedding APIs take inputs in batches. Bulk loaders write rows in batches. The input is often too big to hold in memory, or never ends at all. So batching has to be lazy.

Write a generator batched(iterable, n):

• Yield lists of n consecutive items from iterable.
• The last list can be shorter. Never yield an empty list.
• Be lazy: pull only as many items as the next batch needs. It must work on an infinite iterator.
• If n < 1, raise ValueError.

Don’t use itertools.batched. Build it yourself.

>>> list(batched("abcdefg", 3))
[['a', 'b', 'c'], ['d', 'e', 'f'], ['g']]
>>> from itertools import count, islice
>>> list(islice(batched(count(), 2), 3))
[[0, 1], [2, 3], [4, 5]]

Approach

Take one iterator and slice n items off it at a time.

from itertools import islice


def batched(iterable, n):
    if n < 1:
        raise ValueError("n must be at least 1")
    it = iter(iterable)
    while True:
        chunk = list(islice(it, n))
        if not chunk:
            return
        yield chunk

iter(iterable) matters. Without it, islice on a list would start from the front every time and loop forever. With it, each islice picks up where the last one stopped.

The ValueError only fires when the caller first asks for a batch, because a generator body doesn’t run until then. If you need it to fail at call time, split this into a plain function that checks n and returns an inner generator.

Complexity

O(N) over the whole stream. Memory is O(n): one batch at a time, however long the input is.

Why this matters

Pipelines are chains of generators. A reader yields rows, a batcher groups them, a writer sends each batch to an embedding model or a database. Memory stays flat because nothing upstream runs ahead of what downstream has asked for. That is backpressure, in its simplest form. Python 3.12 added itertools.batched (it yields tuples) for exactly this job.