方寸 Portal

AI Engineer Portal

Your personal operating system for career transition.

Private mode
Back to Learning Paths
Lesson

Data pipelines for AI ingestion

Processing documents for ingestion, streaming data transformations, generators for memory efficiency, and batch vs stream processing tradeoffs.

65 min
python-for-ai-engineersphase-1portfolio

Why this matters

AI systems consume data at unusual scales: millions of tokens per ingestion run, thousands of documents per embedding batch, streaming responses that need to be parsed and forwarded in real time. Naive approaches — loading everything into memory, processing records one by one in a tight loop, materializing intermediate lists everywhere — work fine in notebooks and fail quietly in production.

The patterns in this lesson let you build pipelines that are memory-efficient, composable, and debuggable.

Core concepts

Generators for memory efficiency

A list comprehension materializes all results in memory at once. A generator yields one item at a time, consuming constant memory regardless of input size.

# This loads all 100K documents into memory simultaneously
all_chunks = [chunk for doc in documents for chunk in split_doc(doc)]

# This yields one chunk at a time, using O(1) memory
def stream_chunks(documents):
    for doc in documents:
        yield from split_doc(doc)

For large document collections, the generator version lets you process a 10GB JSONL file on a machine with 2GB of RAM.

Generator pipelines with composition

Generators compose cleanly into pipeline stages:

from pathlib import Path
from typing import Iterator
import json


def read_jsonl(path: Path) -> Iterator[dict]:
    with path.open() as f:
        for line in f:
            line = line.strip()
            if line:
                yield json.loads(line)


def filter_valid(records: Iterator[dict]) -> Iterator[dict]:
    for record in records:
        if record.get("body") and len(record["body"]) >= 100:
            yield record


def normalize(records: Iterator[dict]) -> Iterator[dict]:
    for record in records:
        yield {
            "id": record["id"],
            "text": record["body"].strip(),
            "metadata": {"source": record.get("source", "unknown")},
        }


def write_jsonl(records: Iterator[dict], path: Path) -> int:
    count = 0
    with path.open("w") as f:
        for record in records:
            f.write(json.dumps(record) + "\n")
            count += 1
    return count


# Compose into a pipeline:
pipeline = normalize(filter_valid(read_jsonl(Path("docs.jsonl"))))
total = write_jsonl(pipeline, Path("docs-clean.jsonl"))
print(f"Wrote {total} records")

Each stage is a generator that lazily pulls from the previous one. Memory stays constant. Adding a new transformation means inserting one function in the chain.

Batch processing for embedding

Embedding APIs accept batches. Processing one document at a time is ~100x slower than batching:

from typing import Iterator


def batch(items: Iterator[dict], size: int) -> Iterator[list[dict]]:
    batch = []
    for item in items:
        batch.append(item)
        if len(batch) >= size:
            yield batch
            batch = []
    if batch:
        yield batch


async def embed_pipeline(records: Iterator[dict], client, batch_size: int = 100) -> Iterator[dict]:
    for chunk in batch(records, batch_size):
        texts = [r["text"] for r in chunk]
        embeddings = await client.embed(texts)  # one API call for the batch
        for record, embedding in zip(chunk, embeddings):
            yield {**record, "embedding": embedding}

For 100K documents with average 256 tokens each, processing in batches of 100 means 1,000 API calls instead of 100,000.

Streaming document processing

When an LLM generates long outputs, streaming lets you forward tokens to the user as they arrive rather than waiting for the full response:

import asyncio


async def stream_to_output(response_stream, output_queue: asyncio.Queue) -> str:
    full_text = ""
    async for chunk in response_stream:
        token = chunk.choices[0].delta.content or ""
        full_text += token
        await output_queue.put(token)  # forward to consumer
    await output_queue.put(None)  # sentinel: stream ended
    return full_text

The producer streams tokens to a queue. The consumer (a WebSocket sender, a file writer, a downstream processor) drains the queue independently.

Batch vs stream: when to use each

Use batch processing when:

  • you control both ends of the pipeline (offline ingestion, evaluation runs)
  • throughput matters more than latency
  • partial results are not useful until the full batch completes

Use streaming when:

  • you need to show partial results to a user (LLM responses, progress indicators)
  • the total output size is unbounded
  • downstream processing should start before all input is available (pipeline stages)

Working example

A complete document ingestion pipeline combining generators, batching, and streaming write:

import json
import logging
from pathlib import Path
from typing import Iterator

logger = logging.getLogger(__name__)


def read_raw_docs(path: Path) -> Iterator[dict]:
    for i, line in enumerate(path.open()):
        try:
            yield json.loads(line.strip())
        except json.JSONDecodeError as exc:
            logger.warning("parse_error", extra={"line": i, "error": str(exc)})


def chunk_doc(doc: dict, max_chars: int = 1000) -> Iterator[dict]:
    text = doc.get("body", "")
    for i in range(0, len(text), max_chars):
        yield {
            "doc_id": doc["id"],
            "chunk_index": i // max_chars,
            "text": text[i : i + max_chars],
            "metadata": doc.get("metadata", {}),
        }


def stream_chunks(docs: Iterator[dict]) -> Iterator[dict]:
    for doc in docs:
        yield from chunk_doc(doc)


def write_chunks(chunks: Iterator[dict], path: Path) -> int:
    count = 0
    with path.open("w") as f:
        for chunk in chunks:
            f.write(json.dumps(chunk) + "\n")
            count += 1
    return count


def run_ingestion(input_path: Path, output_path: Path) -> None:
    docs = read_raw_docs(input_path)
    chunks = stream_chunks(docs)
    total = write_chunks(chunks, output_path)
    logger.info("ingestion_complete", extra={"chunks": total})
    print(f"Ingested {total} chunks")

The entire pipeline processes one document at a time. A 10GB input file produces output without ever holding more than one document in memory.

Common mistakes

Collecting all results before writing. results = list(pipeline) followed by writing defeats the memory benefits of generators. Stream results to the output file or database as they are produced.

No backpressure in streaming pipelines. If a producer generates items faster than a consumer can process them, an unbounded queue grows without limit. Use asyncio.Queue(maxsize=N) to create backpressure.

Ignoring partial failures in batch jobs. A single malformed document in a 100K batch should not abort the run. Catch exceptions per record, log them, and continue.

Re-reading large files multiple times. If your pipeline needs to filter then transform, compose those operations into a single pass. Reading a 10GB file twice is unnecessary.

Try it yourself

  1. Write a generator pipeline that reads a JSONL file, filters records with missing required fields, normalizes the text, and writes the cleaned records to a new JSONL file. Verify that it handles a 100K-row file without memory issues.
  2. Add a batching stage to an embedding workflow. Compare wall-clock time for batch size 1, 10, and 100.
  3. Implement a streaming LLM response handler that accumulates the full text while forwarding tokens to a queue. Write a test that verifies the final accumulated text matches the concatenation of all streamed tokens.
Lesson Deep Dive

Ask a follow-up question about this lesson and get an AI-powered explanation.

Loading previous questions...