ulysses-you commented on code in PR #57153:
URL: https://github.com/apache/spark/pull/57153#discussion_r3577455915
##########
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:
Done
##########
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:
Done
--
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]
