From 090e25e6362135558074a981daf9eb659b7b9f76 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Thu, 9 Jul 2026 11:34:39 +0800 Subject: [PATCH 1/2] [SPARK-32750][SQL] Support whole-stage codegen for SortAggregate with 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 --- .../apache/spark/sql/internal/SQLConf.scala | 9 + .../SortAggregateBenchmark-jdk21-results.txt | 72 ++++++ .../SortAggregateBenchmark-jdk25-results.txt | 72 ++++++ .../SortAggregateBenchmark-results.txt | 72 ++++++ .../aggregate/AggregateCodegenSupport.scala | 42 +++- .../aggregate/SortAggregateExec.scala | 208 +++++++++++++++++- .../execution/WholeStageCodegenSuite.scala | 100 +++++++++ .../benchmark/SortAggregateBenchmark.scala | 148 +++++++++++++ 8 files changed, 706 insertions(+), 17 deletions(-) create mode 100644 sql/core/benchmarks/SortAggregateBenchmark-jdk21-results.txt create mode 100644 sql/core/benchmarks/SortAggregateBenchmark-jdk25-results.txt create mode 100644 sql/core/benchmarks/SortAggregateBenchmark-results.txt create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/SortAggregateBenchmark.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 6fe7a957a1216..c089c1690920b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -3699,6 +3699,15 @@ object SQLConf { .booleanConf .createWithDefault(true) + val ENABLE_SORT_AGGREGATE_CODEGEN_WITH_KEYS = + buildConf("spark.sql.codegen.aggregate.sortAggregate.withKeys.enabled") + .internal() + .doc("When true, enable code-gen for sort aggregate with grouping keys. Takes effect only " + + s"when ${ENABLE_SORT_AGGREGATE_CODEGEN.key} is enabled.") + .version("4.3.0") + .booleanConf + .createWithDefault(true) + val ENABLE_FULL_OUTER_SHUFFLED_HASH_JOIN_CODEGEN = buildConf("spark.sql.codegen.join.fullOuterShuffledHashJoin.enabled") .internal() diff --git a/sql/core/benchmarks/SortAggregateBenchmark-jdk21-results.txt b/sql/core/benchmarks/SortAggregateBenchmark-jdk21-results.txt new file mode 100644 index 0000000000000..221ff18c190f9 --- /dev/null +++ b/sql/core/benchmarks/SortAggregateBenchmark-jdk21-results.txt @@ -0,0 +1,72 @@ +================================================================================================ +sort aggregate without grouping +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 9V74 80-Core Processor +sort agg w/o group: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 16835 16940 149 31.1 32.1 1.0X +codegen = T 444 456 10 1182.0 0.8 38.0X + + +================================================================================================ +sort aggregate with linear keys +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 9V74 80-Core Processor +sort agg w linear keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 21718 22249 751 3.9 258.9 1.0X +codegen = T 17500 17624 114 4.8 208.6 1.2X + + +================================================================================================ +sort aggregate with randomized keys +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 9V74 80-Core Processor +sort agg w randomized keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 32644 32806 229 2.6 389.1 1.0X +codegen = T 29288 29539 182 2.9 349.1 1.1X + + +================================================================================================ +sort aggregate with string key +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 9V74 80-Core Processor +sort agg w string key: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 6726 6767 57 3.1 320.7 1.0X +codegen = T 5885 5947 79 3.6 280.6 1.1X + + +================================================================================================ +sort aggregate with decimal key +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 9V74 80-Core Processor +sort agg w decimal key: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 3371 3426 78 6.2 160.7 1.0X +codegen = T 2530 2563 41 8.3 120.6 1.3X + + +================================================================================================ +sort aggregate with multiple key types +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 9V74 80-Core Processor +sort agg w multiple keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 6586 6604 26 3.2 314.0 1.0X +codegen = T 6035 6050 13 3.5 287.8 1.1X + + diff --git a/sql/core/benchmarks/SortAggregateBenchmark-jdk25-results.txt b/sql/core/benchmarks/SortAggregateBenchmark-jdk25-results.txt new file mode 100644 index 0000000000000..12cc5a0fe825b --- /dev/null +++ b/sql/core/benchmarks/SortAggregateBenchmark-jdk25-results.txt @@ -0,0 +1,72 @@ +================================================================================================ +sort aggregate without grouping +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w/o group: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 16207 17142 1323 32.4 30.9 1.0X +codegen = T 415 422 5 1264.8 0.8 39.1X + + +================================================================================================ +sort aggregate with linear keys +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w linear keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 19884 20041 222 4.2 237.0 1.0X +codegen = T 15733 16010 468 5.3 187.6 1.3X + + +================================================================================================ +sort aggregate with randomized keys +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w randomized keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 28860 28899 56 2.9 344.0 1.0X +codegen = T 25706 25781 58 3.3 306.4 1.1X + + +================================================================================================ +sort aggregate with string key +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w string key: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 6927 6934 11 3.0 330.3 1.0X +codegen = T 5836 5902 45 3.6 278.3 1.2X + + +================================================================================================ +sort aggregate with decimal key +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w decimal key: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 3182 3288 150 6.6 151.7 1.0X +codegen = T 2390 2438 45 8.8 113.9 1.3X + + +================================================================================================ +sort aggregate with multiple key types +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w multiple keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 6647 6662 21 3.2 317.0 1.0X +codegen = T 6057 6080 27 3.5 288.8 1.1X + + diff --git a/sql/core/benchmarks/SortAggregateBenchmark-results.txt b/sql/core/benchmarks/SortAggregateBenchmark-results.txt new file mode 100644 index 0000000000000..422de3a6713fe --- /dev/null +++ b/sql/core/benchmarks/SortAggregateBenchmark-results.txt @@ -0,0 +1,72 @@ +================================================================================================ +sort aggregate without grouping +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w/o group: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 17331 17444 160 30.3 33.1 1.0X +codegen = T 881 885 3 595.1 1.7 19.7X + + +================================================================================================ +sort aggregate with linear keys +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w linear keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 20923 21112 267 4.0 249.4 1.0X +codegen = T 17299 18455 680 4.8 206.2 1.2X + + +================================================================================================ +sort aggregate with randomized keys +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w randomized keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 30049 30198 210 2.8 358.2 1.0X +codegen = T 26696 26804 101 3.1 318.2 1.1X + + +================================================================================================ +sort aggregate with string key +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w string key: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 6220 6264 63 3.4 296.6 1.0X +codegen = T 5632 5682 95 3.7 268.6 1.1X + + +================================================================================================ +sort aggregate with decimal key +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w decimal key: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 3220 3239 28 6.5 153.5 1.0X +codegen = T 2482 2529 56 8.4 118.4 1.3X + + +================================================================================================ +sort aggregate with multiple key types +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1018-azure +AMD EPYC 7763 64-Core Processor +sort agg w multiple keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = F 6378 6379 2 3.3 304.1 1.0X +codegen = T 5660 5754 70 3.7 269.9 1.1X + + diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggregateCodegenSupport.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggregateCodegenSupport.scala index 352388a6d8a10..884a8188a5aef 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggregateCodegenSupport.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggregateCodegenSupport.scala @@ -45,9 +45,10 @@ trait AggregateCodegenSupport /** * The variables are used as aggregation buffers and each aggregate function has one or more - * ExprCode to initialize its buffer slots. Only used for aggregation without keys. + * ExprCode to initialize its buffer slots. Used for aggregation without keys, and for sort-based + * aggregation with keys (where a single group is aggregated at a time). */ - private var bufVars: Seq[Seq[ExprCode]] = _ + protected var bufVars: Seq[Seq[ExprCode]] = _ /** * Whether this operator needs to build hash table. @@ -94,13 +95,13 @@ trait AggregateCodegenSupport override def usedInputs: AttributeSet = inputSet /** - * The generated code for `doProduce` call when aggregate does not have grouping keys. + * Creates global mutable state variables to hold the aggregation buffer, one nested sequence of + * `ExprCode` per aggregate function. The returned `ExprCode`s carry the code that (re)initializes + * the buffer slots with the aggregate functions' initial values; running that code resets the + * buffer for a new group. The result is also stored in `bufVars` so that `doConsume` can update + * the same variables. */ - private def doProduceWithoutKeys(ctx: CodegenContext): String = { - val initAgg = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "initAgg") - // The generated function doesn't have input row in the code context. - ctx.INPUT_ROW = null - + protected def createAggBufVars(ctx: CodegenContext): Seq[Seq[ExprCode]] = { // generate variables for aggregation buffer val functions = aggregateExpressions.map(_.aggregateFunction.asInstanceOf[DeclarativeAggregate]) val initExpr = functions.map(f => f.initialValues) @@ -121,9 +122,21 @@ trait AggregateCodegenSupport JavaCode.global(value, e.dataType)) } } - val flatBufVars = bufVars.flatten - val initBufVar = evaluateVariables(flatBufVars) + bufVars + } + /** + * The generated code for `doProduce` call when aggregate does not have grouping keys. + */ + private def doProduceWithoutKeys(ctx: CodegenContext): String = { + val initAgg = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "initAgg") + // The generated function doesn't have input row in the code context. + ctx.INPUT_ROW = null + // generate variables for aggregation buffer + val flatBufVars = createAggBufVars(ctx).flatten + val initBufVar = evaluateVariables(flatBufVars) + val functions = + aggregateExpressions.map(_.aggregateFunction.asInstanceOf[DeclarativeAggregate]) // generate variables for output val (resultVars, genResult) = if (modes.contains(Final) || modes.contains(Complete)) { // evaluate aggregate results @@ -195,6 +208,15 @@ trait AggregateCodegenSupport * The generated code for `doConsume` call when aggregate does not have grouping keys. */ private def doConsumeWithoutKeys(ctx: CodegenContext, input: Seq[ExprCode]): String = { + generateAggBufferUpdateCode(ctx, input) + } + + /** + * The generated code that evaluates the aggregate functions for one input row and updates the + * aggregation buffer held in `bufVars`. Shared by the no-keys path and the sort-based with-keys + * path, both of which keep the aggregation buffer in global mutable state variables. + */ + protected def generateAggBufferUpdateCode(ctx: CodegenContext, input: Seq[ExprCode]): String = { // only have DeclarativeAggregate val functions = aggregateExpressions.map(_.aggregateFunction.asInstanceOf[DeclarativeAggregate]) val inputAttrs = functions.flatMap(_.aggBufferAttributes) ++ inputAttributes diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala index 06f87af50eb5d..80afd61a98f6e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala @@ -17,14 +17,15 @@ package org.apache.spark.sql.execution.aggregate -import org.apache.spark.SparkUnsupportedOperationException import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.BindReferences.bindReferences import org.apache.spark.sql.catalyst.expressions.aggregate._ -import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode} +import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode, GenerateUnsafeProjection} +import org.apache.spark.sql.catalyst.util.UnsafeRowUtils import org.apache.spark.sql.catalyst.util.truncatedString -import org.apache.spark.sql.execution.{OrderPreservingUnaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.{CodegenSupport, OrderPreservingUnaryExecNode, SparkPlan} import org.apache.spark.sql.execution.metric.SQLMetrics import org.apache.spark.sql.internal.SQLConf @@ -94,19 +95,212 @@ case class SortAggregateExec( } override def supportCodegen: Boolean = { - // TODO(SPARK-32750): Support sort aggregate code-gen with grouping keys super.supportCodegen && conf.getConf(SQLConf.ENABLE_SORT_AGGREGATE_CODEGEN) && - groupingExpressions.isEmpty + (groupingExpressions.isEmpty || supportCodegenWithKeys) + } + + private def supportCodegenWithKeys: Boolean = { + conf.getConf(SQLConf.ENABLE_SORT_AGGREGATE_CODEGEN_WITH_KEYS) && + groupingExpressions.forall(e => UnsafeRowUtils.isBinaryStable(e.dataType)) } protected override def needHashTable: Boolean = false + // For the with-keys path, results are produced incrementally while scanning the sorted input: a + // group's result is emitted (appended to the output buffer) as soon as the next group starts. The + // child's producing loop must therefore honor `shouldStop()`, so it yields right after a result + // row is buffered and resumes scanning on the next `processNext` call (see `doProduceWithKeys`). + // Otherwise the child would scan the whole partition in one go, and every emitted row would alias + // 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 + + override protected def canCheckLimitNotReached: Boolean = false + + // The global UnsafeRow holding the grouping key of the group currently being aggregated. + private var currentGroupingKeyTerm: String = _ + + // The global boolean flag indicating whether the current group has been started, i.e. at least + // one input row has been processed. + private var initGroupTerm: String = _ + + // The code that (re)initializes the aggregation buffer variables to the initial values of the + // aggregate functions. Used to reset the buffer when a new group starts. + private var reInitBufferCode: String = _ + + // The name of the generated function that outputs the result of the current group. + private var outputFuncName: String = _ + + /** + * Generate the code for output. The aggregation buffer is held in the global `bufVars` and the + * grouping key in `currentGroupingKeyTerm`, both populated while scanning the current group. + * @return function name for the result code. + */ + private def generateResultFunctionForKeys(ctx: CodegenContext): String = { + val funcName = ctx.freshName("doAggregateWithKeysOutput") + val numOutput = metricTerm(ctx, "numOutputRows") + val flatBufVars = bufVars.flatten + val groupingAttributes = groupingExpressions.map(_.toAttribute) + + val body = + if (modes.contains(Final) || modes.contains(Complete)) { + // generate output using resultExpressions + ctx.currentVars = null + ctx.INPUT_ROW = currentGroupingKeyTerm + val keyVars = groupingExpressions.zipWithIndex.map { case (e, i) => + BoundReference(i, e.dataType, e.nullable).genCode(ctx) + } + val evaluateKeyVars = evaluateVariables(keyVars) + // evaluate the aggregation result from the buffer variables + ctx.currentVars = flatBufVars + ctx.INPUT_ROW = null + val functions = + aggregateExpressions.map(_.aggregateFunction.asInstanceOf[DeclarativeAggregate]) + val aggResults = bindReferences( + functions.map(_.evaluateExpression), + aggregateBufferAttributes).map(_.genCode(ctx)) + val evaluateAggResults = evaluateVariables(aggResults) + // generate the final result + ctx.currentVars = keyVars ++ aggResults + val inputAttrs = groupingAttributes ++ aggregateAttributes + val resultVars = bindReferences[Expression]( + resultExpressions, + inputAttrs).map(_.genCode(ctx)) + val evaluateNondeterministicResults = + evaluateNondeterministicVariables(output, resultVars, resultExpressions) + s""" + |$evaluateKeyVars + |$evaluateAggResults + |$evaluateNondeterministicResults + |${consume(ctx, resultVars)} + """.stripMargin + } else if (modes.contains(Partial) || modes.contains(PartialMerge)) { + // resultExpressions are Attributes of groupingExpressions and aggregateBufferAttributes. + assert(resultExpressions.forall(_.isInstanceOf[Attribute])) + assert(resultExpressions.length == + groupingExpressions.length + aggregateBufferAttributes.length) + + ctx.currentVars = null + ctx.INPUT_ROW = currentGroupingKeyTerm + val keyVars = groupingExpressions.zipWithIndex.map { case (e, i) => + BoundReference(i, e.dataType, e.nullable).genCode(ctx) + } + val evaluateKeyVars = evaluateVariables(keyVars) + + // the aggregation buffer values are output directly + ctx.currentVars = keyVars ++ flatBufVars + ctx.INPUT_ROW = null + val inputAttrs = resultExpressions.map(_.toAttribute) + val resultVars = bindReferences[Expression]( + resultExpressions, + inputAttrs).map(_.genCode(ctx)) + s""" + |$evaluateKeyVars + |${consume(ctx, resultVars)} + """.stripMargin + } else { + // generate result based on grouping key + ctx.INPUT_ROW = currentGroupingKeyTerm + ctx.currentVars = null + val resultVars = bindReferences[Expression]( + resultExpressions, + groupingAttributes).map(_.genCode(ctx)) + val evaluateNondeterministicResults = + evaluateNondeterministicVariables(output, resultVars, resultExpressions) + s""" + |$evaluateNondeterministicResults + |${consume(ctx, resultVars)} + """.stripMargin + } + ctx.addNewFunction(funcName, + s""" + |private void $funcName() throws java.io.IOException { + | $numOutput.add(1); + | $body + |} + """.stripMargin) + } + protected override def doProduceWithKeys(ctx: CodegenContext): String = { - throw new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3170") + ctx.INPUT_ROW = null + // Generate the global variables for the aggregation buffer of the current group. Capture the + // buffer initialization code (and clear it from `bufVars`, so the result/update code does not + // re-emit it); it is reused in `doConsumeWithKeys` to reset the buffer when a new group starts. + reInitBufferCode = evaluateVariables(createAggBufVars(ctx).flatten) + // Global state to track the current group. Inline the grouping-key row (rather than letting it + // be compacted into a shared array) so its term is a plain variable name: it is used as + // `ctx.INPUT_ROW` when generating the output, and an array-subscript term there would break the + // generated code when key expressions are extracted into split functions. + currentGroupingKeyTerm = + ctx.addMutableState("UnsafeRow", "currentGroupingKey", forceInline = true) + initGroupTerm = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "initGroup") + // Whether the whole sorted input has been consumed. + val noMoreInputTerm = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "noMoreInputTerm") + // Generate the output function before `child.produce`, so that `doConsumeWithKeys` can call it + // when it detects a group boundary. + outputFuncName = generateResultFunctionForKeys(ctx) + + val doAgg = ctx.freshName("doAggregateWithKeys") + val doAggFuncName = ctx.addNewFunction(doAgg, + s""" + |private void $doAgg() throws java.io.IOException { + | ${child.asInstanceOf[CodegenSupport].produce(ctx, this)} + |} + """.stripMargin) + + // Sort-based aggregation consumes the sorted input row by row, emitting a group's result as + // soon as the next group starts (see `doConsumeWithKeys`). Emitting a row appends to the output + // buffer, which makes the child's `shouldStop()` return true, so the child's producing loop + // returns to here mid-scan. We therefore must be able to resume scanning across multiple + // `processNext` invocations: + // - `$noMoreInputTerm` guards against re-running once the input is fully consumed; + // - after `$doAggFuncName()` returns, `shouldStop()` distinguishes a real end-of-input (the + // child's loop exhausted, so `shouldStop()` is false) from a mid-scan pause (`shouldStop()` + // is true because an output row is buffered). Only on a real end-of-input do we mark the + // scan done and flush the last group. If the last input row also triggered an output (so + // `shouldStop()` is still true at exhaustion), the next `processNext` re-enters, the child's + // loop produces nothing, `shouldStop()` is then false, and the last group is flushed. + s""" + |if (!$noMoreInputTerm) { + | $doAggFuncName(); + | if (!shouldStop()) { + | $noMoreInputTerm = true; + | if ($initGroupTerm) { + | $outputFuncName(); + | } + | } + |} + """.stripMargin } protected override def doConsumeWithKeys(ctx: CodegenContext, input: Seq[ExprCode]): String = { - throw new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3170") + // Create the grouping key. `ctx.currentVars` is still set to `input` here. + val groupingUnsafeRowKeyCode = GenerateUnsafeProjection.createCode( + ctx, bindReferences[Expression](groupingExpressions, child.output)) + val groupingUnsafeRowKey = groupingUnsafeRowKeyCode.value + // The code to update the aggregation buffer with the current input row. + val updateBufferCode = generateAggBufferUpdateCode(ctx, input) + + // `reInitBufferCode` was captured in `doProduceWithKeys`; it (re)initializes the aggregation + // buffer variables to the aggregate functions' initial values, resetting them for a new group. + // The input is sorted by the grouping key, so rows of the same group are contiguous. When the + // grouping key changes, the current group is complete: output it and reset the buffer for the + // new group. Group equality uses the binary representation of the key, which is valid because + // `supportCodegen` restricts grouping keys to binary-stable types. + s""" + |${groupingUnsafeRowKeyCode.code} + |if (!$initGroupTerm) { + | $initGroupTerm = true; + | $currentGroupingKeyTerm = $groupingUnsafeRowKey.copy(); + | $reInitBufferCode + |} else if (!$currentGroupingKeyTerm.equals($groupingUnsafeRowKey)) { + | $outputFuncName(); + | $currentGroupingKeyTerm = $groupingUnsafeRowKey.copy(); + | $reInitBufferCode + |} + |$updateBufferCode + """.stripMargin } override def simpleString(maxFields: Int): String = toString(verbose = false, maxFields) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala index bcd2f5369932e..e908e5e35a7c8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala @@ -68,6 +68,106 @@ class WholeStageCodegenSuite extends SharedSparkSession } } + // Runs `query` on `data` with sort aggregate forced and its code-gen enabled, asserts the plan + // actually uses a code-gen'd SortAggregateExec, and checks the result matches the interpreted + // (code-gen disabled) result. + private def checkSortAggregateCodegen( + data: Dataset[Row])(query: Dataset[Row] => Dataset[Row]): Unit = { + // Disable both hash-based aggregate operators so the planner always picks SortAggregateExec. + val forceSortAggregate = Seq( + SQLConf.USE_HASH_AGG.key -> "false", + SQLConf.USE_OBJECT_HASH_AGG.key -> "false") + val expected = withSQLConf( + (forceSortAggregate :+ (SQLConf.ENABLE_SORT_AGGREGATE_CODEGEN.key -> "false")): _*) { + val df = query(data) + 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}") + df.collect() + } + withSQLConf( + (forceSortAggregate :+ (SQLConf.ENABLE_SORT_AGGREGATE_CODEGEN.key -> "true")): _*) { + val df = query(data) + 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}") + checkAnswer(df, expected) + } + } + + test("SPARK-32750: SortAggregate code-gen with grouping keys") { + val data = spark.range(200).selectExpr( + "id", + "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") + + // Exercise a variety of shapes: multiple/numeric/string grouping keys, null keys and values, + // single-row groups, a single all-rows group, and a downstream limit (which exercises the + // resumable `shouldStop` path in the generated produce loop). + checkSortAggregateCodegen(data) { + _.groupBy("k1", "k2") + .agg(count(col("v")), sum(col("v")), max(col("v")), min(col("v"))) + .orderBy("k1", "k2") + } + // expression grouping keys, aggregates over expressions, and arithmetic on the results + checkSortAggregateCodegen(data) { + _.groupBy((col("k1") + col("k2")).as("k")) + .agg( + (sum(col("v") * lit(2)) + lit(1)).as("weighted"), + (max(col("v")) - min(col("v"))).as("spread"), + count(col("v")).as("cnt")) + .orderBy("k") + } + // aggregates with FILTER (WHERE) clauses + checkSortAggregateCodegen(data) { + _.groupBy("k1") + .agg( + expr("sum(v) FILTER (WHERE v > 50)"), + expr("count(v) FILTER (WHERE s IS NOT NULL)"), + avg(col("v"))) + .orderBy("k1") + } + // string grouping key with nulls, followed by a HAVING-style filter on the aggregate + checkSortAggregateCodegen(data) { + _.groupBy("s").agg(count(col("v")).as("cnt"), sum(col("v")).as("total")) + .where(col("cnt") > 2) + .orderBy("s") + } + // count(distinct ...): rewritten to a two-round aggregation, both of which are sort aggregates + checkSortAggregateCodegen(data) { + _.groupBy("k2").agg(countDistinct(col("v")), sum(col("v"))).orderBy("k2") + } + // every row is its own group + checkSortAggregateCodegen(data)(_.groupBy("id").agg(max(col("v"))).orderBy("id")) + // a single group covering all rows + checkSortAggregateCodegen(data)(_.groupBy(lit(1)).agg(sum(col("v")), avg(col("v")))) + // downstream limit on top of the aggregate + checkSortAggregateCodegen(data)(_.groupBy("k1").agg(sum(col("v"))).orderBy("k1").limit(3)) + } + + test("SPARK-32750: SortAggregate code-gen with grouping keys - empty input") { + // No rows: a grouped aggregate over empty input must produce no output rows. + val data = spark.range(0).selectExpr("id", "id % 3 as k", "id as v") + checkSortAggregateCodegen(data)(_.groupBy("k").agg(sum(col("v")), count(col("v"))).orderBy("k")) + } + + test("SPARK-32750: SortAggregate code-gen with grouping keys - single partition") { + // Force a single partition so all groups are produced within one task's scan. + val data = spark.range(50).repartition(1).selectExpr("id", "id % 4 as k", "id as v") + checkSortAggregateCodegen(data) { + _.groupBy((col("k") + lit(1)).as("k")) + .agg( + (sum(col("v")) + max(col("v"))).as("mixed"), + avg(col("v") * col("v")).as("avg_sq"), + expr("count(v) FILTER (WHERE v % 2 = 0)").as("evens")) + .orderBy("k") + } + } + testWithWholeStageCodegenOnAndOff("GenerateExec should be" + " included in WholeStageCodegen") { codegenEnabled => import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/SortAggregateBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/SortAggregateBenchmark.scala new file mode 100644 index 0000000000000..b8b9af3ea3101 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/SortAggregateBenchmark.scala @@ -0,0 +1,148 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.benchmark + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.internal.SQLConf + +/** + * Benchmark to measure performance for sort-based aggregate, focusing on the whole-stage + * code-gen path. Hash-map based aggregation is intentionally excluded; both + * `spark.sql.execution.useHashAggregateExec` and `spark.sql.execution.useObjectHashAggregateExec` + * are disabled so the planner always picks + * [[org.apache.spark.sql.execution.aggregate.SortAggregateExec]]. + * + * To run this benchmark: + * {{{ + * 1. without sbt: bin/spark-submit --class + * --jars , + * 2. build/sbt "sql/Test/runMain " + * 3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/Test/runMain " + * Results will be written to "benchmarks/SortAggregateBenchmark-results.txt". + * }}} + */ +object SortAggregateBenchmark extends SqlBasedBenchmark { + + // Force the planner to pick SortAggregateExec by disabling both hash-based aggregate operators. + private val forceSortAggregate: Map[String, String] = Map( + SQLConf.USE_HASH_AGG.key -> "false", + SQLConf.USE_OBJECT_HASH_AGG.key -> "false") + + /** + * Adds the two cases we care about for a sort aggregate. Whole-stage code-gen stays enabled in + * both so the child pipeline (scan, sort) is code-gen'd either way; only the sort aggregate's own + * code-gen is toggled via `ENABLE_SORT_AGGREGATE_CODEGEN`, which isolates its contribution: + * - code-gen off: sort aggregate falls back to the interpreted `SortBasedAggregationIterator`; + * - code-gen on: code-gen'd `SortAggregateExec`. + */ + private def addGroupingKeyCases(benchmark: Benchmark)(f: () => Unit): Unit = { + benchmark.addCase("codegen = F", numIters = 2) { _ => + withSQLConf( + (forceSortAggregate ++ Map( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.ENABLE_SORT_AGGREGATE_CODEGEN.key -> "false")).toSeq: _*) { + f() + } + } + + benchmark.addCase("codegen = T", numIters = 5) { _ => + withSQLConf( + (forceSortAggregate ++ Map( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.ENABLE_SORT_AGGREGATE_CODEGEN.key -> "true")).toSeq: _*) { + f() + } + } + } + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + runBenchmark("sort aggregate without grouping") { + val N = 500L << 20 + val benchmark = new Benchmark("sort agg w/o group", N, output = output) + addGroupingKeyCases(benchmark) { () => + spark.range(N).selectExpr("sum(id)").noop() + } + benchmark.run() + } + + runBenchmark("sort aggregate with linear keys") { + 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. + spark.range(N).selectExpr("id", "(id & 65535) as k") + .sortWithinPartitions("k").createOrReplaceTempView("linear") + addGroupingKeyCases(benchmark) { () => + spark.sql("select k, count(*), sum(id), max(id) from linear group by k").noop() + } + benchmark.run() + } + + runBenchmark("sort aggregate with randomized keys") { + val N = 20 << 22 + val benchmark = new Benchmark("sort agg w randomized keys", N, output = output) + spark.range(N).selectExpr("id", "floor(rand() * 10000) as k") + .sortWithinPartitions("k").createOrReplaceTempView("rand_keys") + addGroupingKeyCases(benchmark) { () => + spark.sql("select k, count(*), sum(id), max(id) from rand_keys group by k").noop() + } + benchmark.run() + } + + runBenchmark("sort aggregate with string key") { + val N = 20 << 20 + val benchmark = new Benchmark("sort agg w string key", N, output = output) + spark.range(N).selectExpr("id", "cast(id & 1023 as string) as k") + .sortWithinPartitions("k").createOrReplaceTempView("string_key") + addGroupingKeyCases(benchmark) { () => + spark.sql("select k, count(*), sum(id), max(id) from string_key group by k").noop() + } + benchmark.run() + } + + runBenchmark("sort aggregate with decimal key") { + val N = 20 << 20 + val benchmark = new Benchmark("sort agg w decimal key", N, output = output) + spark.range(N).selectExpr("id", "cast(id & 65535 as decimal(18, 0)) as k") + .sortWithinPartitions("k").createOrReplaceTempView("decimal_key") + addGroupingKeyCases(benchmark) { () => + spark.sql("select k, count(*), sum(id), max(id) from decimal_key group by k").noop() + } + benchmark.run() + } + + runBenchmark("sort aggregate with multiple key types") { + val N = 20 << 20 + val benchmark = new Benchmark("sort agg w multiple keys", N, output = output) + spark.range(N) + .selectExpr( + "id", + "(id & 1023) as k1", + "cast(id & 1023 as string) as k2", + "cast(id & 1023 as int) as k3", + "id > 1023 as k4") + .sortWithinPartitions("k1", "k2", "k3", "k4") + .createOrReplaceTempView("multi_keys") + addGroupingKeyCases(benchmark) { () => + spark.sql("select k1, k2, k3, k4, count(*), sum(id), max(id) " + + "from multi_keys group by k1, k2, k3, k4").noop() + } + benchmark.run() + } + } +} From 06fe0911c1a80d5905cadad8eb6c1e08676c3324 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Thu, 9 Jul 2026 14:10:58 +0800 Subject: [PATCH 2/2] withBindingPolicy --- .../src/main/scala/org/apache/spark/sql/internal/SQLConf.scala | 1 + 1 file changed, 1 insertion(+) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index c089c1690920b..d23f16391681c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -3705,6 +3705,7 @@ object SQLConf { .doc("When true, enable code-gen for sort aggregate with grouping keys. Takes effect only " + s"when ${ENABLE_SORT_AGGREGATE_CODEGEN.key} is enabled.") .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) .booleanConf .createWithDefault(true)