Search Tech Journey

Find topics, journeys and posts

6-month learning plan47 / 130
back to blog
data engineeringintermediate 55m read

S047 · Spark — RDD, DataFrame, Jobs/Stages/Shuffles

The most misunderstood distributed engine in the world. Learn what actually happens when you write .groupBy().agg() — the DAG, the shuffle, the stages, the spills — and why 90 % of Spark performance is avoiding shuffles.

🗄️Data EngineeringM05 · Data Engineering· Session 047 of 130 90 min

🎯 Trace a PySpark job from action → DAG → stages → tasks → shuffle files → output, and optimise the top three bottlenecks: skew, small files, unnecessary shuffles.

Why this session exists

Everyone writes Spark; almost nobody knows what Spark actually does. You get told "use DataFrames not RDDs, use broadcast joins, cache your DataFrames" — but if you can't explain why in terms of the DAG, stages, and shuffle files, you're cargo-culting. Ninety minutes of understanding the execution model is worth a decade of Stack Overflow copy-paste. This session is that ninety minutes.

You will be able to
  • Explain lazy evaluation, transformations vs actions, and why .count() finally triggers work.
  • Trace a job → stages (broken by shuffle boundaries) → tasks (one per partition).
  • Diagnose data skew from the Spark UI and fix it with salting or AQE skew join.
  • Explain why broadcast joins avoid the shuffle and when they blow up the driver.
  • Read the Spark UI's ‘Stages’ tab and identify the top 3 causes of slowness.

Prerequisites

  • S040 — Joins & Set Operations (Spark's SQL model is standard SQL).
  • S045 — Data modelling (partition strategy = physical modelling).
  • S046 — Batch vs streaming (Spark straddles both).


(a) Intuition · 5 min

A distributed cooking brigade
🌍 Real world

Imagine you're catering a wedding for 10 000 guests. You have 100 chefs. You could give each chef a random 100 dishes and let them go — but then whoever gets ‘slice 3 000 onions’ finishes 3 hours after everyone else. That's skew.

Now imagine each dish requires ingredients from different pantries. Chefs waste 40 % of the day walking between pantries. That's the shuffle. Fix: pre-sort ingredients by which chef will use them, so each chef gets a full crate delivered. That pre-sort is expensive but done once.

Finally: the head chef writes the whole menu on a whiteboard before anyone starts. Only when someone asks ‘is the wedding cake done?’ do they actually begin cooking, and only the parts needed for cake. That's lazy evaluation.

💻 Code world

Spark: transformations (map, filter, groupBy, join) build a plan; actions (count, collect, write) execute it. The scheduler breaks the plan into stages at every shuffle boundary. Each stage runs as many parallel tasks as there are partitions. Between stages, data is shuffled: every task writes N output files (one per reduce partition), and the next stage's tasks read the matching files across the cluster.

90 % of Spark performance is: (1) reduce shuffle volume, (2) fix skew, (3) match parallelism to the data. Everything else is decoration.

The four ideas that unlock everything

Spark's execution model in one breath
  • Lazy — transformations are recorded, not executed, until an action.
  • Immutable + partitioned — every DataFrame is a set of partitions distributed across executors.
  • Shuffles cost network I/O — every groupBy, join, distinct, and repartition creates one.
  • DataFrames > RDDs — the Catalyst optimizer rewrites your query into a better plan; RDDs skip Catalyst entirely.

A quick history so you know why the world looks like this

  1. 2009
    Spark started at UC Berkeley AMPLab
    Matei Zaharia's response to Hadoop MapReduce being 100× too slow for iterative ML.
  2. 2013
    Spark 0.7 — RDDs
    First open-source release. In-memory + fault tolerance via lineage. Kills Hadoop.
  3. 2015
    DataFrames + Catalyst
    Structured API + cost-based optimizer. Now you can write SQL and beat hand-written RDD code.
  4. 2016
    Structured Streaming
    Same DataFrame API on unbounded data. Micro-batch by default; continuous mode experimental.
  5. 2020
    Spark 3.0 · Adaptive Query Execution (AQE)
    Runtime re-optimisation: coalesce partitions, dynamic broadcast, skew join handling.
  6. 2023
    Spark 3.5 · Connect
    Client-server split — thin PySpark client talks to a remote Spark server. Cleaner ops.

(b) Visual walkthrough · 15 min

From code to cluster in five layers

Anatomy of a shuffle

Each map task writes M shuffle files (one per reduce partition). Each reduce task fetches its file from every map task. Total network I/O = N × M small file reads. This is the cost centre.

Join strategies Spark can pick

Broadcast Hash Join

Small table shipped to every executor.

  • No shuffle of the big side
  • Small side must fit in memory (default \<10MB, tune spark.sql.autoBroadcastJoinThreshold)
  • 10-100× faster when applicable
  • AQE promotes to broadcast dynamically
Shuffle Hash Join

Both sides shuffled by join key.

  • Works for any size
  • N × M shuffle files
  • Slow for large joins; usually beaten by sort-merge
Sort-Merge Join

Both sides shuffled + sorted by join key.

  • Default for big × big
  • Spills gracefully to disk
  • Enables downstream range operations for free
Broadcast Nested Loop

The fallback that means you did something wrong.

  • No equi-join condition
  • O(N × M) — deadly at scale
  • Refactor the query if you see this

The Spark UI reading order

1
Jobs tab · pick the slow job

Look at duration and how many stages it has.

2
Stages tab · pick the slow stage

Sort by ‘shuffle read’ or ‘duration’. Focus on the biggest number.

3
Task metrics table

‘Max time’ vs ‘Median time’ — if max is 10× median, you have skew.

4
SQL / DataFrame tab

Physical plan tree. Look for BroadcastHashJoin (good) vs SortMergeJoin (fine) vs BroadcastNestedLoopJoin (bad).

5
Executors tab

GC time > 10 % = memory pressure; add executors or reduce partition size.


Common misconception
✗ What most people think

"Spark is slow on this dataset, so I need a bigger cluster. More executors means more parallelism means faster."

✓ What is actually true

Most Spark slowness is skew or shuffle, and neither is fixed by adding machines. If one partition holds 80% of the rows for a key, the stage finishes when that one task finishes — and a hundred idle executors do not help. Adding capacity fixes the case where every task is busy and there are not enough slots, which is the minority of real Spark performance problems.

Why the myth is so sticky

Because scaling out did work the first few times: early pipelines are genuinely CPU- or slot-bound, so more executors made them faster and the rule got learned. It fails silently later because the Spark UI shows the stage running, not that 199 of 200 tasks finished in seconds and one has been running for an hour. Wall-clock looks like "slow job", not "one hot key", unless you go look at the task duration distribution.

Prove it to yourself

Find the skew before touching cluster size — the task duration spread tells you immediately:

# 1. Look at key distribution on the join/group column
(df.groupBy('join_key').count()
   .orderBy('count', ascending=False)
   .show(20))
# a top key orders of magnitude above the rest = skew

# 2. In the Spark UI, open the slow stage and look at the task
#    duration summary: if max is far above the 75th percentile, it is skew,
#    not capacity. More executors will not change the max.

# 3. Fixes that actually work:
spark.conf.set('spark.sql.adaptive.enabled', 'true')
spark.conf.set('spark.sql.adaptive.skewJoin.enabled', 'true')  # AQE splits skewed partitions

# broadcast the small side to eliminate the shuffle entirely
from pyspark.sql.functions import broadcast
big.join(broadcast(small), 'key')

# or salt the hot key so it spreads across partitions
From first principles
Start with the question

Why does Spark split a job at shuffle boundaries specifically? Stages, tasks and shuffles look like three separate concepts — they are one.

  1. 1
    Transformations are either narrow (each output partition depends on exactly one input partition — map, filter) or wide (an output partition depends on many input partitions — groupBy, join, repartition).
    forced by · whether data must move between partitions is a property of the operation's semantics, not of its implementation
  2. 2
    A chain of narrow transformations can run entirely within one partition, on one machine, with no coordination and no network.
    forced by · all the data that output needs is already local
  3. 3
    So Spark fuses consecutive narrow transformations into a single task that streams a partition through all of them in one pass — no intermediate materialisation.
    forced by · materialising between each step would cost memory and I/O for no benefit when the data does not move
  4. 4
    A wide transformation breaks this: an output partition cannot begin until every upstream partition has contributed its share, which requires a global data exchange.
    forced by · you cannot know you have all rows for a key until every producer has finished emitting
  5. 5
    That exchange requires a barrier — all map-side tasks must complete and write their output before any reduce-side task can read it — and a barrier is exactly what a stage boundary is.
    forced by · partial input would produce partial and therefore wrong aggregates
⇒ Therefore

Therefore stages are not a scheduling convenience: a stage is precisely a maximal run of narrow transformations, and every stage boundary is a shuffle with a synchronisation barrier and a write-then-read of intermediate data to disk and across the network.

And note what this predicts: the number of stages in your job equals the number of shuffles plus one, so you can read shuffle count straight off the Spark UI without profiling anything. It also predicts why a straggler is so expensive — the barrier means the whole stage waits for the slowest task — and why broadcast joins are transformative: they convert a wide dependency into a narrow one, removing the barrier and the shuffle entirely.

Mental modelA plan of stages separated by barriers

Picture your job as a small number of stages laid end to end. Inside a stage, hundreds of independent tasks each stream one partition through a fused chain of operations, touching no network. Between stages sits a barrier: every task writes its shuffle output, everything stops, then the next stage reads.

All Spark performance work is one of three things: reduce the number of barriers (fewer shuffles), reduce the data crossing each barrier (filter and aggregate before shuffling), or make tasks finish at the same time (fix skew). Cluster size only helps a fourth, rarer case: not enough slots for the tasks you have.

  • Transformations are lazy; nothing executes until an action. This is what lets Catalyst reorder, push down predicates and prune columns — and why a bug may surface at an action far from the line that caused it.
  • Filter and project as early as possible. Every column and row dropped before a shuffle is bytes not crossing the network.
  • cache() only pays off when a DataFrame is used more than once and recomputation is expensive. Otherwise it is memory pressure and eviction for nothing.
  • Partition count should be a small multiple of total cores. Too few and cores idle; too many and per-task overhead and tiny output files dominate.
🔔 Fires when you see

Fire this model the moment you see: a stage where max task time is far above the median · a job that got slower after data grew, not proportionally but suddenly · thousands of tiny output files · repeated recomputation of the same DataFrame · an OOM on a single executor while the rest are idle · collect() on anything large · a join where one side would comfortably fit in memory.

The tradeoff

You are joining a very large fact table to a dimension table. Shuffle join, broadcast join, or pre-bucket both tables?

Shuffle hash / sort-merge join
+ you gain works at any size on both sides with no assumptions; the general-purpose default that will not fall over as the dimension grows
− you pay both sides are shuffled across the network, which is usually the single most expensive operation in the job, and it is where skew does its damage
pick when both sides are large, or the smaller side's size is unbounded or unknown — the safe default
Broadcast join
+ you gain eliminates the shuffle for the large side entirely: the small table is shipped to every executor and the join becomes a narrow, map-side operation. Frequently an order-of-magnitude improvement
− you pay the small side must fit in every executor's memory, multiplied by executor count for network cost; when a dimension quietly grows past the threshold, the job goes from fast to OOM with no gradual warning
pick when the small side is reliably small and bounded — a dimension with a known row count ceiling. Set the threshold deliberately rather than relying on the default estimate
Pre-bucketed / pre-sorted tables
+ you gain the shuffle is paid once at write time; every subsequent join on the bucket key is shuffle-free, and the benefit compounds across every job that uses those tables
− you pay both tables must be bucketed on the same key with the same bucket count, which is a schema-level commitment; changing the key or count means rewriting the data; and it helps only joins on that key
pick when the same large-to-large join runs repeatedly on a schedule — a daily fact-to-fact join is the canonical case
What a senior engineer actually does

Broadcast whenever you can prove the small side is bounded, and prove it rather than assume it — check the actual size and set autoBroadcastJoinThreshold explicitly. The classic production incident is a dimension that was small when the job was written, crossed the threshold months later, and turned a reliable job into an intermittent OOM that nobody connects to a table growing.

For repeated large-to-large joins, bucketing is the highest-leverage change available because it moves the shuffle from every run to one run. And enable Adaptive Query Execution — it re-plans using actual runtime statistics rather than estimates, which is exactly the fix for the cardinality-estimation problem, and it handles skew splitting automatically.


(c) Hands-on · 25 min

A single script that demonstrates lazy evaluation, shuffle, broadcast join, skew, and salting. Save as spark_demo.py.

# spark_demo.py — reproduce the ‘six things Spark actually does’ in one script.
# Run with: spark-submit --master local[4] spark_demo.py
from pyspark.sql import SparkSession, functions as F, Window as W
from pyspark.sql.types import StructType, StructField, StringType, IntegerType
 
spark = (
    SparkSession.builder
    .appName("spark_demo")
    .config("spark.sql.adaptive.enabled", "true")
    .config("spark.sql.adaptive.skewJoin.enabled", "true")
    .config("spark.sql.autoBroadcastJoinThreshold", 10 * 1024 * 1024)  # 10 MB
    .config("spark.sql.shuffle.partitions", "8")   # low for demo readability
    .getOrCreate()
)
spark.sparkContext.setLogLevel("WARN")
 
# ------------------------------------------------------------------
# 1. Lazy evaluation — no work happens until an action
# ------------------------------------------------------------------
orders = spark.range(0, 1_000_000).selectExpr(
    "id AS order_id",
    "cast(rand(42) * 100 AS int) AS customer_id",
    "cast(rand(1) * 1000 AS decimal(10,2)) AS amount",
    "cast(rand(2) * 5 AS int) AS status_code"
)
# Nothing has executed. Physical plan is only computed on action.
orders.explain(mode="formatted")   # <- lazy print
 
# ------------------------------------------------------------------
# 2. Trigger an action — SEE the DAG light up in the UI at localhost:4040
# ------------------------------------------------------------------
print("row count:", orders.count())
 
# ------------------------------------------------------------------
# 3. Shuffle: groupBy triggers a hash exchange
# ------------------------------------------------------------------
by_cust = orders.groupBy("customer_id").agg(F.sum("amount").alias("total"))
by_cust.explain()
by_cust.show(5)
 
# ------------------------------------------------------------------
# 4. Broadcast join — small dim table shipped to every executor
# ------------------------------------------------------------------
customers = spark.createDataFrame(
    [(i, f"cust_{i}", ["Hyd","Blr","Del","Mum"][i % 4]) for i in range(100)],
    "customer_id INT, name STRING, city STRING"
)
joined = orders.join(F.broadcast(customers), "customer_id")
joined.explain()   # should show BroadcastHashJoin
print("joined:", joined.count())
 
# ------------------------------------------------------------------
# 5. Skew — 90 % of orders go to customer 7. Watch the stage.
# ------------------------------------------------------------------
skewed = orders.withColumn(
    "customer_id",
    F.when(F.rand(seed=99) < 0.9, F.lit(7)).otherwise(F.col("customer_id"))
)
skewed.groupBy("customer_id").count().show(5)
 
# ------------------------------------------------------------------
# 6. Salting — the classic fix for skewed group / join
# ------------------------------------------------------------------
SALT = 16
salted = (
    skewed
    .withColumn("salt", (F.rand(seed=101) * SALT).cast("int"))
    .withColumn("customer_id_salted", F.concat_ws("_", "customer_id", "salt"))
)
pre = salted.groupBy("customer_id_salted", "customer_id").agg(F.count("*").alias("n"))
# Roll up salted keys back
final = pre.groupBy("customer_id").agg(F.sum("n").alias("n_total"))
final.orderBy(F.desc("n_total")).show(5)
 
spark.stop()

What each block does

Anatomy of the script

adaptive.enabled + skewJoin.enabled
AQE is off in ancient Spark; turn it on explicitly. Skew join splits large partitions dynamically.
config
autoBroadcastJoinThreshold=10MB
Any dim table under 10 MB is auto-broadcast without you writing F.broadcast(). Bump this if you know you have small dims up to ~100 MB.
config
spark.sql.shuffle.partitions=8
Default is 200 — too many for a laptop demo. In prod, tune to ~cores × 2-3 or leave AQE to coalesce.
config
orders.explain()
Prints logical + physical plan. No actual computation. Great for verifying broadcast joins before you run.
debug
F.broadcast(customers)
Explicit hint. Only use when you know the table is small AND the optimizer doesn't figure it out.
hint
Salting pattern
Two-phase agg: (1) group by (key, salt) to spread load, (2) group by key to combine. Standard fix for hot-key skew.
skew-fix
Try itFeel skew and fix it

Run the script, then open http://localhost:4040. Look at the stage for the skewed groupBy. In Summary Metrics for Tasks, find Duration:

  • Median ~200 ms, Max ~5 s → skew ratio of 25×. That single task is the bottleneck; adding executors won't help.
  • After salting, the ratio drops to <2×.

That single UI reading is 80 % of Spark debugging.

💡 Hint · Open http://localhost:4040 in a browser BEFORE spark-submit finishes. Click Stages, click the slow one, look at ‘Summary Metrics for Tasks’ — the Max vs 75th percentile ratio IS your skew number.

(d) Production reality · 15 min

War story Netflix· 2019500 TB Iceberg tables
🔥 What broke

Netflix had petabyte-scale Spark jobs on Parquet + Hive that took 6 hours nightly. Every schema change required rewriting years of partitions. Skew on ‘big customer’ ids caused sporadic 10-hour tails.

🧯 The fix

Adopted Apache Iceberg (which they open-sourced): hidden partitioning, atomic snapshot writes, partition-evolution. Combined with Spark 3.0 AQE for automatic skew handling. Nightly job dropped from 6 h to under 90 min and became recoverable at any partition.

🎓 Lesson to steal
Spark performance is 30 % engine and 70 % table format. Iceberg / Delta / Hudi are the modern default; plain Parquet-on-S3 is legacy.
Post-mortem
War story Uber· 2020Michelangelo ML platform
🔥 What broke

Feature-engineering Spark job joined a 5 TB events table with a 200 GB user-features table. Job ran 4 h. Investigation: the user-features table had 3 users (out of 100 M) with 10 M events each. Sort-merge join sent one reducer 10 M rows and the rest 100 K rows. Classic skew.

🧯 The fix

Two-phase salted join: added a random 1..N salt to the events table, replicated the top-3 hot user rows N times, joined on (user_id, salt). Runtime dropped to 25 minutes. Spark 3.0 AQE now does this automatically for detected skewed keys — Uber contributed the technique upstream.

🎓 Lesson to steal
Skew is the #1 Spark performance problem. AQE handles it in 3.0+; before 3.0 you salt manually.
Post-mortem
War story Airbnbreported in ‘Data Infrastructure at Airbnb’ post
🔥 What broke

Analysts wrote df.repartition(10000) to ‘speed things up’. Result: 10 000 tiny tasks, each processing a few KB, dominated by task-launch overhead. Job slower than the un-repartitioned version.

🧯 The fix

Trained analysts on partition sizing (target ~128-256 MB per partition post-shuffle). AQE's coalescePartitions now handles this automatically when the true output is small.

🎓 Lesson to steal
‘More parallelism’ isn't free. Each task has ~50-100 ms of scheduling overhead. Aim for 100-1000 tasks per stage, not 10 000+.

Where this shows up in the rest of the plan

Spark is the compute layer of everything data
S048 · Kafka
Structured Streaming ingests from Kafka.
S049 · Stream processing
Same DataFrame API, unbounded input.
S050 · Airflow
Orchestrates Spark jobs on schedules.
S051 · dbt
dbt on Spark = SQL-only alternative to raw PySpark.
S078 · Consistency & consensus
Iceberg's transactional commit is a distributed-consensus primitive.
S099 · Feature stores
Spark computes offline features overnight; online store serves them at request time.

(e) Recall + stretch · 10 min

Recall — click each to reveal · click to reveal
★ = stretch question

Explain-out-loud test

  1. Lazy evaluation — one sentence + why it matters.
  2. What is a shuffle — including the file-count math.
  3. How to spot skew in the UI — the one number you look at.

What comes next

Hub: The 6-Month Learning Plan


Part of a 130-session evergreen learning series. Session structure: intuition → visual → hands-on → production war stories → recall. Duration: 90 minutes.