Reusing Shuffle Output
Spark skips a shuffle in 2 situations. Inside 1 query, it plans none when the data already sits partitioned the way the next operation needs. Across actions, it skips a map stage whose shuffle files are still on disk from an earlier job.
Partitioning already in place
A groupBy("region") needs every row of a region in 1 partition. After repartition("region") that already holds: the repartition hashed every row by region. Before adding a shuffle, Spark compares what each operation needs with how its input is partitioned, and here they match, so the aggregation runs inside the repartition's reduce tasks with no second shuffle.
The animation repartitions 9 orders by region into 3 partitions, then counts orders per region. Count the shuffles as it plays: there is 1, and each count is computed in the task where its region's rows landed.
orders.repartition("region").groupBy("region").count().collect()| order_id | region |
|---|---|
| 1 | US |
| 2 | EU |
| 3 | US |
| 4 | APAC |
| 5 | EU |
| 6 | LATAM |
| 7 | US |
| 8 | APAC |
| 9 | EU |
What counts as a match
| Code | Shuffles | Why |
|---|---|---|
| `repartition("region").groupBy("region")` | 1 | Hash on region already gathers each region |
| `repartition(50, "region").groupBy("region")` | 1 | Same match, and the groupBy inherits the 50 partitions |
| `repartition("region").groupBy("region", "status")` | 1 | A region's rows sit together, so each region and status pair does too |
| `repartition("region", "status").groupBy("region")` | 2 | Status spreads 1 region over several partitions |
| `repartition(50).groupBy("region")` | 2 | Round-robin put every region everywhere |
| `groupBy("region").count().groupBy("region").max("count")` | 1 | The first groupBy left the rows hashed on region |
| `repartition("region").orderBy("region")` | 2 | A sort needs ranges, which a hash cannot provide |
Spark's test is whether the existing layout puts all rows that share the new grouping key in 1 partition. Hashing on a subset of the grouping columns passes it, hashing on extra columns fails it, and round-robin always fails. When the layout passes, its partition count carries forward, so repartition(50, "region") sets the task count of everything after it up to the next shuffle.
Shuffle files that outlive their job
Each action re-runs its lineage from the source, as lazy evaluation showed, with 1 exception. After a map stage finishes, its files stay on the executors' disks and the driver keeps their locations. When a later job in the same application needs that exact shuffle, the scheduler finds the map output complete and skips the stage: the Jobs tab lists it as skipped, and its tasks never run.
by_artist = (plays
.groupBy("artist_id")
.agg(F.sum("ms_played").alias("ms")))
first = by_artist.collect()
again = by_artist.collect()The first collect() runs 400 map tasks that read plays and write partial sums, then 200 reduce tasks. The second runs the same query, finds all 400 map outputs registered, skips stage 0, and runs only the 200 reducers, so plays is never read a second time.
Where the reuse ends
The reuse belongs to 1 query. by_artist.filter(F.col("ms") > 0).collect() is a new query that plans its own shuffle, so both stages run again even though the rows feeding the shuffle are identical. The files are also temporary: Spark deletes them some time after the driver drops its last reference to the query that wrote them, and an executor that dies usually takes its files with it. Reuse you control on purpose comes from caching, covered in a later lesson.
Which of these 3 jobs run the 400 map tasks that read plays?
totals = plays.groupBy("artist_id").agg(F.sum("ms_played").alias("ms"))
a = totals.collect()
b = totals.collect()
c = totals.filter(F.col("ms") > 0).collect()- Only the first job; the second and third reuse its shuffle files
- All 3, since shuffle files are deleted when the job that wrote them ends
- The first and the third; the second skips its map stage
- The first and the second; the third reads earlier files through the filter
Capstone: the whole shuffle at once
With spark.sql.shuffle.partitions at 200, how many shuffles does this job run, and how many tasks run the aggregation?
daily = (plays
.repartition(100, "listener_id")
.groupBy("listener_id", "day")
.agg(F.sum("ms_played").alias("ms"))
.orderBy("ms"))
daily.write.parquet("/out/daily")- Two shuffles, and the aggregation runs inside the repartition's 100 tasks
- Three shuffles, and the aggregation runs as 200 tasks of its own
- Two shuffles, and the aggregation runs as 200 tasks after the sort
- One shuffle, since the sort reuses the hash layout on listener_id
The last question joins the partial pass from section 2 to the layout reuse in this one. A repartition has no aggregate to apply, so its map tasks write every row they receive, as they read it. Whatever runs after it, partial pass included, runs in the repartition's reduce tasks, on rows that have already crossed the network.
A teammate adds repartition("country") before groupBy("country").count() on the 1 billion plays in 400 partitions, so the groupBy adds no shuffle. Compared with the groupBy alone, how many rows cross the shuffle?
- The same, at most 72,000, because there is still exactly 1 shuffle
- Half as many, because the groupBy and the repartition share 1 shuffle
- None, because the groupBy reads the repartition's files where they sit
- All 1 billion, because the partial count now runs only after the shuffle
About this lesson
This section is part of the Shuffles: Intermediate lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.