Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ulysses-you wants to merge 2 commits into
base: master
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
090e25e
[SPARK-32750][SQL] Support whole-stage codegen for SortAggregate with…
ulysses-you Jul 9, 2026
06fe091
withBindingPolicy
ulysses-you Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3699,6 +3699,16 @@ 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")
.withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
.booleanConf
.createWithDefault(true)

val ENABLE_FULL_OUTER_SHUFFLED_HASH_JOIN_CODEGEN =
buildConf("spark.sql.codegen.join.fullOuterShuffledHashJoin.enabled")
.internal()
Expand Down
72 changes: 72 additions & 0 deletions sql/core/benchmarks/SortAggregateBenchmark-jdk21-results.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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


72 changes: 72 additions & 0 deletions sql/core/benchmarks/SortAggregateBenchmark-jdk25-results.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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


72 changes: 72 additions & 0 deletions sql/core/benchmarks/SortAggregateBenchmark-results.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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


42 changes: 32 additions & 10 deletions ...ore/src/main/scala/org/apache/spark/sql/execution/aggregate/AggregateCodegenSupport.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading