Combining Before the Shuffle

The cheapest shuffle byte is the one never written. For count, sum, min, max, avg and most other built-in aggregates, Spark computes the result in 2 passes. A partial pass runs inside each map task before the shuffle write and collapses that task's rows to 1 row per key it saw, carrying a running result. A final pass runs in the reducer after the read and merges each key's partial rows. Every DataFrame groupBy using these aggregates gets the partial pass without being asked.

Both outputs give the reducer what it needs: it adds US 4 to the partial US counts from every other map task and reaches the same total. The partial version writes 2 rows where the raw version writes 6. For avg, the partial row carries a sum and a count, since an average of averages would be wrong.

What decides the saving

A map task can write at most 1 row per distinct key in its own partition, so the saving rests on a single ratio: rows per key within 1 task. Take a year of 1 billion plays read as 400 partitions, 2.5 million rows per map task, and group them 3 ways.

The same groupBy code cuts the shuffle by a factor of about 14,000 for country, 50 for artist, and almost nothing for listener: a task holding 2.5 million rows in arrival order sees nearly every listener once. Layout matters as much as the key. Had the files been written so each partition held a contiguous range of listeners, each task would see all of a listener's plays together, around 10 on average, and the partial pass would cut the shuffle about 10 times.

Two copies of the 1 billion plays each sit in 400 partitions. Copy A is in arrival order. Copy B was written so each partition holds a contiguous range of listener_id values. Which is true of groupBy("listener_id").count()?

• Both shuffle the same rows, since the partial pass depends only on how many distinct keys exist
• Copy A shuffles fewer rows, since arrival order spreads each listener's plays across many tasks
• Copy B shuffles fewer rows, since each task sees a listener's plays together and writes 1 row for them
• Neither copy combines, since count is computed only after the reducers have read every row

Aggregates that combine poorly

What the partial row holds

• count, sum, min, max and avg: a fixed few numbers per key, however many rows fed it.
• collect_list and collect_set: the values themselves, so rows collapse while the bytes stay close to the raw values.
• countDistinct: only repeated (key, value) pairs can be dropped, so Spark shuffles twice, once on key and value, then once on key.
• approx_count_distinct: a fixed-size sketch, so the job shuffles once, with a default relative standard deviation of 5%.

exact = (plays
         .groupBy("artist_id")
         .agg(F.countDistinct("listener_id").alias("listeners")))

approx = (plays
          .groupBy("artist_id")
          .agg(F.approx_count_distinct("listener_id").alias("listeners")))

Run as written, exact takes 2 shuffles and 3 stages: the first shuffle groups by (artist_id, listener_id) so duplicate pairs meet and drop, and the second groups the surviving pairs by artist_id to count them. approx takes 1 shuffle and 2 stages, and each map task writes 1 sketch per artist it saw, the same size whatever the listener count.

Check the ratio on real data

The cell below counts rows and distinct values for 3 columns of the practice transactions table. Run it, then divide rows by each key count. The quotient is the most a partial pass could shrink a groupBy on that column if 1 task saw every row; how the rows are split across tasks decides how close a real job gets.

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.