Python for Data Engineering

Beyond notebooks and one-off scripts

Most people learn Python in a notebook: cells you run top to bottom, variables floating in memory, a chart at the end. That's a wonderful way to explore data. It's a terrible way to ship it. A notebook hides state (you ran cell 7 before cell 4, who knows), can't be tested, and is painful to schedule or review. Production data code looks different: it lives in modules (plain .py files), it's organized into functions, and it's versioned in git so every change is reviewable and reversible.

Recall the uv project you set up in Course 2. That layout is the starting point for everything here:

my-pipeline/
├── pyproject.toml        # dependencies + project metadata (uv manages this)
├── uv.lock               # exact pinned versions, committed to git
├── src/
│   └── my_pipeline/
│       ├── __init__.py   # marks this folder as an importable package
│       ├── extract.py    # one module per stage
│       ├── transform.py
│       └── load.py
└── tests/
    └── test_transform.py # tests live next door, not inside the code

A module is just a file. An import pulls names from one file into another (from my_pipeline.transform import clean_rows). A package is a folder with an __init__.py. None of this is exotic — it's how you turn a wall of cells into named, reusable pieces. The payoff: each function can be tested in isolation, imported by the scheduler, and reviewed in a pull request line by line.

Type hints: catching shape drift early

Python doesn't force you to declare types, but you can annotate them. A type hint is a note to readers and to tools about what a function expects and returns. It changes nothing at runtime — Python won't stop you passing the wrong type — but a type checker reads the hints and flags mismatches before you run anything.

This matters enormously in data code, where the bugs are almost always shape and type drift: a column that was an int yesterday arrives as a str today, a function that promised a list quietly returns None, a dict is missing a key. Types make those promises explicit.

from datetime import date

def daily_total(rows: list[dict[str, float]], day: date) -> float:
    """Sum the 'amount' field across all rows for one day."""
    total = 0.0
    for row in rows:
        total += row["amount"]      # type checker knows this is a float
    return total

Read the signature like a contract: "give me a list of dicts whose values are floats, plus a date; I'll give you back a float." If somewhere else you call daily_total(df, "2026-06-21") — passing a DataFrame and a string — a type checker catches it instantly, instead of you discovering it via a confusing error in production.

Error handling: fail loudly, not silently

When something goes wrong, Python raises an exception — an object describing the failure that travels up the call stack until something catches it or the program stops. You catch exceptions with try/except, and clean up with finally (which runs no matter what):

def parse_amount(raw: str) -> float:
    try:
        return float(raw)
    except ValueError:
        # the WRONG thing to do is `return 0.0` and move on — that
        # silently corrupts your data. Re-raise with context instead.
        raise ValueError(f"could not parse amount from {raw!r}")
    finally:
        pass  # e.g. close a file or connection here, always

The cardinal sin in data engineering is the silent swallow: except: pass. It turns a loud, debuggable crash into a quiet wrong number that nobody notices for weeks. A pipeline that fails is annoying; a pipeline that succeeds with corrupt output is a disaster, because everyone downstream trusts it. Fail loudly. Let bad data stop the line.

For pipelines, define custom exceptions per stage so failures are self-labeling:

class PipelineError(Exception):
    """Base class for everything our pipeline raises."""

class ExtractError(PipelineError): ...
class TransformError(PipelineError): ...
class LoadError(PipelineError): ...

Now a stack trace that ends in TransformError tells you where in the flow it broke before you read a single line. Here's how to decide what to do when something fails:

Iterators & generators: data bigger than memory

Here's a scenario you'll hit constantly: a 50 GB CSV and a laptop with 16 GB of RAM. rows = open("huge.csv").readlines() tries to load all 50 GB into memory at once and crashes (that OOM from the Start Here table). The fix is laziness: process one row at a time, holding only that row in memory, and let the rest stay on disk until you reach it.

Python's tool for this is the generator — a function that uses yield instead of return. Each yield hands back one value and pauses, resuming where it left off when you ask for the next. It produces values one at a time, on demand, never building the whole list.

Let's prove the memory stays flat. This generator yields squares forever; we consume a billion of them and watch memory not grow:

import tracemalloc

def squares(n: int):
    """A generator: yields one value at a time, never builds a list."""
    for i in range(n):
        yield i * i

tracemalloc.start()
total = 0
for sq in squares(1_000_000_000):   # a BILLION values
    total += sq
    if total > 10_000_000:          # stop early; the point is memory, not the sum
        break
current, peak = tracemalloc.get_traced_memory()
print(f"peak memory: {peak / 1024:.1f} KB")
print(f"summed up to: {total}")

This is why generators are core to data engineering, not a nice-to-have. Streaming a file row by row, paging through an API, transforming records as they flow — all of it is generators. The pattern below reads any file lazily; it works the same on a 1 KB file and a 1 TB file:

def read_rows(path: str):
    """Yield one line at a time. Memory stays flat for any file size."""
    with open(path) as f:
        for line in f:          # iterating a file object is ALREADY lazy
            yield line.rstrip("\n")

# Count non-empty rows in a huge file without ever loading it:
count = sum(1 for row in read_rows("huge.csv") if row)
print(count)

The data libraries: which one, when

Plain Python (loops, dicts, generators) is perfect for glue, streaming, and oddly-shaped data. But for tabular work — filtering, joining, grouping millions of rows — you want a library that does the heavy lifting in fast compiled code. Three dominate, and they're good at different things:

Notice the family resemblance: pandas and polars give you a Python DataFrame API; DuckDB gives you SQL (everything you learned in Chapter 05) that runs straight over CSV and Parquet files. They overlap heavily, and you'll mix them — read with DuckDB, refine with polars, hand a small result to pandas for a chart.

Testing with pytest

"It ran without errors" is not the same as "it produced the right numbers." Data engineers test transform logic — the functions that clean, reshape, and compute — because that's where silent correctness bugs hide. A test pins down the expected output for a known input, so when you change the code later, you find out immediately if you broke the meaning.

The standard tool is pytest. The convention: put functions named test_* in files named test_*.py, and assert what should be true. Here's a transform and its test:

def normalize_region(region: str) -> str:
    """Lowercase and strip a region code. The thing we want to trust."""
    return region.strip().lower()

def test_normalize_region():
    assert normalize_region("  US ") == "us"
    assert normalize_region("Eu") == "eu"
    assert normalize_region("apac") == "apac"

Run the whole suite from your project root:

uv run pytest

Structuring a pipeline: extract, transform, load

Pull the threads together. A pipeline is three jobs — extract (get the data), transform (clean and reshape it), load (write it somewhere) — and the trick is to make each one a separate pure function. A pure function depends only on its inputs and returns a value with no hidden side effects. Pure functions are testable (same input → same output, every time) and they compose cleanly.

def extract(path: str):
    """Yield raw rows lazily — memory stays flat (see #iterators)."""
    with open(path) as f:
        header = f.readline().rstrip("\n").split(",")
        for line in f:
            yield dict(zip(header, line.rstrip("\n").split(",")))

def transform(rows):
    """Pure: clean each row. No I/O, no surprises — easy to test."""
    for row in rows:
        yield {
            "region": row["region"].strip().lower(),
            "amount": float(row["amount"]),
        }

def load(rows, conn) -> int:
    """The ONLY function with a side effect: writing out."""
    n = 0
    for row in rows:
        conn.execute("INSERT INTO sales VALUES (?, ?)",
                     (row["region"], row["amount"]))
        n += 1
    return n

# Wire them together — data streams through, one row at a time:
# load(transform(extract("sales.csv")), conn)

Two things are happening here. First, it's a stream: extract and transform are generators, so a row flows extract → transform → load and is forgotten before the next one starts — constant memory, any file size. Second, the messy stuff (file I/O, the database) is isolated in extract and load, leaving transform a pure function you can test with a plain list of dicts and no database at all.

✓ Check yourself

• Why is a notebook a poor home for production data code — name two concrete reasons?
• What's the danger of except: pass, and what should you do instead?
• How does a generator keep memory flat on a file bigger than RAM?
• Which would you reach for: a 200-row CSV needing a quick chart vs. a join across three 5 GB Parquet files?
• Why are pure functions easier to test, and how does that connect to idempotency?

def clean_rows(rows: list[dict]):
    """Yield cleaned rows; drop rows with a missing region."""
    for row in rows:
        region = row.get("region", "").strip().lower()
        if not region:
            continue                      # skip junk, don't crash on it
        yield {"region": region, "amount": float(row["amount"])}

from clean import clean_rows

def test_clean_rows():
    raw = [
        {"region": "  US ", "amount": "120"},
        {"region": "eu",    "amount": "80.5"},
        {"region": "",      "amount": "999"},   # should be dropped
    ]
    out = list(clean_rows(raw))               # materialize the generator
    assert out == [
        {"region": "us", "amount": 120.0},
        {"region": "eu", "amount": 80.5},
    ]
    assert len(out) == 2                       # the empty-region row was skipped

Next

You can write structured, tested, memory-aware Python — now let's make the files it reads and writes fast and compact, from CSV's pitfalls to columnar Parquet. → Files & Serialization at Scale