Data Engineering for AI Systems from First PrinciplesPost 7 of 21 · 21 of 21 written
- How to use this guide
- 01Why AI runs on data engineering
- 02Representing data: bytes, types, schemas, formats
- 03Storage: where bytes live and what reading them costs
- 04Databases from first principles
- 05Data modeling: shaping data for its readers
- 06Distributed systems: the physics of many machines
- 07Batch processing: transforming data in bulk
- 08Streaming: data that never ends
- 09Architectures: warehouse, lake, lakehouse
- 10Pipelines and orchestration
- 11Data quality, contracts, lineage and observability
- 12Data engineering for machine learning
- 13Data engineering for LLMs
- 14Data engineering for RAG and AI applications
- 15Feeding GPUs: data loading for training
- 16Governance, privacy and security
- 17Cost and performance reasoning
- 18The tool landscape, mapped to principles
- 19Learning path and hands-on projects
- 20Question bank for fluency
- 21Glossary
7. Batch processing: transforming data in bulk
Batch processing runs a computation over a bounded dataset in one go. Its core moves, split the data, process the pieces in parallel, regroup by key, combine, have not changed since MapReduce in 2004; only the engines got faster and friendlier.
Bounded versus unbounded data
A bounded dataset has an end: yesterday’s logs, a snapshot of a table, a crawl. An unbounded dataset never ends: clicks arriving forever. Batch processing handles bounded data; streaming (section 8) handles unbounded data. Many systems treat batch as a special case of streaming with a known end.
The MapReduce idea
- Map: apply a function to each record independently, emitting key/value pairs. For word count: each word becomes (word, 1).
- Shuffle: move all pairs with the same key to the same machine. This crosses the network.
- Reduce: combine all values for each key: (word, total).
Three lessons outlived the original system. Send computation to the data, not data to the computation. Make functions deterministic, so a failed task can simply be re-run. And the shuffle is almost always the expensive part.
Apache Spark
Spark generalized MapReduce and kept intermediate data in memory. The concepts you need:
- Driver and executors: the driver plans the job; executors on worker machines run tasks.
- Partitions: the unit of parallelism. One task processes one partition.
- Lazy evaluation: transformations (filter, select, join) only build a plan. Nothing runs until an action (count, write, collect).
- Narrow versus wide transformations: narrow ones (filter, map) work within a partition. Wide ones (groupBy, join, distinct) need a shuffle and start a new stage.
- The optimizer (Catalyst) rewrites DataFrame and SQL plans: pushing filters down, pruning columns, choosing join strategies.
- Fault tolerance through lineage: Spark remembers how each partition was computed, and recomputes lost ones instead of replicating them.
Shuffles and how to avoid them
A shuffle serializes data, writes it to local disk, sends it across the network and reads it back. To shuffle less:
- Filter rows and select columns as early as possible.
- Use a broadcast join when one side is small: copy it to every executor, and the large side never moves.
- Pre-partition or bucket tables on join keys that are used repeatedly.
- Aggregate before joining, when the logic allows.
Join strategies at scale
| Strategy | How it works | Use when |
|---|---|---|
| Broadcast hash join | Small table copied to every worker | One side fits comfortably in memory |
| Shuffle hash join | Both sides shuffled by key, hash table per partition | Medium-sized sides |
| Sort-merge join | Both sides shuffled and sorted by key, then merged | Two large tables, the default for big joins |
Skew, the silent job killer
If one key holds 30% of the data, one task does 30% of the work while hundreds of others sit idle. Remedies:
- Salting: append a random suffix to the hot key, aggregate in two stages, then remove the suffix.
- Adaptive query execution (AQE in Spark 3): detects skewed partitions at runtime and splits them.
- Treat hot keys separately: process the few giant keys with a different strategy.
In the Spark UI, one task taking far longer than its siblings almost always means skew.
Partition sizing and memory
A common rule of thumb is 100 to 200 MB per partition. Too many tiny partitions waste time on scheduling overhead. Too few large ones limit parallelism and run out of memory. When memory runs short, engines spill to disk, which is slow but survivable; when spilling is impossible, the task fails with an out-of-memory error.
Why modern engines are fast: vectorized execution
Old engines processed one row at a time through a chain of function calls. Modern engines process batches of column values: a tight loop over 2,048 integers that the CPU can run with SIMD instructions and keep in cache. This, plus columnar files, is why DuckDB, Polars, ClickHouse and Spark’s newer engines are an order of magnitude faster than row-at-a-time systems.
The single-node renaissance
| Tool | What it is | Reach for it when |
|---|---|---|
| pandas | Eager, in-memory DataFrame | Small data, quick exploration, ecosystem compatibility |
| Polars | Rust DataFrame with a lazy, optimized API | Medium data on one machine, fast transforms in Python |
| DuckDB | In-process analytical SQL engine | SQL over Parquet, CSV or S3 without any server |
| Spark | Distributed engine on a cluster | Data or compute clearly exceeds one machine |
| Ray | Distributed Python tasks and actors | Python and GPU-heavy work: batch inference, preprocessing |
The same query in two single-node tools, reading Parquet straight from storage:
import duckdb
duckdb.sql("""
SELECT language, count(*) AS docs, avg(token_count) AS avg_tokens
FROM 's3://my-bucket/docs/*.parquet'
WHERE fetched_at >= '2026-10-01'
GROUP BY language
""").show()
import polars as pl
(pl.scan_parquet("s3://my-bucket/docs/*.parquet") # lazy: nothing read yet
.filter(pl.col("fetched_at") >= pl.datetime(2026, 10, 1))
.group_by("language")
.agg(pl.len().alias("docs"), pl.col("token_count").mean())
.collect()) # plan optimized, then runIn both, the engine reads only the needed columns and skips row groups outside the date range. That is Parquet pruning from section 2, paying off.
Batch inference is a batch job
Running a model over millions of items, such as embedding a corpus or scoring documents for quality, is batch processing with a GPU in the middle. The concerns are the same plus a few more:
- Keep the GPU busy: batch inputs, overlap CPU preprocessing with GPU work.
- Stages need different hardware: parsing on CPUs, the model on GPUs. Ray Data and similar tools pipeline these stages.
- Checkpoint progress, so a failure at 80% resumes rather than restarts.
- Write results idempotently, keyed by item ID.
Rules for well-behaved batch jobs
- Treat inputs as immutable.
- Make outputs a deterministic function of inputs and code version.
- Partition outputs by time or by run, and overwrite whole partitions.
- Make every run re-runnable for any past date.
- Log row counts in and out of every step; a sudden change is your first alarm.
Questions to ask
- Does this fit on one machine with DuckDB or Polars?
- Where are the shuffles, and can any be avoided?
- Is any key skewed?
- What happens if the job dies at 90%?
- Is the output identical if I run it twice?