[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 4 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>
3 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