Reducing Spark SQL Shuffle with Partial Aggregation Before Joins

Overview

A common reporting query joins a fact table to a dimension table, groups by a dimension attribute, and adds up a measure:

SELECT drivers.region,
       SUM(trips.fare_amount) AS total_fare,
       COUNT(*)               AS trip_count
FROM trips
JOIN drivers ON trips.driver_id = drivers.driver_id
GROUP BY drivers.region;

trips has billions of rows and drivers has millions, but the query returns only one row per region. Spark’s plan joins each trip row to its driver, often shuffling it across the network, and then collapses the joined rows into those few regional totals.

Quanton adds a rule that aggregates the trips by driver_id before the join, carries the intermediate sums and counts through it, and finishes the totals afterwards. The idea is decades old: Yan and Larson’s Performing Group-By before Join (ICDE 1994) proves when a group-by can move past a join.

The problem: the join sits between the rows and the totals

Spark already splits an aggregate into a partial aggregate in each task and a final aggregate after the shuffle. In this query, both steps run after the join, because region comes from the drivers table:

Aggregate [region], [SUM(fare_amount), COUNT(*)]
  Join Inner, trips.driver_id = drivers.driver_id
    trips      ← billions of rows, all of them go through the join
    drivers

A driver with thousands of trips sends thousands of rows through the join, and the aggregate folds them straight back into one sum.

The idea: aggregate first, then join

A sum per region is a sum of per-driver sums, so Quanton computes those subtotals first. If one partition contains 1,000 trips for the same driver, the partial aggregate produces one row with that driver’s fare subtotal and a trip count of 1,000. The join matches that row to the driver’s region, and the final aggregate adds the subtotals and counts for each region:

Final aggregate [region], [combine fare subtotals, add trip counts]
  Join Inner, trips.driver_id = drivers.driver_id
    Partial aggregate [driver_id]
      outputs: driver_id, fare_sum, fare_is_empty, trip_count
      trips
    drivers

Spark stores these partial results in aggregation buffers, the intermediate values a function needs to finish later (a decimal SUM carries a running sum and an “is empty” flag, a COUNT a count). The join matches buffer rows instead of raw trips, and the final aggregate merges them.

The partial aggregate runs inside each partition and adds no shuffle, so a driver whose trips span several partitions yields several buffer rows, and the final aggregate merges them. The more rows share a join key in a partition, the bigger the win.

Keeping the answer exact

Duplicate matches. If driver 1 appears in both North and South, the join copies driver 1’s single buffer row into both regions, exactly as it would have copied the raw rows.

Rows the join throws away. The partial aggregate runs before the join, so it also evaluates trips that have no matching driver, rows the original plan drops before any aggregate sees them. If an expression can fail on such a row, the rewrite would turn a query that succeeds into one that fails. With ANSI mode on, SUM(10 / fare_amount) raises DIVIDE_BY_ZERO on an unmatched trip with a zero fare, and an integral SUM can raise ARITHMETIC_OVERFLOW.

Quanton checks both the aggregate function and its input expressions before moving it before the join. Functions such as COUNT, MIN, MAX, AVG and decimal SUM are eligible only when their inputs and the rest of the query pass those checks. A CASE WHEN expression must pass the same checks; it does not make a SUM automatically eligible. An integral SUM under ANSI does not qualify because it can overflow.

All or nothing. If one aggregate can’t be pushed, the rule leaves the whole query alone.

When the rule does not apply

The rule declines, and Spark plans the query as usual, when any of these hold:

• An aggregate uses DISTINCT or a FILTER clause, or reads columns from both sides of the join.
• The query uses an unsupported aggregate function, such as collect_list, collect_set, percentile, sketches or a user-defined aggregate function (UDAF).
• A function depends on the order in which rows arrive: FIRST, LAST and ANY_VALUE.
• An aggregate input can raise an error on a row the join drops, as described above.
• The join is not an inner join, has no plain column equality between its two sides, or carries a join hint.
• An expression in the aggregate, the join condition or a projection between them is nondeterministic, meaning it can return different values on repeated evaluations, like rand().
• The join key the partial aggregate would group on is a floating-point column.

When it does not help

A valid rewrite is not always a faster one. The partial aggregate costs CPU and memory, and it only pays off when it collapses many rows into few:

• Few repeated keys per partition. If most join keys appear once or twice in each partition, the partial aggregate emits almost as many rows as it reads.
• Wide buffers. Each aggregate adds one or more buffer columns. Many aggregates over a nearly unique key can produce buffer rows wider than the raw rows they replace.
• Filters after the join. A filter on a fact column evaluated after the join adds that column to the pushed grouping key, which splits the groups and shrinks the reduction.
• A small fact side. When the fact side is already small, or a broadcast join avoids its shuffle, there is little shuffle volume left to save.

The execution has to keep up

The rule alone does not speed up a query. The shuffle still carries partial results between executors, so writing and reading those results must be efficient. A buffer has several fields, and an aggregate expression in Spark returns one value, so the natural way to carry it is one struct column. Quanton’s native shuffle, though, writes struct columns row by row. Giving each buffer field its own column lets the shuffle process columnar batches instead of serializing each struct row separately. For the shuffle carrying partial results in q23a:

Metric	Struct column	Flat columns	Change
Time to split (write)¹	29.9 m	4.1 m	7.3×
Time to deserialize (read)¹	9.2 m	17.2 s	32×
Shuffle bytes written	101.6 GiB	69.8 GiB	−31%
Sum of task shuffle times¹	33.8 m	7.2 m	4.7×

¹ Summed across all 1,307 tasks of the exchange, not elapsed time.

Turning it on

The rule is off by default. Use these settings to enable it:

spark.quanton.sql.optimizer.pushAggregateThroughJoin.enabled=true
spark.quanton.sql.columnar.backend.velox.propagateIgnoreNullKeys=false

Queries that fail any of the checks above are planned the way Spark always plans them, so enabling the rule does not change their results.

Impact

When many rows share a join key within each partition, partial aggregation can reduce the data sent through the shuffle and join. Quanton computes partial results first, carries Spark’s aggregation buffers through the join so the answer stays exact, and ships them as separate columns for efficient shuffle processing.

1. TPC-DS at 10 TB. Our reported results show a 221-second reduction in total runtime, led by q23a and q23b. In that comparison, q23a ran 2.1× faster and q23b 2.3× faster.
2. Production ETL. For our users’ heavy ETL jobs that run large joins, shuffle volume dropped by 2–3× and job durations by 70%.