Lab 01 — Spark Execution Simulator

Phase: 06 — Apache Spark Internals | Difficulty: ⭐⭐⭐⭐☆ | Time: 5–7 hours

Round 3 of the interview hands you "a Spark job on EMR runs for 9 hours, spills heavily, produces millions of small Parquet files, and makes Athena scan too much." This lab builds the calculator and diagnostics behind every sentence of that answer: stage planning at shuffle boundaries, skew detection, AQE coalescing, broadcast decisions, shuffle sizing, spill estimation, and a small-file report.

What you build

  • plan_stages / num_stages — split a logical plan into stages at shuffle (wide) boundaries (the thing that makes Spark Spark)
  • detect_skew — partitions far above the median → the "one task runs forever" bug
  • coalesce_partitions — AQE's merge of small post-shuffle partitions (and why a fat partition stands alone — coalesce never splits)
  • should_broadcast — broadcast-hash-join decision (skip the shuffle for a small side)
  • shuffle_partitions — size parallelism by bytes (the fix for the wrong default 200)
  • estimate_spill — when a partition exceeds execution memory → spill (a partition-count problem, not just a memory one)
  • small_file_report — the compaction pitch ("1M files → 8k") that fixes Athena cost (P10)

Key concepts

ConceptWhat to understand
Stage boundarya shuffle materializes data; narrow ops pipeline within a stage
Skewa few partitions >> median dominate stage time; fix by salt/AQE/broadcast
AQE coalescemerge small partitions post-shuffle; never split (skew-join does that)
Broadcast joinship the small side everywhere; avoid the shuffle
Spillpartition > execution memory → disk; more partitions = less spill
Small filesthe #1 cause of scan-cost and metadata blowups (P08–P10)

Run

pip install -r requirements.txt
pytest test_lab.py -v
LAB_MODULE=solution pytest test_lab.py -v

Success criteria

  • All 16 tests pass.
  • You can narrate the Round-3 answer end to end using these functions.
  • You can explain why spill is fixed by more partitions as often as by more memory.
  • You can explain why a >target partition is left alone by coalesce but split by skew-join.

Extensions

  • Add salting: split a hot key into key#0..key#n, re-aggregate, and show the skew factor (P01) drop — the manual version of AQE skew-join.
  • Model AQE dynamic join switching (a sort-merge join that becomes a broadcast join once runtime stats show one side is small).
  • Add an idempotent write planner: stage-dir → atomic commit (foreshadows P09 table-format commits) so a retried job doesn't double-write.