Spark Execution Model — Jobs, Stages, Shuffles

Why this session matters

Part of a 36-topic learning series on engineering, ML, and LLM systems. Each session is a 90-minute deep-dive on one topic — designed so anyone can pick it up cold. Every two topics are followed by a revision session with recall prompts and hands-on drills.

Part 1: Spark Part 1 — Driver, Executors, RDDs, Lazy Evaluation

Why this session matters

It builds on the rhythm of one focused topic, paced so you have time to actually absorb it rather than rush.

Agenda

• Cluster topology — driver, executors, cluster manager, shuffle service
• RDD lineage and why Spark is lazy by default
• DataFrame vs Dataset vs RDD — when each one is the right answer
• Tasks, stages, jobs — the unit of work and where they get scheduled
• Local hands-on — read 5 GB Parquet, inspect a physical plan

Pre-read (assume zero background — watch/skim in order)

Videos (in this order, ~80 min total — do this before the session, not during):

1. MapReduce Explained — ByteByteGo — 8 min. The problem Spark solves. If you don't know what "shuffle" means in distributed data, start here.
2. Apache Spark in 100 Seconds — Fireship — 2 min. Vocabulary primer.
3. Spark Architecture Explained — Afaque Ahmad — 20 min. Driver / executor / cluster manager; DAG scheduler; job → stage → task.
4. Deep Dive into Spark SQL's Catalyst Optimizer — Databricks — 30 min (skim). How Spark rewrites your query into a physical plan.
5. Spark Performance Tuning — Databricks / Daniel Tomes — 25 min (skim second half). Skews, spills, AQE — the words you'll hear all session.

Notes to skim (if any term above felt fuzzy):

• New to distributed computing? Anatomy of a Spark Job — 10-min written walkthrough with pictures.
• Never seen a DataFrame? PySpark Quickstart — 5-min hands-on.
• Reference materials (keep tabs open, don't read cover-to-cover): Spark Tuning Guide (official) · Learning Spark 2e (free PDF) · Databricks Catalyst blog.

Prerequisite check before you start: you should be comfortable with (a) SQL GROUP BY, (b) what a "partition" of data means conceptually, (c) reading Python code with method chaining (df.filter(…).groupBy(…).agg(…)). If (b) is shaky, watch the MapReduce video twice.

Deep dive

1. Where Spark still wins

Spark is the workhorse for petabyte ETL and feature engineering because:

• One API for batch + streaming + ML.
• Runs on YARN, Kubernetes, or standalone — portable.
• Catalyst (the query optimiser) gets you most of the way without hand-tuning.
• It survived 10+ years and absorbed every lesson from MapReduce, Hive, and Tez.

Understanding the execution model is the difference between a 9-minute job and a 9-hour job. It's also the most common deep-dive in senior DE interviews.

2. Cluster topology

┌──────────────┐
│  Driver      │  holds SparkContext, plans queries,
│  (your code) │  schedules tasks, tracks state
└──────┬───────┘
       │
┌──────▼────────────────┐
│  Cluster Manager      │  YARN / Kubernetes / standalone
│  (allocates           │  hands out containers to the
│   executor JVMs)      │  driver on request
└──────┬────────────────┘
       │
   ┌───┴────┬────────┬────────┐
   ▼        ▼        ▼        ▼
 Executor Executor Executor Executor   (long-lived JVMs,
   1        2        3        N         hold cached partitions,
                                        run tasks on cores)
   │         │         │         │
   └─────────┴────┬────┴─────────┘
                  ▼
            shuffle service
            (writes/reads shuffle files between stages)

• Driver — your main(). Death of the driver = job dies. Keep its memory tight; avoid collect() of giant DataFrames.
• Executor — JVM with cores and memory. Holds cached partitions and runs tasks. Death of an executor = its tasks re-run elsewhere thanks to lineage.
• Cluster manager — schedules containers. Most production shops run on YARN or k8s.
• Shuffle service — external process per node that lets executors come and go without losing shuffle files (essential for dynamic allocation).

3. The RDD — lineage and laziness

An RDD (Resilient Distributed Dataset) is a plan for how to compute a partitioned collection, plus the lineage to recompute lost partitions. Two operation types:

• Transformations (map, filter, join, groupByKey) — lazy. They build the DAG.
• Actions (count, collect, save) — eager. They trigger execution.

This is why df.filter(...).filter(...).count() only runs once across the data — Spark fuses the filters into one pass at execution time.

4. Tasks → stages → jobs

A groupBy().count() produces 2 stages: stage 0 reads + partial aggregates locally; stage 1 (after shuffle) does the final aggregate.

5. Narrow vs wide transformations

• Narrow — output partition depends on exactly one input partition. map, filter, mapPartitions, narrow select. No shuffle, no network.
• Wide — output depends on multiple input partitions. groupByKey, join on un-bucketed keys, reduceByKey. Forces a shuffle — write to disk, network transfer, read back, merge.

Shuffles are 10–100× more expensive than narrow ops. Two reliable ways to flip a wide into a narrow:

1. Co-partition the inputs (repartition(N, key) on both sides ahead of a join — same N and same hash function).
2. Bucket at write time (df.write.bucketBy(N, "key")) so future joins are shuffle-free.

6. DataFrame vs Dataset vs RDD

In Python: just use DataFrame + SQL. RDD is now an escape hatch.

7. Catalyst — the four stages

When you write spark.sql("SELECT ...") or DataFrame ops, Catalyst walks:

1. Parsed logical plan — a tree, possibly invalid.
2. Analyzed logical plan — resolved against the catalog. Column names known.
3. Optimized logical plan — predicate pushdown, projection pruning, constant folding, join reordering.
4. Physical plan — choose the actual operator (SortMergeJoin vs BroadcastHashJoin), apply codegen.

Always check with:

df.explain("formatted")          # readable physical plan
df.explain("cost")               # with cost estimates (Spark 3+)

The "Exchange" operators in explain() are your shuffles. Count them; minimise them.

8. Local hands-on (the actual deliverable for this session)

docker run -it --rm -p 4040:4040 jupyter/pyspark-notebook

In a notebook:

from pyspark.sql import SparkSession, functions as F
spark = (SparkSession.builder
         .appName("s02").config("spark.sql.shuffle.partitions", 16)
         .getOrCreate())

# Use NYC taxi green-tripdata (free, ~150 MB/month) — fetch a year ≈ 2 GB.
df = spark.read.parquet("s3a://nyc-tlc/trip data/yellow_tripdata_2024-*.parquet")
df.printSchema()
print(df.count())

# Pretend join: aggregate by pickup zone, join to a zone lookup
zones = spark.read.csv("taxi_zone_lookup.csv", header=True)
agg = df.groupBy("PULocationID").agg(F.count("*").alias("trips"))
result = agg.join(zones, F.col("PULocationID") == F.col("LocationID"))
result.explain("formatted")
result.write.mode("overwrite").parquet("/tmp/out")

Open the Spark UI at http://localhost:4040. Look at:

• The stages tab — count of shuffles.
• The SQL tab — the physical plan with row counts at each operator.
• The executors tab — task time vs GC time vs shuffle read/write.

9. Common gotchas

• collect() on a big DataFrame → OOM on the driver. Use take(n) or write to disk.
• groupByKey (RDD) — never; use reduceByKey so partial aggregation happens in the map stage.
• UDF in Python — serialises every row to Python, breaks codegen. Use built-in functions or pandas UDFs.
• spark.sql.shuffle.partitions = 200 (default) for a 5 GB table = too many. Tune to ~(input GB × 2) for a starting point. With AQE (Spark 3+), Spark coalesces post-shuffle for you.

10. What's next (the next session — Spark Part 2)

• Shuffle anatomy — sort vs hash-based, bypass merge sort
• Catalyst rule examples
• AQE — adaptive coalescing, skew handling, join strategy switch
• Tuning recipes (memory, parallelism, broadcast threshold)

Reading material

Books:

• Learning Spark, 2nd ed. — Damji, Wenig, Das, Lee (free PDF from Databricks)
• Spark: The Definitive Guide — Bill Chambers, Matei Zaharia
• High Performance Spark — Holden Karau, Rachel Warren

Papers:

• Resilient Distributed Datasets (Zaharia et al., NSDI 2012) — the foundational RDD paper.
• Spark SQL: Relational Data Processing in Spark (SIGMOD 2015) — Catalyst + DataFrame paper.

Official docs:

• Apache Spark — Tuning Guide
• Apache Spark — Configuration reference
• PySpark API reference

Blog posts:

• Deep Dive into Spark SQL's Catalyst Optimizer — Databricks
• Adaptive Query Execution in Spark 3 — Databricks

In-depth research material

• apache/spark — github.com/apache/spark — ~40k ★, the canonical source. Start with core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala.
• Mastering Apache Spark (Jacek Laskowski) — free book — internals walk-through, kept up to date.
• Spark Summit / Data + AI Summit talks — best engineering content per minute on Spark.
• Photon: A Fast Query Engine for Lakehouse Systems (SIGMOD 2022) — Databricks' C++ Spark engine paper.
• Spark Architecture & internals — Databricks Engineering blog — long-form posts on shuffle, AQE, Photon.
• SE-Radio episode on Spark with Matei Zaharia — origin-story interview.

Videos

• Advanced Apache Spark Training — Sameer Farooqui (Databricks) — Spark Summit · 5 h 58 min — the legendary deep-dive on cluster topology, RDDs, DAGs, the scheduler, shuffle, and tuning. Watch the first 90 min for this session.
• Apache Spark Core — Deep Dive — Proper Optimization — Daniel Tomes, Databricks · 1 h 30 min — production tuning recipes from a Databricks SA who actually runs huge jobs.
• Apache Spark Core — Practical Optimization (Daniel Tomes) — Databricks · 53 min — the cleaned-up follow-up; same author, fresher numbers.
• A Deep Dive into Spark SQL's Catalyst Optimizer — Yin Huai (Databricks) · 28 min — the original architect explaining the four-stage plan pipeline.
• Exploring Wikipedia With Apache Spark — Advanced Training, Part 1 — Sameer Farooqui, Databricks · 2 h 37 min — hands-on, with the cluster UI open beside the code. Sample sections you care about.

LeetCode — Group Anagrams

• Link: https://leetcode.com/problems/group-anagrams/
• Difficulty: Medium
• Why this problem: Sort the string or use a 26-count signature as the hash-map key.
• Time-box: 30 minutes. Look up the editorial only after.

Assignment / Deliverables

Give yourself a clean 2-hour window and complete all of these before moving on:

1. Read the deep-dive above end-to-end — no skimming. Take notes in your own words.
2. Solve the LeetCode problem below without help first. Only look at the hint after 15 focused minutes; only look at editorial after 30. Log your time.
3. Reproduce one code snippet locally. Pick the snippet that felt least obvious and get it running in a scratch file / notebook.
4. Draw the core diagram from memory. Paper, whiteboard, or tldraw — doesn't matter. If you can't, re-read section 2 and try again.
5. Write a 3-line takeaway in your prep journal: what surprised you, what you still don't understand, what you'd read next.
6. Skim one item from the Reading material section. Bookmark the rest for the weekend.
7. Commit any code + notes to your prep repo with message session-NN: <one-line summary>.

Stretch (optional, +30 min): explain today's topic to a rubber duck / a friend / a voice note. If you can't teach it in 5 minutes, you don't own it yet — flag it and revisit next weekend.

Post-session checklist

By the end of this session you should be able to:

• Draw the cluster topology and name what each box does.
• Define narrow vs wide transformation; give 3 examples of each.
• Walk through the 4 Catalyst stages on a small SQL example.
• Read a .explain('formatted') output and identify shuffle boundaries.
• Tune spark.sql.shuffle.partitions for a known input size.
• Pick the right API (DataFrame vs RDD) for a given workload.

Generated from sessions_data.py + content_part*.py. To edit a video / leetcode / title, edit the data file and re-run write_sessions.py.

Part 2: Spark Part 2 — Shuffles, Catalyst, AQE, Tuning

Why this session matters

It builds on the rhythm of one focused topic, paced so you have time to actually absorb it rather than rush.

Agenda

• Shuffle anatomy — map-side write, network transfer, reduce-side read
• Sort-based vs hash-based shuffle, bypass merge sort
• Catalyst optimisations you can actually trigger or block
• AQE — dynamic coalesce, skew splitting, join strategy switch
• A tuning checklist that survives most production jobs

Deep dive

1. Where shuffles come from

A shuffle is forced whenever output partitions must be re-derived from input partitions that have moved keys around. Three common triggers:

• groupByKey / reduceByKey — keys must end up on the same reducer.
• join on un-bucketed columns — matching keys must co-locate.
• repartition(N) / repartition(N, col) — explicit redistribution.

In the physical plan, shuffles appear as Exchange operators. Every Exchange = disk write on the map side + network transfer + disk read on the reduce side. Easily the most expensive thing Spark does.

2. Anatomy of a shuffle (sort-based, the modern default)

  Map side                                 Reduce side
  --------                                 -----------
  task 1: partition data,                  task 1: fetch its slice
          sort by (key, partition),                  from every map task,
          spill to disk if needed                    merge-sort,
  task 2: same                                       feed to reducer
  …                                         task 2: …

Each map task writes one shuffle file plus an index. The reducer issues range reads against every map output. With M map tasks and R reducers, you get up to M×R network requests — a lot at scale ("reducer fan-in" can spike NIC).

3. Bypass merge sort

When reducers < spark.shuffle.sort.bypassMergeThreshold (default 200), Spark uses hash-based shuffle: one file per reducer per map task, no per-task sort. Cheaper for low-reducer jobs but creates many small files. AQE makes this less relevant; default is fine for most.

4. Skew — the silent job killer

If one key has 50% of the rows (think: country='US'), one reducer task gets 50% of the work. The job is gated by the slowest task. Look for it in the Spark UI:

• One executor task taking 20 min while others finish in 30 s.
• Massive spill (memory) and spill (disk) on a single task.

Fixes:

• AQE skew join (Spark 3+): set spark.sql.adaptive.skewJoin.enabled = true. Detects skewed partitions post-shuffle and splits them into multiple tasks. Set spark.sql.adaptive.skewJoin.skewedPartitionFactor = 5 and spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes = 256MB.
• Salting — add a random 0..N suffix to the hot key on both sides; explode rows on the other side. Costly but always works.
• Broadcast the small side if it fits (< spark.sql.autoBroadcastJoinThreshold, default 10 MB — raise to 100 MB for joins to dimensions).

5. Catalyst optimisations you can trigger or block

Things that block Catalyst:

• Python UDFs — opaque, prevent codegen and pushdown.
• from_json(col, schema) without schema — prevents projection pruning.
• select(*) early — prevents projection pruning; only select what you need.

6. AQE — Adaptive Query Execution (Spark 3+, default on in 3.2+)

AQE re-plans during execution using actual shuffle stats:

1. Dynamic coalesce — after a shuffle, if many small partitions, coalesce to spark.sql.adaptive.advisoryPartitionSizeInBytes (default 64 MB). Replaces the old "tune shuffle.partitions by hand" ritual.
2. Skew join handling — splits skewed partitions (see above).
3. Join strategy switch — if one side of a SortMergeJoin turns out small after filtering, switch to broadcast.

Let AQE do its thing and only override when the SQL UI shows it making a bad call.

7. Memory & storage knobs that matter

Rule of thumb for a 1 TB job on 50 executors with 8 cores each:

• spark.executor.memory = 16g, memoryOverhead = 4g
• spark.sql.shuffle.partitions = 2000 (rough, AQE will coalesce)
• Bump autoBroadcastJoinThreshold to 100 MB if joining dimensions
• Enable AQE: spark.sql.adaptive.enabled = true

8. Tuning checklist (use this on every slow job)

1. Open Spark UI → SQL tab → click the job.
2. Look at the DAG. Count Exchange operators — each is a shuffle.
3. Stage tab: which task is the slowest? Skew = 1 task >> others.
4. Executor tab: GC time / total time. > 10% → raise heap or split executors.
5. Storage tab: are you caching things you don't reuse?
6. Check Spark UI → SQL → cost. CBO making the right join order?
7. If broadcasting, confirm the small side was actually broadcast (look for BroadcastHashJoin in the plan).
8. If a join is skewed, verify AQE is on and skewedPartitionFactor is reasonable.

9. Real production levers (from large jobs)

• Replacing a groupBy + collect_list with a mapPartitions + per-partition aggregator to avoid a giant shuffle.
• repartitionByRange(num_partitions, col) before a sort-heavy write to avoid global sort.
• Bucketing the largest fact table at write time so future joins are shuffle-free.
• Materialising intermediate stages with df.persist(StorageLevel.MEMORY_AND_DISK_SER) when they're used 3+ times in the DAG.
• Pinning spark.sql.files.maxPartitionBytes to align with target output file size to avoid massive shuffle write fan-out.

10. Code snippets you'll actually use

from pyspark.sql import functions as F

# 1. Force broadcast (small dim)
orders.join(F.broadcast(customers), "customer_id")

# 2. Salt a known-skewed key
N = 16
left  = left.withColumn("salt", (F.rand() * N).cast("int"))
right = right.withColumn("salt", F.explode(F.array([F.lit(i) for i in range(N)])))
left.join(right, ["key", "salt"])

# 3. Inspect actual physical plan
df.explain("formatted")   # operator tree
df.explain("cost")        # with cost estimates (CBO)

# 4. Repartition by range before a sort-heavy write
(df.repartitionByRange(2000, "event_ts")
   .write.mode("overwrite").partitionBy("event_date").parquet(out))

11. What's next (the next session — Kafka Part 1)

Not Spark — next DE session jumps to Kafka. Spark + Kafka is the streaming workhorse pair; we cover Flink/Spark structured streaming in the next session.

Reading material

Books:

• Learning Spark, 2nd ed. — Damji et al. (chs. 3 & 4 on the optimiser; ch. 6 on tuning)
• Spark: The Definitive Guide — Chambers, Zaharia (ch. 19: Performance Tuning)
• High Performance Spark — Holden Karau, Rachel Warren — the whole book is shuffle / Catalyst tuning.

Papers:

• Spark SQL: Relational Data Processing in Spark (SIGMOD 2015) — the Catalyst paper.
• Adaptive Query Execution: Speeding Up Spark SQL at Runtime (Databricks, 2020) — the AQE writeup.
• Photon: A Fast Query Engine for Lakehouse Systems (SIGMOD 2022) — next-gen vectorised engine.

Official docs:

• Spark — Tuning Guide
• Spark — SQL Performance Tuning
• Spark — Adaptive Query Execution

Blog posts:

• Deep dive into Catalyst — Databricks
• Skew join handling in AQE — Databricks
• Understanding Spark Shuffle — Cloudera blog — anatomy of sort vs hash shuffle.

In-depth research material

• apache/spark/sql/catalyst — github.com/apache/spark/tree/master/sql/catalyst — read Optimizer.scala, RuleExecutor.scala.
• Mastering Spark SQL (Jacek Laskowski) — the most-detailed online reference.
• Spark Internals (Sameer Farooqui) — slides — slide deck companion to the legendary video.
• Databricks Engineering — AQE evolution series
• Apache Spark Architecture — Cloudera deep dive — internals of stage boundaries.
• Spark Summit talks — "Tuning Spark" playlist — Sameer + Daniel Tomes back catalog.

Videos

• Apache Spark Core — Deep Dive — Proper Optimization — Daniel Tomes, Databricks · 1 h 30 min — the canonical talk on shuffle anatomy, skew, partitioning, AQE.
• A Deep Dive into Spark SQL's Catalyst Optimizer — Yin Huai (Databricks) · 28 min — the four-stage Catalyst plan from the original architect.
• A Deep Dive into the Catalyst Optimizer — Hands-on — Herman van Hovell, Databricks · 23 min — adds rule-writing examples.
• Apache Spark Core — Practical Optimization — Daniel Tomes · 53 min — the second-iteration tuning talk with newer AQE numbers.
• Deep Dive Into Catalyst: Apache Spark 2.0's Optimizer — Spark Summit · 30 min — useful historical lens — the optimizer's evolution from 1.x → 3.x.

LeetCode — Top K Frequent Elements

• Link: https://leetcode.com/problems/top-k-frequent-elements/
• Difficulty: Medium
• Why this problem: Hash-map count then heap of size k; bucket-sort gives O(n).
• Time-box: 30 minutes. Look up the editorial only after.

Assignment / Deliverables

Give yourself a clean 2-hour window and complete all of these before moving on:

1. Read the deep-dive above end-to-end — no skimming. Take notes in your own words.
2. Solve the LeetCode problem below without help first. Only look at the hint after 15 focused minutes; only look at editorial after 30. Log your time.
3. Reproduce one code snippet locally. Pick the snippet that felt least obvious and get it running in a scratch file / notebook.
4. Draw the core diagram from memory. Paper, whiteboard, or tldraw — doesn't matter. If you can't, re-read section 2 and try again.
5. Write a 3-line takeaway in your prep journal: what surprised you, what you still don't understand, what you'd read next.
6. Skim one item from the Reading material section. Bookmark the rest for the weekend.
7. Commit any code + notes to your prep repo with message session-NN: <one-line summary>.

Stretch (optional, +30 min): explain today's topic to a rubber duck / a friend / a voice note. If you can't teach it in 5 minutes, you don't own it yet — flag it and revisit next weekend.

Post-session checklist

By the end of this session you should be able to:

• Trace a shuffle through map→network→reduce, naming every step.
• Diagnose skew from the Spark UI and pick the right fix (AQE / salt / broadcast).
• Walk through AQE's three big optimisations.
• List 5 Catalyst optimisations and how to keep them firing.
• Tune executor.memory, shuffle.partitions, broadcast threshold for a 1 TB job.
• Solve top-k-frequent-elements two ways (heap + bucket-sort) — same shape as a reduce-side aggregate.

Generated from sessions_data.py + content_part*.py. To edit a video / leetcode / title, edit the data file and re-run write_sessions.py.