ulysses-you commented on code in PR #57153:
URL: https://github.com/apache/spark/pull/57153#discussion_r3568404163
##########
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:
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.
##########
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")
Review 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-gen`d `SortAggregateExec` when disabled while the result
stays correct.
Also updated the "How was this patch tested?" section so it matches what is
actually covered.
--
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]
