[SPARK-32750][SQL] Support whole-stage codegen for SortAggregate with grouping keys#57153
[SPARK-32750][SQL] Support whole-stage codegen for SortAggregate with grouping keys#57153ulysses-you wants to merge 5 commits into
Conversation
… grouping keys Add 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. A new internal config spark.sql.codegen.aggregate.sortAggregate.withKeys.enabled (default true) gates this path and takes effect only when spark.sql.codegen.aggregate.sortAggregate.enabled is enabled. Co-Authored-By: Claude <noreply@anthropic.com>
|
.cc @viirya @dongjoon-hyun @cloud-fan if you have time to take a look, thank you! |
|
Long time no see. |
| val doAgg = ctx.freshName("doAggregateWithKeys") | ||
| val doAggFuncName = ctx.addNewFunction(doAgg, | ||
| s""" | ||
| |private void $doAgg() throws java.io.IOException { |
There was a problem hiding this comment.
The $doAgg wrapper here doesn't take an int partitionIndex parameter, unlike the no-keys path, HashAggregateExec, and SortExec. The child's produce emits bare partitionIndex references (e.g. addToSorter(partitionIndex) for a Sort child). When the outer generated class exceeds GENERATED_CLASS_SIZE_THRESHOLD (1MB) and addNewFunction spills $doAgg into a private class NestedClass, that bare reference resolves to the protected BufferedRowIterator.partitionIndex field, which the inner (non-subclass) class can't access, so Janino compilation throws IllegalAccessError. That's an Error, not NonFatal, so the codegen fallback in WholeStageCodegenExec won't catch it and the task fails. Suggest matching the other three wrappers: private void $doAgg(int partitionIndex), called as $doAggFuncName(partitionIndex).
There was a problem hiding this comment.
Good catch, fixed in abe40dd. The doAggregateWithKeys wrapper now takes int partitionIndex and is called as $doAggFuncName(partitionIndex), matching the no-keys path, HashAggregateExec, and SortExec. This prevents the bare partitionIndex in the child produce (e.g. sort_addToSorter_0(partitionIndex)) from resolving to the protected BufferedRowIterator.partitionIndex field once addNewFunction spills the helper into a nested class past the 1MB threshold.
| } | ||
|
|
||
| protected override def doProduceWithKeys(ctx: CodegenContext): String = { | ||
| throw new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3170") |
There was a problem hiding this comment.
This PR removes the only two throw sites for _LEGACY_ERROR_TEMP_3170 in SortAggregateExec, leaving the entry at error-conditions.json:11601 unreferenced. `
There was a problem hiding this comment.
Removed the now-unreferenced _LEGACY_ERROR_TEMP_3170 entry from error-conditions.json in abe40dd. Confirmed no other references remain in sql/ or common/, and SparkThrowableSuite passes.
| } | ||
|
|
||
| test("SPARK-32750: SortAggregate code-gen with grouping keys") { | ||
| val data = spark.range(200).selectExpr( |
There was a problem hiding this comment.
supportCodegenWithKeys gates only on isBinaryStable, and Double/Float/Decimal are all binary-stable, so they take this new path. Group boundaries use raw UnsafeRow.equals, so float-key correctness rides on the planner's NormalizeFloatingNumbers folding -0.0/NaN first, but the new tests only cover long/string keys with no float grouping-key correctness assertion (decimal appears only in the benchmark). Suggest adding a float grouping-key case (with
There was a problem hiding this comment.
Added coverage in abe40dd. The new "float/double grouping keys" test includes 0.0/-0.0 (must group together), NaN plus a non-canonical NaN bit pattern (all must group together), and nulls, checking single-key (float, double) and multi-key cases. This verifies the NormalizeFloatingNumbers canonicalization holds under the UnsafeRow.equals group-boundary check. Also added a separate decimal grouping-key test (with nulls).
| |$evaluateKeyVars | ||
| |${consume(ctx, resultVars)} | ||
| """.stripMargin | ||
| } else { |
There was a problem hiding this comment.
The third branch of generateResultFunctionForKeys (no aggregate functions, grouping-only) is codegen'd, but every new test carries at least one aggregate function, so this branch and its empty bufVars/reInitBufferCode path are uncovered. Suggest adding a groupBy("k").agg() or distinct case.
There was a problem hiding this comment.
Added a "no aggregate functions" test in abe40dd using DISTINCT (single- and multi-key), which lowers to a grouping aggregate with no aggregate functions. This exercises the third branch of generateResultFunctionForKeys and its empty bufVars/reInitBufferCode path.
cloud-fan
left a comment
cloud-fan
left a comment
There was a problem hiding this comment.
1 blocking, 4 non-blocking.
Solid PR. The incremental-emission design and its equivalence with the interpreted SortBasedAggregationIterator check out. @LuciferYang's four inline comments are all valid — I independently confirmed each and agree; I'm not re-posting on those same lines to avoid duplicating the threads. One new finding below.
Correctness (1)
- SortAggregateExec.scala:247: (blocking) agrees with @LuciferYang's thread — the
$doAggwrapper omits theint partitionIndexparameter that the no-keys path (AggregateCodegenSupport.scala:171) andHashAggregateExecpass for nested-class safety. Once the outer class spills past the 1MB threshold andaddNewFunctionmoves$doAgginto a nested class, the barepartitionIndexreference (e.g.sort_addToSorter_0(partitionIndex), present in this PR's own generated-code sample) resolves to the protectedBufferedRowIterator.partitionIndexfield, which the non-subclass nested class can't access —IllegalAccessErrorat compile. It's anError, notNonFatal, so the codegen fallback won't catch it and the task fails. Fix:private void $doAgg(int partitionIndex)called as$doAggFuncName(partitionIndex).
Design / architecture (1)
- WholeStageCodegenSuite.scala:106: the PR description overstates test coverage — see inline.
Suggestions (3)
- WholeStageCodegenSuite.scala:101 (existing @LuciferYang thread): add a float grouping-key correctness test —
isBinaryStableadmits float/double/decimal keys, whose-0.0/NaNhandling rides on upstreamNormalizeFloatingNumbers. - SortAggregateExec.scala:202 (existing @LuciferYang thread): the grouping-only (no aggregate function) branch of
generateResultFunctionForKeysis uncovered — every new test has ≥1 aggregate. - SortAggregateExec.scala:105 (existing @LuciferYang thread): remove the now-dead
_LEGACY_ERROR_TEMP_3170entry aterror-conditions.json:11612(confirmed: no other references remain insql/orcommon/).
Verification
Traced the incremental-emission safety since the operator drops the blocking-operator default (needStopCheck = true, canCheckLimitNotReached = false): BufferedRowIterator.shouldStop() returns true whenever currentRows is non-empty and hasNext only calls processNext once it's drained, so at most one output row is buffered before it's consumed — the single reused output-writer buffer is never aliased across emitted rows. The noMoreInputTerm + shouldStop() logic correctly resumes the scan across processNext calls and flushes the last group at real end-of-input (including the last-row-triggers-output case). Group-boundary equality via UnsafeRow.equals under the isBinaryStable gate matches the interpreted iterator's groupKeyEqualityCheck exactly. No correctness issue found in the core mechanism.
| "id % 7 as k1", | ||
| "id % 3 as k2", | ||
| "case when id % 5 = 0 then null else id end as v", | ||
| "case when id % 4 = 0 then null else cast(id % 11 as string) end as s") |
There was a problem hiding this comment.
The PR description claims the new tests cover "decimal grouping keys", "split aggregate functions", and "the config gate", but the diff doesn't back this:
- No new test uses a decimal (or float) grouping key — the test columns here are long
k1/k2and strings; decimal only appears in the benchmark. - No new test sets
CODEGEN_METHOD_SPLIT_THRESHOLD/CODEGEN_SPLIT_AGGREGATE_FUNC, so with the default threshold (1024) these small queries never reach the split-aggregate path. - No new test references
ENABLE_SORT_AGGREGATE_CODEGEN_WITH_KEYS, so the new config gate is never toggled.
Either add the coverage or correct the description. (The decimal half overlaps @LuciferYang's thread just below on float/-0.0/NaN correctness, which is the sharper concern.)
There was a problem hiding this comment.
Thanks, addressed in abe40dd by adding the missing coverage rather than trimming the description:
- decimal grouping-key test (with nulls), plus a float/double test covering
-0.0/NaN(the sharper concern from the thread below). - split-aggregate test toggling
CODEGEN_SPLIT_AGGREGATE_FUNC=true+CODEGEN_METHOD_SPLIT_THRESHOLD=1. - config-gate test toggling
ENABLE_SORT_AGGREGATE_CODEGEN_WITH_KEYS, asserting no code-gendSortAggregateExec` when disabled while the result stays correct.
Also updated the "How was this patch tested?" section so it matches what is actually covered.
- doProduceWithKeys: give the doAggregateWithKeys wrapper an `int partitionIndex` parameter (matching the no-keys path, HashAggregateExec and SortExec) so bare partitionIndex references in the child's produce resolve to the local instead of the protected BufferedRowIterator field when the helper is spilled to a nested class. - Remove the now-unreferenced _LEGACY_ERROR_TEMP_3170 error condition. - WholeStageCodegenSuite: add coverage for float/double keys (-0.0/NaN), decimal keys, grouping-only (no aggregate function) aggregates, split aggregate functions, and the with-keys config gate. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
cloud-fan
left a comment
cloud-fan
left a comment
There was a problem hiding this comment.
5 addressed, 0 remaining, 2 new. (2 = 0 newly introduced, 2 late catches - my own misses from the earlier round.)
0 blocking, 2 non-blocking, 0 nits.
The blocking nested-class-safety fix ($doAgg(int partitionIndex)) and all four non-blocking follow-ups from the prior round are correctly addressed - re-verified. Two minor late catches below, neither gates merge.
Correctness (1)
- WholeStageCodegenSuite.scala:86: inverted assertion failure message in the
checkSortAggregateCodegenhelper - see inline
Suggestions (1)
- SortAggregateExec.scala:119:
canCheckLimitNotReachedoverride is unconditional - see inline
Verification
Re-confirmed the incremental-emission safety this round: needStopCheck = true (with keys) propagates the stop check up the child produce chain (e.g. SortExec at if (shouldStop()) return;), and BufferedRowIterator.shouldStop() returns true iff currentRows is non-empty while hasNext drains before re-calling processNext - so at most one output row is buffered before consumption and the single reused UnsafeRowWriter is never aliased across emitted rows. The noMoreInput + shouldStop() driver correctly distinguishes real end-of-input from a mid-scan pause and flushes the final group (including the last-row-triggers-output case). Group-boundary equality via UnsafeRow.equals under the isBinaryStable gate matches the interpreted SortBasedAggregationIterator.groupKeyEqualityCheck exactly. No correctness issue in the core mechanism.
| assert(!df.queryExecution.executedPlan.exists(p => | ||
| p.isInstanceOf[WholeStageCodegenExec] && | ||
| p.asInstanceOf[WholeStageCodegenExec].child.isInstanceOf[SortAggregateExec]), | ||
| s"Expected a code-gen'd SortAggregateExec in:\n${df.queryExecution.executedPlan}") |
There was a problem hiding this comment.
This is the interpreted/expected branch (codegen disabled), so the message should say the plan must not be code-gen'd - it currently reads the inverse of what the assertion checks. The config-gate test at :248 has the correct wording.
| s"Expected a code-gen'd SortAggregateExec in:\n${df.queryExecution.executedPlan}") | |
| s"Expected no code-gen'd SortAggregateExec in:\n${df.queryExecution.executedPlan}") |
| // default (no stop check) still applies there. | ||
| override def needStopCheck: Boolean = groupingExpressions.nonEmpty | ||
|
|
||
| override protected def canCheckLimitNotReached: Boolean = false |
There was a problem hiding this comment.
This override is unconditional, so it also flips the no-keys path (which otherwise inherits true from BlockingOperatorWithCodegen). It's inert - neither produce path calls limitNotReachedCond - but for consistency with needStopCheck just above (gated on groupingExpressions.nonEmpty), consider gating this the same way. Optional.
| override protected def canCheckLimitNotReached: Boolean = false | |
| override protected def canCheckLimitNotReached: Boolean = groupingExpressions.nonEmpty |
…NotReached - WholeStageCodegenSuite: the interpreted/expected branch asserts the plan is NOT code-gen'd, so correct the failure message to say so. - SortAggregateExec: gate canCheckLimitNotReached on groupingExpressions.nonEmpty for consistency with needStopCheck (avoids flipping the no-keys path). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
cloud-fan
left a comment
cloud-fan
left a comment
There was a problem hiding this comment.
2 addressed, 0 remaining, 0 new.
0 blocking, 0 non-blocking, 0 nits.
Clean re-review. Both round-2 follow-ups are correctly addressed — the inverted assertion message (WholeStageCodegenSuite.scala:86) and the canCheckLimitNotReached gating (SortAggregateExec.scala:119). No new issues.
Verification
Re-confirmed the two fixes: the corrected assertion message now matches its assert(!codegen'd) in the code-gen-disabled branch; and canCheckLimitNotReached = groupingExpressions.nonEmpty is inert for this operator (it never calls limitNotReachedCond — only HashAggregateExec does) and now mirrors needStopCheck. Core mechanism unchanged and re-verified: codegen group-boundary equality (UnsafeRow.equals under the isBinaryStable gate) matches SortBasedAggregationIterator.initialize exactly, and the incremental-emission path buffers at most one row before consumption, so the reused output writer is never aliased. LGTM.
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
A few post-approval suggestions, none blocking.
| // the single reused output row buffer. The with-keys path is thus not fully blocking. The | ||
| // without-keys path produces its single result only after the scan completes, so the blocking | ||
| // default (no stop check) still applies there. | ||
| override def needStopCheck: Boolean = groupingExpressions.nonEmpty |
There was a problem hiding this comment.
The comment block here explains why needStopCheck must flip for the with-keys path, but the operator still inherits needCopyResult = false from BlockingOperatorWithCodegen, whose stated justification ("blocking operators keep the data in some buffer") no longer applies once this path streams results out mid-scan. It's still safe, but for a different reason: the stop-check discipline guarantees at most one output row is buffered before it's consumed, and any in-stage parent that multiplies rows (e.g. joins) declares needCopyResult = true itself. Since that argument is non-obvious and this comment is where a future reader would look, suggest adding a sentence covering needCopyResult as well.
There was a problem hiding this comment.
Good point. Added a paragraph to the comment block spelling this out: we keep the inherited needCopyResult = false, but its original justification no longer applies once the with-keys path streams results mid-scan. It stays safe for a different reason -- the stop-check discipline guarantees at most one output row is buffered before it's consumed, so the reused output row is never aliased, and any in-stage parent that multiplies rows (e.g. a join) declares needCopyResult = true itself. Done in 874cc2a.
| } | ||
| } | ||
|
|
||
| test("SPARK-32750: SortAggregate code-gen with float/double grouping keys") { |
There was a problem hiding this comment.
All the new tests exercise the supportCodegenWithKeys gate in the true direction; the isBinaryStable = false fallback is never hit (the config-gate test disables the path via conf, not via the type check). Suggest one case with a non-binary collated string grouping key (e.g. UTF8_LCASE), asserting no code-gen'd SortAggregateExec appears in the plan while the result stays correct - that pins the type-based fallback to the interpreted path, which is the half of the gate a future type-support change is most likely to break.
There was a problem hiding this comment.
Added a "non-binary collated key" test that groups by a UTF8_LCASE string key (not binary-stable), asserting no code-gen'd SortAggregateExec appears in the plan (the interpreted fallback is used) while the collation-aware result stays correct ('a'/'A' and 'b'/'B' each collapse to one group, 'c' stays alone). This pins the isBinaryStable = false half of the gate. Done in 874cc2a.
| val N = 20 << 22 | ||
| val benchmark = new Benchmark("sort agg w linear keys", N, output = output) | ||
| // The child of a sort aggregate must be sorted by the grouping keys. Sorting the whole input | ||
| // dominates, so pre-sort once into a temp view and aggregate over the already-ordered rows. |
There was a problem hiding this comment.
Nit: "pre-sort once into a temp view" reads as if the sort cost is paid once, but a temp view is lazy - the sort re-executes on every iteration (and without the explicit sortWithinPartitions the planner would insert an equivalent SortExec below the aggregate anyway). What the setup actually achieves is that both cases pay the same sort cost, isolating the aggregate's own contribution - which also means the reported 1.2~1.3x understates the aggregate-only speedup. Suggest rewording the comment along those lines.
There was a problem hiding this comment.
Reworded along those lines: the temp view is lazy so the sort re-executes each iteration (and without the explicit sortWithinPartitions the planner would insert an equivalent SortExec below the aggregate anyway); pinning the sort just makes both cases pay the same sort cost, isolating the aggregate's own contribution -- which means the reported speedup understates the aggregate-only gain. Done in 874cc2a.
- Document why needCopyResult=false stays safe for the streaming with-keys path in the needStopCheck comment block. - Add a non-binary collated (UTF8_LCASE) grouping-key test asserting the interpreted fallback (no code-gen'd SortAggregateExec) with a correct collation-aware result. - Reword the pre-sort comment in SortAggregateBenchmark to clarify the temp view is lazy and the setup only equalizes sort cost across cases. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
LGTM. All three of my post-approval suggestions are addressed in 874cc2a -- the needCopyResult rationale is now documented in the comment block, the non-binary-collated-key test pins the isBinaryStable = false fallback to the interpreted path, and the benchmark comment is corrected. I ran the SPARK-32750 test group locally (all 9 pass, including the new collation test). The incremental-emission mechanism was already verified in the earlier rounds. Thanks @ulysses-you!
|
thank you @cloud-fan @LuciferYang @viirya for review! |
… grouping keys
### 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 grouping keys,
float/double keys (including `-0.0` and `NaN`), decimal keys, null keys and values,
single-row groups, a single all-rows group, grouping-only aggregates (no aggregate
function, e.g. `DISTINCT`), 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.
The generated code of a simple query: `select id , count(*) from t1 group by id`:
```
/* 001 */ public Object generate(Object[] references) {
/* 002 */ return new GeneratedIteratorForCodegenStage1(references);
/* 003 */ }
/* 004 */
/* 005 */ // codegenStageId=1
/* 006 */ final class GeneratedIteratorForCodegenStage1 extends org.apache.spark.sql.execution.BufferedRowIterator {
/* 007 */ private Object[] references;
/* 008 */ private scala.collection.Iterator[] inputs;
/* 009 */ private boolean sortAgg_bufIsNull_0;
/* 010 */ private long sortAgg_bufValue_0;
/* 011 */ private UnsafeRow sortAgg_currentGroupingKey_0;
/* 012 */ private boolean sortAgg_initGroup_0;
/* 013 */ private boolean sortAgg_noMoreInputTerm_0;
/* 014 */ private boolean sort_needToSort_0;
/* 015 */ private org.apache.spark.sql.execution.UnsafeExternalRowSorter sort_sorter_0;
/* 016 */ private org.apache.spark.executor.TaskMetrics sort_metrics_0;
/* 017 */ private scala.collection.Iterator<UnsafeRow> sort_sortedIter_0;
/* 018 */ private int columnartorow_batchIdx_0;
/* 019 */ private org.apache.spark.sql.vectorized.ColumnarBatch[] columnartorow_mutableStateArray_1 = new org.apache.spark.sql.vectorized.ColumnarBatch[1];
/* 020 */ private org.apache.spark.sql.execution.vectorized.OnHeapColumnVector[] columnartorow_mutableStateArray_2 = new org.apache.spark.sql.execution.vectorized.OnHeapColumnVector[1];
/* 021 */ private scala.collection.Iterator[] columnartorow_mutableStateArray_0 = new scala.collection.Iterator[1];
/* 022 */ private org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[] sortAgg_mutableStateArray_0 = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter[3];
/* 023 */
/* 024 */ public GeneratedIteratorForCodegenStage1(Object[] references) {
/* 025 */ this.references = references;
/* 026 */ }
/* 027 */
/* 028 */ public void init(int index, scala.collection.Iterator[] inputs) {
/* 029 */ partitionIndex = index;
/* 030 */ this.inputs = inputs;
/* 031 */
/* 032 */ sortAgg_mutableStateArray_0[0] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(2, 0);
/* 033 */ sort_needToSort_0 = true;
/* 034 */ sort_sorter_0 = ((org.apache.spark.sql.execution.SortExec) references[1] /* plan */).createSorter();
/* 035 */ sort_metrics_0 = org.apache.spark.TaskContext.get().taskMetrics();
/* 036 */ columnartorow_mutableStateArray_0[0] = inputs[0];
/* 037 */ sortAgg_mutableStateArray_0[1] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
/* 038 */ sortAgg_mutableStateArray_0[2] = new org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter(1, 0);
/* 039 */
/* 040 */ }
/* 041 */
/* 042 */ private void sortAgg_doConsume_0(InternalRow sort_outputRow_0, long sortAgg_expr_0_0, boolean sortAgg_exprIsNull_0_0) throws java.io.IOException {
/* 043 */ sortAgg_mutableStateArray_0[2].reset();
/* 044 */
/* 045 */ sortAgg_mutableStateArray_0[2].zeroOutNullBytes();
/* 046 */
/* 047 */ sortAgg_mutableStateArray_0[2].writeNullable(0, sortAgg_expr_0_0, sortAgg_exprIsNull_0_0);
/* 048 */ if (!sortAgg_initGroup_0) {
/* 049 */ sortAgg_initGroup_0 = true;
/* 050 */ sortAgg_currentGroupingKey_0 = (sortAgg_mutableStateArray_0[2].getRow()).copy();
/* 051 */ sortAgg_bufIsNull_0 = false;
/* 052 */ sortAgg_bufValue_0 = 0L;
/* 053 */ } else if (!sortAgg_currentGroupingKey_0.equals((sortAgg_mutableStateArray_0[2].getRow()))) {
/* 054 */ sortAgg_doAggregateWithKeysOutput_0();
/* 055 */ sortAgg_currentGroupingKey_0 = (sortAgg_mutableStateArray_0[2].getRow()).copy();
/* 056 */ sortAgg_bufIsNull_0 = false;
/* 057 */ sortAgg_bufValue_0 = 0L;
/* 058 */ }
/* 059 */
/* 060 */ // do aggregate
/* 061 */ // common sub-expressions
/* 062 */
/* 063 */ // evaluate aggregate functions and update aggregation buffers
/* 064 */
/* 065 */ long sortAgg_value_7 = -1L;
/* 066 */
/* 067 */ sortAgg_value_7 = org.apache.spark.sql.catalyst.util.MathUtils.addExact(sortAgg_bufValue_0, 1L, ((org.apache.spark.sql.catalyst.trees.SQLQueryContext) references[7] /* errCtx */));
/* 068 */
/* 069 */ sortAgg_bufIsNull_0 = false;
/* 070 */ sortAgg_bufValue_0 = sortAgg_value_7;
/* 071 */
/* 072 */ }
/* 073 */
/* 074 */ private void wholestagecodegen_doAggregateWithKeys_0(int partitionIndex) throws java.io.IOException {
/* 075 */ if (sort_needToSort_0) {
/* 076 */ long sort_spillSizeBefore_0 = sort_metrics_0.memoryBytesSpilled();
/* 077 */ sort_addToSorter_0(partitionIndex);
/* 078 */ sort_sortedIter_0 = sort_sorter_0.sort();
/* 079 */ ((org.apache.spark.sql.execution.metric.SQLMetric) references[6] /* sortTime */).add(sort_sorter_0.getSortTimeNanos() / 1000000);
/* 080 */ ((org.apache.spark.sql.execution.metric.SQLMetric) references[4] /* peakMemory */).add(sort_sorter_0.getPeakMemoryUsage());
/* 081 */ ((org.apache.spark.sql.execution.metric.SQLMetric) references[5] /* spillSize */).add(sort_metrics_0.memoryBytesSpilled() - sort_spillSizeBefore_0);
/* 082 */ sort_metrics_0.incPeakExecutionMemory(sort_sorter_0.getPeakMemoryUsage());
/* 083 */ sort_needToSort_0 = false;
/* 084 */ }
/* 085 */
/* 086 */ while ( sort_sortedIter_0.hasNext()) {
/* 087 */ UnsafeRow sort_outputRow_0 = (UnsafeRow)sort_sortedIter_0.next();
/* 088 */
/* 089 */ boolean sort_isNull_0 = sort_outputRow_0.isNullAt(0);
/* 090 */ long sort_value_0 = sort_isNull_0 ?
/* 091 */ -1L : (sort_outputRow_0.getLong(0));
/* 092 */
/* 093 */ sortAgg_doConsume_0(sort_outputRow_0, sort_value_0, sort_isNull_0);
/* 094 */
/* 095 */ if (shouldStop()) return;
/* 096 */ }
/* 097 */
/* 098 */ }
/* 099 */
/* 100 */ private void sort_addToSorter_0(int partitionIndex) throws java.io.IOException {
/* 101 */ if (columnartorow_mutableStateArray_1[0] == null) {
/* 102 */ columnartorow_nextBatch_0();
/* 103 */ }
/* 104 */ while ( columnartorow_mutableStateArray_1[0] != null) {
/* 105 */ int columnartorow_numRows_0 = columnartorow_mutableStateArray_1[0].numRows();
/* 106 */ int columnartorow_localEnd_0 = columnartorow_numRows_0 - columnartorow_batchIdx_0;
/* 107 */ for (int columnartorow_localIdx_0 = 0; columnartorow_localIdx_0 < columnartorow_localEnd_0; columnartorow_localIdx_0++) {
/* 108 */ int columnartorow_rowIdx_0 = columnartorow_batchIdx_0 + columnartorow_localIdx_0;
/* 109 */ boolean columnartorow_isNull_0 = columnartorow_mutableStateArray_2[0].isNullAt(columnartorow_rowIdx_0);
/* 110 */ long columnartorow_value_0 = columnartorow_isNull_0 ? -1L : (columnartorow_mutableStateArray_2[0].getLong(columnartorow_rowIdx_0));
/* 111 */ sortAgg_mutableStateArray_0[1].reset();
/* 112 */
/* 113 */ sortAgg_mutableStateArray_0[1].zeroOutNullBytes();
/* 114 */
/* 115 */ sortAgg_mutableStateArray_0[1].writeNullable(0, columnartorow_value_0, columnartorow_isNull_0);
/* 116 */ sort_sorter_0.insertRow((UnsafeRow)(sortAgg_mutableStateArray_0[1].getRow()));
/* 117 */ // shouldStop check is eliminated
/* 118 */ }
/* 119 */ columnartorow_batchIdx_0 = columnartorow_numRows_0;
/* 120 */ columnartorow_nextBatch_0();
/* 121 */ }
/* 122 */ // clean up resources
/* 123 */ if (columnartorow_mutableStateArray_1[0] != null) {
/* 124 */ columnartorow_mutableStateArray_1[0].close();
/* 125 */ }
/* 126 */
/* 127 */ }
/* 128 */
/* 129 */ protected void processNext() throws java.io.IOException {
/* 130 */ if (!sortAgg_noMoreInputTerm_0) {
/* 131 */ wholestagecodegen_doAggregateWithKeys_0(partitionIndex);
/* 132 */ if (!shouldStop()) {
/* 133 */ sortAgg_noMoreInputTerm_0 = true;
/* 134 */ if (sortAgg_initGroup_0) {
/* 135 */ sortAgg_doAggregateWithKeysOutput_0();
/* 136 */ }
/* 137 */ }
/* 138 */ }
/* 139 */ }
/* 140 */
/* 141 */ private void columnartorow_nextBatch_0() throws java.io.IOException {
/* 142 */ columnartorow_mutableStateArray_1[0] = org.apache.spark.sql.execution.ColumnarToRowExec.advanceBatch(
/* 143 */ columnartorow_mutableStateArray_0[0], columnartorow_mutableStateArray_1[0], ((org.apache.spark.sql.execution.metric.SQLMetric) references[3] /* numInputBatches */), ((org.apache.spark.sql.execution.metric.SQLMetric) references[2] /* numOutputRows */));
/* 144 */ if (columnartorow_mutableStateArray_1[0] != null) {
/* 145 */ columnartorow_batchIdx_0 = 0;
/* 146 */ columnartorow_mutableStateArray_2[0] = (org.apache.spark.sql.execution.vectorized.OnHeapColumnVector) columnartorow_mutableStateArray_1[0].column(0);
/* 147 */
/* 148 */ }
/* 149 */ }
/* 150 */
/* 151 */ private void sortAgg_doAggregateWithKeysOutput_0() throws java.io.IOException {
/* 152 */ ((org.apache.spark.sql.execution.metric.SQLMetric) references[0] /* numOutputRows */).add(1);
/* 153 */
/* 154 */ boolean sortAgg_isNull_1 = sortAgg_currentGroupingKey_0.isNullAt(0);
/* 155 */ long sortAgg_value_1 = sortAgg_isNull_1 ?
/* 156 */ -1L : (sortAgg_currentGroupingKey_0.getLong(0));
/* 157 */
/* 158 */ sortAgg_mutableStateArray_0[0].reset();
/* 159 */
/* 160 */ sortAgg_mutableStateArray_0[0].zeroOutNullBytes();
/* 161 */
/* 162 */ sortAgg_mutableStateArray_0[0].writeNullable(0, sortAgg_value_1, sortAgg_isNull_1);
/* 163 */
/* 164 */ sortAgg_mutableStateArray_0[0].write(1, sortAgg_bufValue_0);
/* 165 */ append((sortAgg_mutableStateArray_0[0].getRow()));
/* 166 */
/* 167 */ }
/* 168 */
/* 169 */ }
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #57153 from ulysses-you/SPARK-32750-sort-agg-codegen.
Authored-by: Xiduo You <ulyssesyou18@gmail.com>
Signed-off-by: Xiduo You <ulyssesyou@apache.org>
(cherry picked from commit 1db5e68)
Signed-off-by: Xiduo You <ulyssesyou@apache.org>
4 participants
What changes were proposed in this pull request?
This PR adds whole-stage code-gen support for
SortAggregateExecwhen it has groupingkeys. Previously only the no-grouping-keys path was code-gen'd; with keys it fell back to
the interpreted
SortBasedAggregationIterator.The implementation:
AggregateCodegenSupportis refactored to share the aggregation-buffer creation(
createAggBufVars) and buffer-update (generateAggBufferUpdateCode) logic between theno-keys path and the new sort-based with-keys path.
SortAggregateExecimplementsdoProduceWithKeys/doConsumeWithKeys. Since the inputis 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
processNextinvocations viashouldStop(), so a completedgroup's output row is not overwritten by the reused output buffer.
supportCodegen), because group boundaries are detected viaUnsafeRow.equals.spark.sql.codegen.aggregate.sortAggregate.withKeys.enabled(default true) gates this path; it takes effect only when
spark.sql.codegen.aggregate.sortAggregate.enabledis 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.
SortAggregateBenchmarkshowsa 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?
WholeStageCodegenSuitecovering: multiple/numeric/string grouping keys,float/double keys (including
-0.0andNaN), decimal keys, null keys and values,single-row groups, a single all-rows group, grouping-only aggregates (no aggregate
function, e.g.
DISTINCT), downstream limit (resumableshouldStoppath), empty input,single partition, split aggregate functions, FILTER clauses, HAVING-style filters, and the
config gate. Each test asserts a code-gen'd
SortAggregateExecin the plan and checks theresult matches the interpreted (code-gen disabled) result.
SortAggregateBenchmarkwith results committed for JDK 17/21/25.The generated code of a simple query:
select id , count(*) from t1 group by id:Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code