ulysses-you opened a new pull request, #57153:
URL: https://github.com/apache/spark/pull/57153
### 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
