Skip to content

[SPARK-32750][SQL] Support whole-stage codegen for SortAggregate with grouping keys#57152

Closed
ulysses-you wants to merge 4 commits into
apache:masterfrom
ulysses-you:apache-agg-codegen
Closed

[SPARK-32750][SQL] Support whole-stage codegen for SortAggregate with grouping keys#57152
ulysses-you wants to merge 4 commits into
apache:masterfrom
ulysses-you:apache-agg-codegen

Conversation

@ulysses-you

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

This PR adds whole-stage code-gen support for SortAggregateExec when it has grouping
keys. Previously only the no-grouping-keys path was code-gen'd; with keys it fell back to
the interpreted SortBasedAggregationIterator.

The implementation:

  • AggregateCodegenSupport is refactored to share the aggregation-buffer creation
    (createAggBufVars) and buffer-update (generateAggBufferUpdateCode) logic between the
    no-keys path and the new sort-based with-keys path.
  • SortAggregateExec implements doProduceWithKeys / doConsumeWithKeys. Since the input
    is sorted by the grouping keys, a group's result is emitted as soon as the next group
    starts (detected by comparing the binary representation of the grouping key). The produce
    loop is resumable across processNext invocations via shouldStop(), so a completed
    group's output row is not overwritten by the reused output buffer.
  • The with-keys path only supports binary-stable grouping-key types (guarded in
    supportCodegen), because group boundaries are detected via UnsafeRow.equals.
  • A new internal config spark.sql.codegen.aggregate.sortAggregate.withKeys.enabled
    (default true) gates this path; it takes effect only when
    spark.sql.codegen.aggregate.sortAggregate.enabled is enabled.

Why are the changes needed?

Sort aggregate with grouping keys was the only remaining aggregate path without whole-stage
code-gen, forcing it through the slower interpreted iterator. SortAggregateBenchmark shows
a consistent 1.2~1.3x speedup for grouped aggregates on the code-gen path.

Does this PR introduce any user-facing change?

No. This is an internal code-gen optimization; results are unchanged. The new config is
internal().

How was this patch tested?

  • New tests in WholeStageCodegenSuite covering: multiple/numeric/string/decimal grouping
    keys, null keys and values, single-row groups, a single all-rows group, downstream limit
    (resumable shouldStop path), empty input, single partition, split aggregate functions,
    FILTER clauses, HAVING-style filters, and the config gate. Each test asserts a code-gen'd
    SortAggregateExec in the plan and checks the result matches the interpreted (code-gen
    disabled) result.
  • Added SortAggregateBenchmark with results committed for JDK 17/21/25.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code

@ulysses-you ulysses-you closed this Jul 9, 2026
@ulysses-you ulysses-you deleted the apache-agg-codegen branch July 9, 2026 03:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant