Emmanuel Corels
← Blog
Automation & Scripting

Stream a large Python data job with yield instead of filling memory

Turn an all-at-once import into a lazy, testable pipeline while preserving cleanup, error context, batching and backpressure.

By Emmanuel Corels

A Python import that works on a 20 MB sample can collapse when production sends 20 GB. The usual failure is architectural: the program constructs a complete list before downstream work begins. A generator changes the contract from “return every result” to “produce the next result when requested.”

Find the eager allocation

def load_rows(path):
    rows = []
    with path.open(newline='', encoding='utf-8') as handle:
        for row in csv.DictReader(handle):
            rows.append(row)
    return rows

Memory grows with the number and size of records, plus Python object overhead. Returning the list also delays the first useful result until the entire file has been read.

Make the source lazy

from collections.abc import Iterator
from pathlib import Path
import csv

def stream_rows(path: Path) -> Iterator[dict[str, str]]:
    with path.open(newline='', encoding='utf-8') as handle:
        reader = csv.DictReader(handle)
        for line_number, row in enumerate(reader, start=2):
            row['_line_number'] = str(line_number)
            yield row

The presence of yield makes this a generator function. Calling it creates a generator object; it does not open and consume the file immediately. Each iteration resumes execution after the previous yield, with local state preserved.

rows = stream_rows(Path('/data/transactions.csv'))
print(rows)          # generator object
first = next(rows)   # opens and reads through the first data row

Build transformations as composable stages

from collections.abc import Iterable, Iterator
from decimal import Decimal, InvalidOperation

def valid_transactions(
    rows: Iterable[dict[str, str]]
) -> Iterator[dict[str, object]]:
    for row in rows:
        try:
            amount = Decimal(row['amount'])
        except (KeyError, InvalidOperation) as exc:
            raise ValueError(
                f"invalid amount at CSV line {row.get('_line_number')}"
            ) from exc
        if amount < 0:
            continue
        yield {**row, 'amount': amount}

Nothing is processed until a consumer requests items. The stages can be tested separately with small lists, yet remain lazy when connected to a file iterator.

Consume without accidentally materializing the stream

count = 0
total = Decimal('0')

for transaction in valid_transactions(
    stream_rows(Path('/data/transactions.csv'))
):
    write_transaction(transaction)
    count += 1
    total += transaction['amount']

print(f'processed={count} total={total}')

Calls such as list(generator), sorted(generator) and some forms of grouping consume the entire iterator and can restore the original memory problem. Document whether a function accepts an iterable once or requires a reusable collection.

Batch at the system boundary

One-row database writes may waste most of the performance gained through streaming. Batch a bounded number of records:

from itertools import islice
from collections.abc import Iterable, Iterator
from typing import TypeVar

T = TypeVar('T')

def batches(items: Iterable[T], size: int) -> Iterator[list[T]]:
    if size < 1:
        raise ValueError('size must be positive')
    iterator = iter(items)
    while batch := list(islice(iterator, size)):
        yield batch

for batch in batches(valid_transactions(stream_rows(path)), 1000):
    repository.insert_many(batch)

Memory now scales with batch size rather than total input. Select the size from measurement: database limits, transaction duration, retry cost and record size all matter.

Preserve resource cleanup when consumers stop early

rows = stream_rows(path)
try:
    for row in rows:
        if should_stop(row):
            break
        process(row)
finally:
    rows.close()

A generator’s with block closes the file when the generator is exhausted or closed. In normal short-lived loops, reference cleanup often happens quickly, but explicit closure is clearer for long-running services and early termination. If the generator owns sockets, cursors or temporary resources, test cancellation paths.

Know the one-pass rule

rows = stream_rows(path)
count = sum(1 for _ in rows)
total = sum(Decimal(row['amount']) for row in rows)  # now empty

A generator iterator is consumed once. Combine aggregates in one pass, recreate the source, or persist an intentional intermediate representation. itertools.tee can split an iterator but may buffer a large gap between consumers, so it is not a free duplicate.

Add observability without destroying laziness

from time import monotonic

started = monotonic()
for count, row in enumerate(pipeline, start=1):
    write_transaction(row)
    if count % 100_000 == 0:
        elapsed = monotonic() - started
        logger.info(
            'import_progress',
            extra={'rows': count, 'rows_per_second': count / elapsed},
        )

Record input identity, checkpoint or line number, processed and rejected counts, throughput and batch retry state. A generator improves memory behavior; it does not automatically provide restartability or exactly-once delivery.

Measure the result

/usr/bin/time -l python import_transactions.py

On macOS, the command reports maximum resident set size among other metrics. Compare the same dataset and downstream behavior before and after. Also measure total runtime, time to first processed record and correctness totals; lower memory with incorrect or silently skipped data is not a win.

SymptomLikely causeCorrection
Memory still growsDownstream list, sort, cache or lagging tee consumerProfile allocations and inspect every stage
File remains openConsumer exits early without closingClose generator or wrap ownership in a context manager
Second loop is emptyIterator already consumedCombine work or recreate source
Throughput is poorPer-record network/database callIntroduce bounded batching

yield does not make a job fast by itself. It makes production incremental, which gives you control over memory, batching, cancellation and backpressure.

Sources: the canonical Stack Overflow explanation (CC BY-SA), Python’s yield-expression reference, the official generator HOWTO, and itertools documentation.

Primary source: Review the official reference ↗