LuciferYang commented on code in PR #57153:
URL: https://github.com/apache/spark/pull/57153#discussion_r3556710033
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala:
##########
@@ -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 {
Review 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)`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala:
##########
@@ -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")
Review 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. `
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala:
##########
@@ -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 {
Review 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.
##########
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(
Review 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
