Generators and working with large data

A list holds everything at once, which is fine until the everything is a ten-gigabyte file. A generator hands you one item at a time and remembers where it got to, so you can work through data far larger than memory, and start handling a model's reply before it has finished arriving. This section covers what an iterator is, how to write a generator, how to batch, and the async version used for streaming responses.

The idea#

A normal function computes everything and returns it all at once. A generator produces one value at a time and pauses in between, remembering where it was. Nothing is computed until you ask for it.

# Returns a list of 1,000,000 numbers. All in memory at once.
def squares_list(n):
    return [i * i for i in range(n)]

# Produces them one at a time. Almost no memory.
def squares_gen(n):
    for i in range(n):
        yield i * i

for value in squares_gen(1_000_000):
    print(value)

Why this matters#

Reading a very large file with read() loads all of it into memory and may fail. Looping over the file object reads one line at a time, so the size of the file stops mattering.

# Loads the whole file
with open("huge.jsonl") as f:
    lines = f.readlines()

# Reads one line at a time
def records(path):
    with open(path) as f:
        for line in f:
            yield json.loads(line)

for record in records("huge.jsonl"):
    process(record)

Generator expressions#

A comprehension with round brackets instead of square ones is a generator. Use it when you are going to loop over the result once and do not need to keep it.

squares = [n * n for n in range(1000)]     # a list, held in memory
squares = (n * n for n in range(1000))     # a generator, produced on demand

total = sum(n * n for n in range(1000))    # no list is ever built

Batching#

Most embedding and inference APIs take batches, so you will regularly need to turn a long stream into fixed-size groups.

from itertools import batched

for chunk in batched(records("huge.jsonl"), 64):
    embed(chunk)        # a tuple of up to 64 items

# Before Python 3.12, written by hand
def batched_manual(items, size):
    batch = []
    for item in items:
        batch.append(item)
        if len(batch) == size:
            yield batch
            batch = []
    if batch:
        yield batch      # do not lose the last partial batch

Streaming responses#

When an API streams a reply, it arrives as an async generator. You loop with async for, and each pass waits for the next piece.

async def collect(stream):
    parts = []
    async for chunk in stream:
        parts.append(chunk.text)
        print(chunk.text, end="", flush=True)   # show it as it arrives
    return "".join(parts)

One consequence worth remembering: you cannot check the shape of a streamed response until it has finished. Any JSON parsing or schema validation has to happen after the loop, not inside it.