PySpark · ML Systems Lab

PySpark Shuffle: The Complete Mental Model

Every wide transformation triggers a shuffle. Most engineers know this. Few have a clear mental model of what actually happens — the partition lifecycle, the disk writes, the network transfers, the sort. Without that model, you're guessing when you tune. With it, you can look at a Spark UI and immediately know what's wrong.

A shuffle happens in three phases: map, shuffle write, shuffle read.

The map phase: Each executor runs its tasks on its input partitions. At the end of a map task, the output is partitioned by the shuffle key (hash of the join/group-by key). Each map task writes one output file per downstream reducer.

Shuffle write: The map output is spilled to local disk. If executor memory is exceeded before the map phase completes, Spark spills partial output to disk and merges later. This is the first place skew kills you: a task with a hot key writes 10× the data of others.

Shuffle read: Each reduce task reads its designated partition from every mapper. If you have 200 mappers and 200 reducers, that's 40,000 small file reads. The network transfer is proportional to data volume, not task count — but the connection overhead scales with task count.

The key metrics in Spark UI:

Shuffle Write: bytes written by map tasks. High = expensive downstream reads.

Shuffle Read: bytes read by reduce tasks. Should approximately equal shuffle write.

Task Duration (max vs median): ratio > 5× = data skew.

GC Time: high GC = executor heap pressure = likely spill incoming.

Spill (memory): amount spilled from RAM to disk. Non-zero = your partitions are too large for executor heap.

The key dials:

`spark.sql.shuffle.partitions` (default: 200): target 128–256 MB per partition after shuffle. For 100GB of shuffled data: 100*1024 / 200 = 512 MB per partition — too large. Try 800.

`spark.sql.autoBroadcastJoinThreshold` (default: 10MB): raise to 500MB if one join side fits.

`spark.sql.adaptive.enabled=true`: AQE dynamically coalesces small partitions and splits large ones at runtime. Enable by default in Spark 3.

AQE is not magic. It can't save a job with a 5000:1 key skew if you haven't enabled AQE skew join (`spark.sql.adaptive.skewJoin.enabled=true`). And even with skew join enabled, extremely hot keys (> 20× median) may require explicit salting.

Continue interactively

Read this post inside ML Systems Lab — with Simplify toggle, interview Q&As, inline glossary, and the MLE Path forward pointer.

Open in MSL →