Skip to content

perf: Drop the Spark-planned sort below a Comet partial sort aggregate #7323

Open

Description

@parthchandra

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 = true to 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 (SortAggregateExec suppor

Steps to reproduce

No response

Expected behavior

No response

Additional context

No response

Willingness to contribute

None

Activity

  1. self-assigned this
    on Oct 29, 2026
  2. changed the title [-]feat: Drop the Spark-planned sort below a Comet partial sort aggregate[/-] on Oct 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions