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.
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.
| Symptom | Likely cause | Correction |
|---|---|---|
| Memory still grows | Downstream list, sort, cache or lagging tee consumer | Profile allocations and inspect every stage |
| File remains open | Consumer exits early without closing | Close generator or wrap ownership in a context manager |
| Second loop is empty | Iterator already consumed | Combine work or recreate source |
| Throughput is poor | Per-record network/database call | Introduce bounded batching |
yielddoes 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 ↗