Eliminating a Shuffle

Concepts covered: Broadcast Join

The cheapest shuffle is the one the plan never contains. A wide operation needs every row sharing a key inside 1 task, and Spark adds a shuffle only when its input isn't already laid out that way. That leaves 2 ways to remove one: drop a wide operation the result doesn't need, or hand the next wide operation input that already meets its requirement.

Delete the shuffle nobody asked for

Pattern	Shuffles as written	After the fix	Why it was waste
`repartition(400)` right before `groupBy("user_id")`	2	1	the round-robin layout is discarded by the hash shuffle the groupBy adds
`orderBy("ts")` before a write no reader consumes in order	1	0	the range shuffle and its extra stage buy an order nobody uses
`select("order_id").distinct()` on a column you know is unique	1	0	Spark cannot know the column is unique, so it shuffles to prove it

Each of these survives code review because it looks harmless: a repartition left over from debugging, a sort added so a sample looked tidy, a distinct added out of caution. Search the code for wide operations whose output the next step reshuffles or ignores.

How many shuffles does this program run, and which one can you delete without changing the rows it returns?

result = (orders
    .repartition(50)
    .filter(F.col("status") == "Completed")
    .groupBy("status")
    .agg(F.sum("profit").alias("profit"))
    .orderBy(F.col("profit").desc()))

result.collect()

Two shuffles; delete the groupBy's, because the orderBy already moves rows by profit
Three shuffles; delete the round-robin one from repartition(50), which the groupBy redoes
Three shuffles; delete the orderBy's, because the groupBy output arrives in profit order
One shuffle; the filter, groupBy and orderBy all run in a single stage after it

Hand over input that is already arranged

Spark records how each DataFrame is partitioned. After repartition("customer_id"), every row for a customer sits in 1 partition, so a groupBy("customer_id") that follows finds its requirement met and adds no shuffle of its own. The board plays exactly that: 2 wide operations in the code, rows exchanged once.

(orders
   .repartition("customer_id")
   .groupBy("customer_id")
   .agg(F.sum("amount").alias("spend"),
        F.count("order_id").alias("order_count"))
   .collect())

orders
order_id	customer_id	amount
1	c1	40
2	c2	25
3	c1	60
4	c3	35
5	c2	50
6	c1	20

The driver collects c1 at 120 over 3 orders, c2 at 75 over 2, and c3 at 35 over 1. Look at what crossed the shuffle: all 6 raw orders, because a repartition moves rows before any aggregation runs. A plain groupBy would have shuffled partial sums. Repartitioning by a key first pays off only when 2 or more wide operations then share it.

One shuffle for 2 aggregations

A hash layout on user_id also satisfies a grouping on user_id plus other columns, since rows that share a (user_id, event_date) pair already share a user_id. The job below counts each user's events per day, then keeps each user's busiest day.

per_day = (events
    .groupBy("user_id", "event_date")
    .agg(F.count("*").alias("n")))

busiest = (per_day
    .groupBy("user_id")
    .agg(F.max("n").alias("busiest_day_events")))

per_day = (events
    .repartition("user_id")
    .groupBy("user_id", "event_date")
    .agg(F.count("*").alias("n")))

busiest = (per_day
    .groupBy("user_id")
    .agg(F.max("n").alias("busiest_day_events")))

Plan	Shuffles	Rows that cross
As written	2	partial counts per (user, day) per map task, then partial maxima per user
Repartitioned	1	every raw event

Suppose a user-day averages 40 events and a user's events for a day tend to land in the same map task. The first plan's first shuffle then carries as little as a fortieth of the rows in the second plan's only shuffle. Fewer shuffles is the goal only when the bytes agree, so price both plans with the method from the previous section before choosing.

The drill hands you a program that already repartitions the seed's order_items by order_id. You choose the grouping columns and aggregates for both steps so the result is each order's largest per-product quantity, with both aggregations reusing that single shuffle.

About this lesson

This section is part of the Shuffles: Advanced 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.