Partitions, shuffles, and performance

⏱ 35 min

Hook

Your job reads a clean 10 GB table, groups it, and writes a summary — and it moves 40 GB across the network to do it, four times the size of the input. Worse, the same DataFrame gets recomputed from scratch three times because you referenced it in three places. Both problems are invisible in the code and obvious in the plan. This lesson teaches you to read that plan and fix both.

Concept

A shuffle is Spark redistributing rows across the cluster so that rows sharing a key end up on the same partition. It is the physical event a wide transformation triggers, and it is the single most expensive thing a batch job normally does: rows are serialized, written to disk, sent over the network, and read back. Every groupBy and every join pays for one.

You can see exactly where it happens. The explain() method prints a DataFrame's physical plan — the actual steps Spark will run — and a shuffle appears in that plan as an Exchange node. Reading plans bottom-up (Spark executes from the leaves toward the root), an Exchange marks the boundary where one stage ends and the reshuffled data feeds the next. Find the Exchange, and you have found your shuffle.

Knowing where the shuffle is, you have three main levers to make a job faster.

Shrink the data before the shuffle. A shuffle's cost scales with the bytes it moves, so anything that makes rows fewer or narrower before the wide step pays off directly. Filter out rows you will not aggregate; select only the columns the aggregation or join needs. Because Catalyst can push some of these down automatically, writing them explicitly and early both helps the optimizer and makes your intent clear. This is why Lesson 2 insisted on ordering narrow steps first.

Broadcast a small side of a join. A normal join shuffles both DataFrames. But if one side is small enough to fit in memory, Spark can send a full copy of it to every executor and join without shuffling the big side at all — a broadcast join. You request it with F.broadcast(small_df), and it turns a two-sided shuffle into no shuffle for the large table. Use it when one side is a small dimension (a lookup of a few thousand rows), not when both sides are large.

Cache a DataFrame you reuse. Because of lazy evaluation, a DataFrame is recomputed from its source every time an action touches it. If you run three actions on the same expensive DataFrame — or reference it in three branches — Spark redoes the whole chain three times. Calling .cache() tells Spark to keep the computed partitions in memory after the first action, so later actions read the stored result instead of recomputing. Cache only what you genuinely reuse: caching a DataFrame you touch once wastes memory for no gain.

Lever	What it does	When to reach for it
Filter and project early	Fewer, narrower rows enter the shuffle	Always, before any wide step
Broadcast join	Removes the shuffle on the large side	One join side is small
Cache	Stores a result so it is not recomputed	A DataFrame is reused by 2+ actions

warning

Caching is not free memory insurance. A cached DataFrame that does not fit in executor memory spills to disk and can make a job slower, not faster, while starving other work of memory. Cache deliberately, and only what you reuse.

🧠 Knowledge check

1. You run df.explain() and see an Exchange node in the physical plan. What does it tell you?

a) The DataFrame has been cached in memoryb) A shuffle happens there — data is redistributed across the cluster, marking a stage boundaryc) The query failed to compiled) Spark is reading the data from disk for the first time

2. A join matches a 200-million-row fact table against a 3,000-row lookup table. What is the cheapest sound strategy?

a) Cache the fact table before the joinb) Broadcast the small lookup table so the large fact table is not shuffledc) groupBy the fact table first to shrink itd) Broadcast the fact table to every executor

Worked example

Let me find a shuffle in a real plan and then cut what it moves. I will aggregate web-session events by country, and use explain() to prove where the shuffle is.

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.appName("shuffles").getOrCreate()

sessions = spark.createDataFrame(
    [
        ("s-1", "IN", "mobile", 5, 1),
        ("s-2", "IN", "web", 2, 0),
        ("s-3", "US", "mobile", 8, 1),
        ("s-4", "US", "web", 3, 1),
        ("s-5", "IN", "mobile", 1, 0),
    ],
    ["session_id", "country", "channel", "page_views", "converted"],
)

by_country = sessions.groupBy("country").agg(
    F.sum("converted").alias("conversions")
)
by_country.explain()

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[country], functions=[sum(converted)])
   +- Exchange hashpartitioning(country, 200)
      +- HashAggregate(keys=[country], functions=[partial_sum(converted)])
         +- LocalTableScan [session_id, country, channel, page_views, converted]

Read it bottom-up. Spark scans the data, does a partial_sum per partition (a partial aggregate that needs no data movement), then hits the Exchange hashpartitioning(country, 200) — that is the shuffle, repartitioning rows by country so each country's partial sums meet — and finally the top HashAggregate finishes the totals. The Exchange is the expensive line.

Now shrink what it moves. Suppose a later step only needs conversions, and only for mobile sessions. Filtering and projecting before the groupBy means the shuffle carries a fraction of the columns and rows:

lean = (
    sessions
    .filter(F.col("channel") == "mobile")
    .select("country", "converted")
    .groupBy("country")
    .agg(F.sum("converted").alias("conversions"))
)
lean.show()

+-------+-----------+
|country|conversions|
+-------+-----------+
|     IN|          1|
|     US|          1|
+-------+-----------+

The Exchange is still there — a groupBy cannot avoid its shuffle — but it now moves two columns of mobile rows instead of five columns of every row. If I then ran several actions on lean, I would call lean.cache() once so Spark computes that shuffle a single time rather than on every action.

Hands-on

Your turn, on a different pipeline. The exercise below gives you a two-stage job over clickstream data with a groupBy and a join against a large table, plus a tiny country-lookup table. Use explain() to locate the Exchange nodes, then reduce the shuffled data by projecting early and broadcasting the lookup.

▶ Spark Forgede202-cut-the-shuffle

Open in Spark Forge ↗

Success criteria: you point to the Exchange node for the groupBy, rewrite the job so the large table is no longer shuffled for the join, and produce the same result with less data moved. Spark Forge checks both your revised plan and your output.

Recap

• You can define a shuffle as Spark redistributing rows by key across the cluster, and name it as the main cost of any wide operation.
• You can locate a shuffle in a physical plan by finding the Exchange node in explain() output.
• You can cut shuffled data by filtering and projecting before the wide step, and remove a join's large-side shuffle with a broadcast join.
• You can cache a DataFrame that multiple actions reuse so Spark computes it once instead of every time.

Next up in the full course: writing these tuned jobs against real Parquet files with partitioning and schema handling — and the lab below has you build the first end-to-end version now.

🧠 Knowledge check

1. You reference the same expensive DataFrame in three separate show() calls and the job is slow. What is the direct fix?

a) Add more columns to the select so Spark does less workb) Call .cache() on the DataFrame so the computed result is reused instead of recomputed for each actionc) Replace the groupBy with a filterd) Broadcast the DataFrame to the driver

2. Which change reduces the bytes a groupBy shuffle moves, without changing the result?

a) Selecting only the grouping key and the aggregated column, and filtering unneeded rows, before the groupByb) Adding a withColumn after the groupByc) Calling collect() instead of show()d) Increasing the number of output columns

📌 Key takeaways

• A shuffle redistributes rows by key across the cluster and is the main cost of every wide operation.
• Find a shuffle in a physical plan as the Exchange node in explain() output.
• Shrink the shuffle by filtering and projecting early; remove a join's large-side shuffle by broadcasting the small side.
• Cache only a DataFrame that multiple actions reuse — never one you touch once.