	In Browser
	StumbleUpon
	del.icio.us
	Google
	Google Buzz
	reddit
	LinkedIn

	Facebook
	Twitter
	Linkedin
	E-Mail

Modern Python > Type Hints and Dataclasses > Python Generators and itertools

Python Generators and itertools

Author: Venkata Sudhakar

A Python generator is a function that uses the yield keyword to produce values one at a time, pausing execution between each value. Unlike a regular function that returns all results in a list, a generator computes each value only when asked and holds only one value in memory at a time. This makes generators ideal for processing large datasets - a generator that reads a 10 GB file yields one line at a time without loading the entire file into memory. In data migration, generators are essential for processing millions of rows in batches without running out of RAM.

The itertools module provides fast, memory-efficient tools for working with iterators. itertools.islice() takes the first N elements from any iterator. itertools.chain() flattens multiple iterables into one. itertools.groupby() groups consecutive elements by key. itertools.takewhile() and dropwhile() take or skip elements while a condition holds. These tools compose cleanly with generators to build efficient data processing pipelines that handle data larger than memory.

The below example shows generators and itertools applied to data migration scenarios: streaming rows from a database, processing in batches, and building a complete extract-transform-load pipeline that never loads all data into memory at once.

import itertools
from typing import Generator, Iterator

# Generator: stream rows from database one batch at a time
def stream_rows(table: str, batch_size: int = 1000) -> Generator[dict, None, None]:
    offset = 0
    total = 5000  # simulate 5000-row table
    while offset < total:
        # In production: cursor.execute("SELECT * FROM ... LIMIT %s OFFSET %s")
        batch = [
            {"id": offset + i, "name": f"Customer_{offset+i}", "status": "ACTIVE"}
            for i in range(min(batch_size, total - offset))
        ]
        for row in batch:
            yield row  # yield one row at a time - O(1) memory
        offset += batch_size

# Transform generator: filter + map without building intermediate list
def transform_rows(rows: Iterator[dict]) -> Generator[dict, None, None]:
    for row in rows:
        if row["status"] == "ACTIVE":
            yield {
                "customer_id": row["id"],
                "full_name": row["name"].upper(),
                "migrated_at": "2024-01-15"
            }

# Helper: batch any iterator into chunks of size n
def batched(iterable, n):
    it = iter(iterable)
    while True:
        batch = list(itertools.islice(it, n))
        if not batch:
            return
        yield batch

# Full pipeline - each stage is lazy, only one batch in memory at a time
raw = stream_rows("customers", batch_size=1000)
transformed = transform_rows(raw)

total = 0
for batch in batched(transformed, 500):
    print(f"Inserting {len(batch)} rows | first_id={batch[0]['customer_id']}")
    total += len(batch)

print(f"Done. Total migrated: {total}")

It gives the following output,

Inserting 500 rows | first_id=0
Inserting 500 rows | first_id=500
Inserting 500 rows | first_id=1000
Inserting 500 rows | first_id=1500
Inserting 500 rows | first_id=2000
Inserting 500 rows | first_id=2500
Inserting 500 rows | first_id=3000
Inserting 500 rows | first_id=3500
Inserting 500 rows | first_id=4000
Inserting 500 rows | first_id=4500
Done. Total migrated: 5000
# Peak memory: ~500 rows at any time regardless of table size

import itertools

# itertools.chain: merge multiple table streams as one
customers = ({"table": "customers", "id": i} for i in range(3))
orders    = ({"table": "orders",    "id": i} for i in range(2))
all_rows  = itertools.chain(customers, orders)
for row in all_rows:
    print(f"{row['table']}:{row['id']}")

print()

# itertools.groupby: group CDC events by operation type
events = sorted([
    {"op": "c", "id": 1001}, {"op": "c", "id": 1002},
    {"op": "u", "id": 999},  {"op": "d", "id": 998}
], key=lambda e: e["op"])

for op, group in itertools.groupby(events, key=lambda e: e["op"]):
    rows = list(group)
    label = {"c": "INSERT", "u": "UPDATE", "d": "DELETE"}.get(op, op)
    print(f"{label}: {len(rows)} rows - IDs: {[r['id'] for r in rows]}")

print()

# itertools.takewhile: process CDC events until lag exceeds threshold
lag_readings = [0, 2, 5, 45, 230, 450, 50, 10, 0]
stable = list(itertools.takewhile(lambda lag: lag < 100, lag_readings))
print(f"Stable readings before lag spike: {stable}")

# Generator expression in sum - no intermediate list
amounts = (r * 9.99 for r in range(1, 1001))
total = sum(amounts)  # sums 1000 values with O(1) memory
print(f"Total amount: ${total:.2f}")

It gives the following output,

customers:0
customers:1
customers:2
orders:0
orders:1

DELETE: 1 rows - IDs: [998]
INSERT: 2 rows - IDs: [1001, 1002]
UPDATE: 1 rows - IDs: [999]

Stable readings before lag spike: [0, 2, 5, 45]
Total amount: $5004.99

Generator expressions vs list comprehensions:

Write [x*2 for x in data] when you need random access or will iterate multiple times. Write (x*2 for x in data) when you will iterate once and the dataset is large. The built-ins sum(), min(), max(), any(), all(), and sorted() all accept generators directly. sum(row["amount"] for row in million_rows) sums without building a list. sorted() is the one exception - it must materialise the generator to sort, so use it only when the dataset fits in memory.

Send your comments, suggestions or queries regarding this site to [email protected].