cloud-fan commented on code in PR #57153:
URL: https://github.com/apache/spark/pull/57153#discussion_r3571616404
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala:
##########
@@ -68,6 +68,192 @@ 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}")
Review 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.
```suggestion
s"Expected no code-gen'd SortAggregateExec
in:\n${df.queryExecution.executedPlan}")
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala:
##########
@@ -94,19 +95,215 @@ 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
Review 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.
```suggestion
override protected def canCheckLimitNotReached: Boolean =
groupingExpressions.nonEmpty
```
--
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]
