Leverage PartialReduce AggregationExec mode to drastically reduce shuffle size

[fmonjalet]

Summary

When DataFusion 53 is released, we should use the PartialReduce mode of AggregateExec (apache/datafusion#20019, kudos to @njsmith) for maximum local data reduction before shuffling.

The expectation is reducing the amount of data exchanged by up to a factor of n_partition_per_task (for example 16x).

Details

Very schematically, here is how a shuffle happens:

My understanding is that we will send duplicate aggregate keys on the network, where we could actually have reduced them to one line locally before the shuffle.

We should leverage PartialReduce to further reduce data before shuffling:

In this example, we'd reduce data exchanged by 4x (because there are 4 partitions per task).

No strong opinion on whether we should add another RepartitionExec layer like in this graph, or if we should run the AggregateExec(PartialReduce) on the "ready to shuffle" partitions. Since shuffles can have hundreds of partitions, I don't know what is the comparative cost of:

• RepartitionExec(Hash, output_partitions=16) + AggregateExec(PartialReduce) on 16 partitions
• vs AggregateExec(PartialReduce) on potentially hundreds of partitions

[nchu05]

Take

[nchu05]

Design Document

Approach 1: Add repartition + partial reduce

Plan shape

RepartitionExec(Hash([group_keys], N×T))   ← original, stays in place
  AggregateExec(mode=PartialReduce)        ← NEW
    RepartitionExec(Hash([group_keys], M)) ← NEW inner, merges S partitions by key
      AggregateExec(mode=Partial)
        ...

Idea

Insert a smaller RepartitionExec below PartialReduce to hash-partition by group key, so PartialReduce sees all copies of the same key in one partition and can merge them. The existing RepartitionExec stays above, re-expanding into the shuffle partitions.

Implementation

• Add a post-processing optimizer pass that:
  • Finds NetworkShuffleExec (network boundary) nodes.
  • Matches RepartitionExec(Hash) → AggregateExec(mode=Partial) underneath.
  • Inserts:
    • A new inner RepartitionExec(Hash, M) wrapping AggregateExec(Partial) (using same hash expressions, fewer partitions).
    • An AggregateExec(mode=PartialReduce) above the new inner repartition, with group expressions rebuilt using positional indices (0..N-1).
  • Rebuilds the original RepartitionExec above PartialReduce, preserving its original partitioning.
  • Replaces the network boundary’s child via with_new_children.
• Run this pass before batch coalescing.
• Add module + wiring in the distributed planner.

---

Approach 2: Add partial reduce only

Plan shape

RepartitionExec(Hash(group_keys, N×T))  ← existing, stays in place
  AggregateExec(PartialReduce)          ← new
    AggregateExec(Partial)                                                                                                                         

Idea

Insert PartialReduce between the existing RepartitionExec and Agg(Partial) with no additional repartition. Each partition is reduced independently.

Implementation

• Add a post-processing optimizer pass that:
  • Finds NetworkShuffleExec nodes.
  • Matches RepartitionExec → Aggregate(Partial) underneath.
  • Inserts a PartialReduce below the repartition.
• Run this pass before batch coalescing.
• Add module + wiring in the distributed planner.

[nchu05]

Local Benchmarking

Approach 1 is optimal. I created a new run script that compared the bytes transferred for queries with large cardinality. This is done by walking the rewritten plan tree and summing bytes_transferred from all NetworkShuffleExec and NetworkCoalesceExec nodes.

Approach 1 vs Baseline - Approach 1 Shows Significant Memory Improvement

Comparing tpch_sf1 results from branch 'leverage-partialreduce-approach1' [prev] with 'datafusion53' [new]
 tpch_sf1 bytes_q1: prev=  31 ms, new=  31 ms, diff=1.00 slower ✖, bytes=603K→838K (0.72x)
 tpch_sf1 bytes_q2: prev=  68 ms, new=  73 ms, diff=1.07 slower ✖, bytes=7661K→12414K (0.62x)
 tpch_sf1 bytes_q3: prev= 108 ms, new=  99 ms, diff=1.09 faster ✔, bytes=35536K→35475K (1.00x)
 tpch_sf1 bytes_q4: prev=  37 ms, new=  39 ms, diff=1.05 slower ✖, bytes=1417K→2255K (0.63x)
 tpch_sf1 bytes_q5: prev=  23 ms, new=  25 ms, diff=1.09 slower ✖, bytes=3981K→6119K (0.65x)
 tpch_sf1 bytes_q6: prev=   8 ms, new=   9 ms, diff=1.12 slower ✖, bytes=313K→368K (0.85x)
 tpch_sf1 bytes_q7: prev=  59 ms, new=  46 ms, diff=1.28 faster ✅, bytes=10896K→14850K (0.73x)
 tpch_sf1 bytes_q8: prev=   2 ms, new=   2 ms, diff=1.00 slower ✖
 tpch_sf1 bytes_q9: prev=   8 ms, new=   8 ms, diff=1.00 slower ✖
 tpch_sf1 bytes_q10: prev=   7 ms, new=   7 ms, diff=1.00 slower ✖, bytes=482K→637K (0.76x)
 tpch_sf1 bytes_q11: prev=  17 ms, new=  15 ms, diff=1.13 faster ✔, bytes=3545K→3535K (1.00x)
 tpch_sf1 bytes_q12: prev=  23 ms, new=  23 ms, diff=1.00 slower ✖, bytes=220K→217K (1.02x)
 tpch_sf1 bytes_q13: prev=  45 ms, new=  47 ms, diff=1.04 slower ✖, bytes=3164K→5198K (0.61x)
 tpch_sf1 bytes_q14: prev=   1 ms, new=   1 ms, diff=1.00 slower ✖
 tpch_sf1 bytes_q15: prev=   2 ms, new=   2 ms, diff=1.00 slower ✖
 tpch_sf1 bytes_q16: prev=   9 ms, new=   9 ms, diff=1.00 slower ✖
 tpch_sf1 bytes_q17: prev=   0 ms, new=   0 ms, diff=NaN slower ✖
 tpch_sf1 bytes_q18: prev= 162 ms, new= 213 ms, diff=1.31 slower ❌, bytes=36689K→90263K (0.41x)
 tpch_sf1 bytes_q19: prev=  13 ms, new=  14 ms, diff=1.08 slower ✖, bytes=750K→1089K (0.69x)
 tpch_sf1 bytes_q20: prev=  21 ms, new=  21 ms, diff=1.00 slower ✖, bytes=319K→376K (0.85x)
 tpch_sf1 bytes_q21: prev=  32 ms, new=  32 ms, diff=1.00 slower ✖, bytes=953K→1449K (0.66x)
 tpch_sf1 bytes_q22: prev=  10 ms, new=  10 ms, diff=1.00 slower ✖, bytes=274K→302K (0.91x)

Baseline vs Approach 2 - No Improvement

Comparing tpch_sf1 results from branch 'datafusion53' [prev] with 'leverage-partialreduce-approach2' [new]
tpch_sf1 bytes_q1: prev=  32 ms, new=  32 ms, diff=1.00 slower ✖, bytes=838K→840K (1.00x)
tpch_sf1 bytes_q2: prev=  69 ms, new=  75 ms, diff=1.09 slower ✖, bytes=12417K→12416K (1.00x)
tpch_sf1 bytes_q3: prev=  97 ms, new= 107 ms, diff=1.10 slower ✖, bytes=35476K→35475K (1.00x)
tpch_sf1 bytes_q4: prev=  37 ms, new=  39 ms, diff=1.05 slower ✖, bytes=2254K→2253K (1.00x)
tpch_sf1 bytes_q5: prev=  23 ms, new=  26 ms, diff=1.13 slower ✖, bytes=6119K→6123K (1.00x)
tpch_sf1 bytes_q6: prev=   8 ms, new=   8 ms, diff=1.00 slower ✖, bytes=368K→370K (0.99x)
tpch_sf1 bytes_q7: prev=  55 ms, new=  54 ms, diff=1.02 faster ✔, bytes=14849K→12925K (1.15x)
tpch_sf1 bytes_q8: prev=   2 ms, new=   2 ms, diff=1.00 slower ✖
tpch_sf1 bytes_q9: prev=   8 ms, new=   9 ms, diff=1.12 slower ✖
tpch_sf1 bytes_q10: prev=   7 ms, new=   8 ms, diff=1.14 slower ✖, bytes=637K→639K (1.00x)
tpch_sf1 bytes_q11: prev=  15 ms, new=  16 ms, diff=1.07 slower ✖, bytes=3535K→3537K (1.00x)
tpch_sf1 bytes_q12: prev=  22 ms, new=  23 ms, diff=1.05 slower ✖, bytes=217K→218K (0.99x)
tpch_sf1 bytes_q13: prev=  45 ms, new=  47 ms, diff=1.04 slower ✖, bytes=5198K→5198K (1.00x)
tpch_sf1 bytes_q14: prev=   1 ms, new=   1 ms, diff=1.00 slower ✖
tpch_sf1 bytes_q15: prev=   2 ms, new=   2 ms, diff=1.00 slower ✖
tpch_sf1 bytes_q16: prev=   9 ms, new=   9 ms, diff=1.00 slower ✖
tpch_sf1 bytes_q17: prev=   0 ms, new=   0 ms, diff=NaN slower ✖
tpch_sf1 bytes_q18: prev= 171 ms, new= 178 ms, diff=1.04 slower ✖, bytes=90263K→53341K (1.69x)
tpch_sf1 bytes_q19: prev=  14 ms, new=  14 ms, diff=1.00 slower ✖, bytes=1089K→1091K (1.00x)
tpch_sf1 bytes_q20: prev=  21 ms, new=  21 ms, diff=1.00 slower ✖, bytes=376K→378K (0.99x)
tpch_sf1 bytes_q21: prev=  32 ms, new=  33 ms, diff=1.03 slower ✖, bytes=1448K→1449K (1.00x)
tpch_sf1 bytes_q22: prev=  10 ms, new=  10 ms, diff=1.00 slower ✖, bytes=302K→304K (0.99x)

Approach 1 vs Approach 2 - Approach 1 is significantly better

Comparing tpch_sf1 results from branch 'leverage-partialreduce-approach1' [prev] with 'leverage-partialreduce-approach2' [new]
 tpch_sf1 bytes_q1: prev=  30 ms, new=  31 ms, diff=1.03 slower ✖, bytes=603K→840K (0.72x)
 tpch_sf1 bytes_q2: prev=  70 ms, new=  74 ms, diff=1.06 slower ✖, bytes=7662K→12420K (0.62x)
 tpch_sf1 bytes_q3: prev= 115 ms, new= 103 ms, diff=1.12 faster ✔, bytes=35537K→35478K (1.00x)
 tpch_sf1 bytes_q4: prev=  37 ms, new=  39 ms, diff=1.05 slower ✖, bytes=1417K→2254K (0.63x)
 tpch_sf1 bytes_q5: prev=  23 ms, new=  24 ms, diff=1.04 slower ✖, bytes=3983K→6120K (0.65x)
 tpch_sf1 bytes_q6: prev=   8 ms, new=   8 ms, diff=1.00 slower ✖, bytes=313K→370K (0.85x)
 tpch_sf1 bytes_q7: prev=  55 ms, new=  56 ms, diff=1.02 slower ✖, bytes=10898K→12925K (0.84x)
 tpch_sf1 bytes_q8: prev=   2 ms, new=   2 ms, diff=1.00 slower ✖
 tpch_sf1 bytes_q9: prev=   8 ms, new=   8 ms, diff=1.00 slower ✖
 tpch_sf1 bytes_q10: prev=   7 ms, new=   7 ms, diff=1.00 slower ✖, bytes=482K→639K (0.75x)
 tpch_sf1 bytes_q11: prev=  16 ms, new=  16 ms, diff=1.00 slower ✖, bytes=3545K→3537K (1.00x)
 tpch_sf1 bytes_q12: prev=  22 ms, new=  23 ms, diff=1.05 slower ✖, bytes=220K→219K (1.01x)
 tpch_sf1 bytes_q13: prev=  46 ms, new=  48 ms, diff=1.04 slower ✖, bytes=3164K→5198K (0.61x)
 tpch_sf1 bytes_q14: prev=   1 ms, new=   2 ms, diff=2.00 slower ❌
 tpch_sf1 bytes_q15: prev=   2 ms, new=   2 ms, diff=1.00 slower ✖
 tpch_sf1 bytes_q16: prev=   9 ms, new=   9 ms, diff=1.00 slower ✖
 tpch_sf1 bytes_q17: prev=   0 ms, new=   0 ms, diff=NaN slower ✖
 tpch_sf1 bytes_q18: prev= 165 ms, new= 165 ms, diff=1.00 slower ✖, bytes=36688K→53339K (0.69x)
 tpch_sf1 bytes_q19: prev=  13 ms, new=  14 ms, diff=1.08 slower ✖, bytes=751K→1091K (0.69x)
 tpch_sf1 bytes_q20: prev=  21 ms, new=  21 ms, diff=1.00 slower ✖, bytes=319K→378K (0.84x)
 tpch_sf1 bytes_q21: prev=  31 ms, new=  32 ms, diff=1.03 slower ✖, bytes=952K→1449K (0.66x)
 tpch_sf1 bytes_q22: prev=  10 ms, new=  10 ms, diff=1.00 slower ✖, bytes=274K→304K (0.90x)

[fmonjalet]

Hi Nathan, and thank you for taking this on!

The first approach looks great.
For the second, I think there is an issue in the design that explains the lack of improvement, you suggested:

RepartitionExec(Hash(group_keys, N×T))  ← existing, stays in place
  AggregateExec(PartialReduce)          ← new
    AggregateExec(Partial) 

And as you noted it does not help much, because it won't reduce data (Partial already did this job correctly 👍).
What you want to try out is:

AggregateExec(PartialReduce)          ← new
  RepartitionExec(Hash(group_keys, N×T))  ← existing, stays in place
    AggregateExec(Partial)

Now the PartialReduce runs on properly partitioned data and can yield data reduction.

Then you can compare with approach 1 and see how it goes! In theory it should have the same number of bytes exchanged. What is unclear to me is which one is less CPU intensive (or fastest), and your bench is perfect to measure this.
I would also suggest you try out on a larger scale factor (at least 10, ideally 100) so that you can ensure your bench are not too noisy. You can aim for an execution that lasts at least one second, ideally multiple.
Also this will be more visible if each workers have many partitions: you can try out with 16 partitions per worker (I think the default test set up may have 2 to 4), which is also closer to what you'd see in production deployments.

Keep us posted! Excited to see the early improvements.

[nchu05]

Thanks for the feedback! I've revised approach 2 to the new design and benchmarked using a scale factor of 50 and 16 partitions per worker. With these changes, I believe both approaches show similar results — leaning towards approach 2 because it adds less complexity?

Bytes Transferred

Query	GROUP BY	Distinct Keys	Baseline	Approach 2	Approach 1	A2 reduction	A1 reduction
bytes_q1	l_suppkey	500K	2,823,894 KB	20,352 KB	20,453 KB	138.8x	138.1x
bytes_q2	l_partkey	10M	3,127,082 KB	386,813 KB	387,759 KB	8.1x	8.1x
bytes_q3	l_orderkey	75M	1,817,756 KB	1,817,678 KB	1,820,543 KB	1.00x	1.00x
bytes_q4	l_suppkey, l_returnflag	1.5M	3,479,439 KB	69,106 KB	69,350 KB	50.4x	50.2x
bytes_q5	o_custkey	5M	828,830 KB	196,661 KB	197,113 KB	4.2x	4.2x
bytes_q6	o_orderdate	2.4K	1,112 KB	716 KB	751 KB	1.6x	1.5x
bytes_q7	o_custkey, o_orderpriority	23M	832,122 KB	599,255 KB	600,688 KB	1.4x	1.4x
bytes_q10	ps_suppkey	500K	323,732 KB	13,680 KB	13,760 KB	23.7x	23.5x
bytes_q13	l_suppkey, l_shipmode	3.5M	3,733,316 KB	161,125 KB	161,608 KB	23.2x	23.1x
bytes_q18	l_partkey, l_suppkey	40M	4,983,339 KB	1,970,153 KB	1,975,506 KB	2.5x	2.5x
bytes_q22	o_clerk	50K	15,657 KB	3,365 KB	3,396 KB	4.7x	4.6x

Query Execution Time

Query	Baseline	Approach 2	Approach 1	A2 speedup	A1 speedup
bytes_q1	5,749 ms	3,171 ms	3,400 ms	1.81x	1.69x
bytes_q4	7,076 ms	4,904 ms	4,890 ms	1.44x	1.45x
bytes_q5	2,064 ms	1,650 ms	1,659 ms	1.25x	1.24x
bytes_q10	678 ms	363 ms	359 ms	1.87x	1.89x
bytes_q13	6,451 ms	5,511 ms	5,442 ms	1.17x	1.19x
bytes_q18	13,072 ms	11,364 ms	11,907 ms	1.15x	1.10x

[fmonjalet]

Awesome, thanks! These numbers look good, I suspect the speedup will be more significant when running on different machines with true network.
Given these I think Approach 2 makes sense, and we can still revise this as needed in the future 👍

[gabotechs]

Approach 2 with @fmonjalet revision looks good! and those numbers look great. I'd expect to see real latency improvements if running in a "real" scenario, with actual network calls.