perf: Drop the Spark-planned sort below a Comet partial sort aggregate #7323
Description
Describe the bug
What is the problem the feature request solves?
Spark plans an aggregate as SortAggregateExec when its buffer can't live in an UnsafeRow and it isn't a TypedImperativeAggregate. One example is min/max over a struct, which now runs natively (#7191). EnsureRequirements then puts a SortExec on the grouping keys below both the partial and the final aggregate. Comet converts both into native sorts.
Comet's native aggregate doesn't need sorted input. It is DataFusion's AggregateExec, which hashes groups when the input is unordered. So for the partial stage, the sort is pure overhead: every input row, with all its columns, is sorted before being reduced to one row per group.
With wide rows this cost dominates. In CometAggregateExpressionBenchmark ("Struct min/max benchmarks": 200k rows, 1,000 groups, 4 KiB string payloads), SELECT MAX(struct(k, wide)) FROM t GROUP BY grp takes:
| Best time | |
|---|---|
| Spark | 258 ms |
| Comet | 398 ms |
The global version without grouping, which has no sort, is at parity. Measured natively on the same data:
| Time | |
|---|---|
| DataFusion sort by the grouping key | ~200 ms |
| Struct max accumulator | ~20 ms |
With 64-byte payloads the same grouped query is 2.6x faster in Comet, so the extra cost grows with row width.
Describe the potential solution
When a CometSortAggregateExec runs in Partial mode, remove the SortExec/CometSort below it that exists only to satisfy SortAggregateExec.requiredChildOrdering, and lthe native aggregate hash its input.
Check the partial's parent before dropping the sort. It's only safe whel relies on its output ordering. Usually the parent is a shuffle, whichdoesn't preserve order, so the sort can go. But when the input is already partitioned by the grouping keys and not ordered by them, Spark plans no exchange between partial andfinal. One example is the output of a shuffled hash join on the grouping kew. There, the final SortAggregateExec relies on the partial's outputordering for its own required ordering, so EnsureRequirements adds no sort before the final. Dropping the sort under the partial would leave the final with unsorted input. Either keep the sort in that case, or move it up between partial and final,r group instead of every input row. When the child is already ordered by the grouping keys, for example a sort-merge join on those keys, Spark adds no sort at all, so there is nothing to remove.
Other points:
- Plan-time decision: the partial aggregate reports
orderedByGroupingKeys = trueto the native planner. That has to become false when the sort is removed, so the native side doesn't assume sorted input. - Memory behavior: a sorted partial emits finished groups after each batch. A hash partial holds every group until it emits early under memory pressure. Comet's partial aggregates already behave this way for
HashAggregateExec, so this should - PartialMerge stages from distinct-aggregate rewrites can likely be treated the same way when they feed a shuffle, with the same parent check.
The final stage is a separate question. There, the sort after the shuffle lets DataFusion use its ordered final stream, which keeps memory bounded. Dropping it would trade sort
time for a hash final that must hold or spill every group. It could be cons behind a config.
Additional context
-
Benchmark cases:
CometAggregateExpressionBenchmark,max_struct_wide/d) and*_global. -
The no-exchange case, confirmed on Spark 4.1 with broadcast joins and AQE turned off. The partial and final sort aggregates both run natively with no exchange between them,
and only the partial has a sort below it:SELECT /*+ SHUFFLE_HASH(b) */ a.k, max(struct(a.v, b.s)) FROM a JOIN b ON a.k = b.k GROUP BY a.k
CometSortAggregate [Final] +- CometSortAggregate [Partial] +- CometSort [k ASC NULLS FIRST] +- CometProject +- CometHashJoin [k], [k], Inner, BuildRight :- CometExchange hashpartitioning(k, 200) +- CometExchange hashpartitioning(k, 200) -
Related: feat: run min and max over struct inputs natively #7191 (native struct min/max), feat: support SortAggregateExec #4565 (
SortAggregateExecsuppor
Steps to reproduce
No response
Expected behavior
No response
Additional context
No response
Willingness to contribute
None
Activity
[-]feat: Drop the Spark-planned sort below a Comet partial sort aggregate[/-]