viirya commented on code in PR #57153:
URL: https://github.com/apache/spark/pull/57153#discussion_r3580898032
##########
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
Review Comment:
The comment block here explains why `needStopCheck` must flip for the
with-keys path, but the operator still inherits `needCopyResult = false` from
`BlockingOperatorWithCodegen`, whose stated justification ("blocking operators
keep the data in some buffer") no longer applies once this path streams results
out mid-scan. It's still safe, but for a different reason: the stop-check
discipline guarantees at most one output row is buffered before it's consumed,
and any in-stage parent that multiplies rows (e.g. joins) declares
`needCopyResult = true` itself. Since that argument is non-obvious and this
comment is where a future reader would look, suggest adding a sentence covering
`needCopyResult` as well.
##########
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 no 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")
+
+ // Exercise a variety of shapes: multiple/numeric/string grouping keys,
null keys and values,
+ // single-row groups, a single all-rows group, and a downstream limit
(which exercises the
+ // resumable `shouldStop` path in the generated produce loop).
+ checkSortAggregateCodegen(data) {
+ _.groupBy("k1", "k2")
+ .agg(count(col("v")), sum(col("v")), max(col("v")), min(col("v")))
+ .orderBy("k1", "k2")
+ }
+ // expression grouping keys, aggregates over expressions, and arithmetic
on the results
+ checkSortAggregateCodegen(data) {
+ _.groupBy((col("k1") + col("k2")).as("k"))
+ .agg(
+ (sum(col("v") * lit(2)) + lit(1)).as("weighted"),
+ (max(col("v")) - min(col("v"))).as("spread"),
+ count(col("v")).as("cnt"))
+ .orderBy("k")
+ }
+ // aggregates with FILTER (WHERE) clauses
+ checkSortAggregateCodegen(data) {
+ _.groupBy("k1")
+ .agg(
+ expr("sum(v) FILTER (WHERE v > 50)"),
+ expr("count(v) FILTER (WHERE s IS NOT NULL)"),
+ avg(col("v")))
+ .orderBy("k1")
+ }
+ // string grouping key with nulls, followed by a HAVING-style filter on
the aggregate
+ checkSortAggregateCodegen(data) {
+ _.groupBy("s").agg(count(col("v")).as("cnt"), sum(col("v")).as("total"))
+ .where(col("cnt") > 2)
+ .orderBy("s")
+ }
+ // count(distinct ...): rewritten to a two-round aggregation, both of
which are sort aggregates
+ checkSortAggregateCodegen(data) {
+ _.groupBy("k2").agg(countDistinct(col("v")), sum(col("v"))).orderBy("k2")
+ }
+ // every row is its own group
+
checkSortAggregateCodegen(data)(_.groupBy("id").agg(max(col("v"))).orderBy("id"))
+ // a single group covering all rows
+ checkSortAggregateCodegen(data)(_.groupBy(lit(1)).agg(sum(col("v")),
avg(col("v"))))
+ // downstream limit on top of the aggregate
+
checkSortAggregateCodegen(data)(_.groupBy("k1").agg(sum(col("v"))).orderBy("k1").limit(3))
+ }
+
+ test("SPARK-32750: SortAggregate code-gen with grouping keys - empty input")
{
+ // No rows: a grouped aggregate over empty input must produce no output
rows.
+ val data = spark.range(0).selectExpr("id", "id % 3 as k", "id as v")
+ checkSortAggregateCodegen(data)(_.groupBy("k").agg(sum(col("v")),
count(col("v"))).orderBy("k"))
+ }
+
+ test("SPARK-32750: SortAggregate code-gen with grouping keys - single
partition") {
+ // Force a single partition so all groups are produced within one task's
scan.
+ val data = spark.range(50).repartition(1).selectExpr("id", "id % 4 as k",
"id as v")
+ checkSortAggregateCodegen(data) {
+ _.groupBy((col("k") + lit(1)).as("k"))
+ .agg(
+ (sum(col("v")) + max(col("v"))).as("mixed"),
+ avg(col("v") * col("v")).as("avg_sq"),
+ expr("count(v) FILTER (WHERE v % 2 = 0)").as("evens"))
+ .orderBy("k")
+ }
+ }
+
+ test("SPARK-32750: SortAggregate code-gen with float/double grouping keys") {
Review Comment:
All the new tests exercise the `supportCodegenWithKeys` gate in the true
direction; the `isBinaryStable = false` fallback is never hit (the config-gate
test disables the path via conf, not via the type check). Suggest one case with
a non-binary collated string grouping key (e.g. `UTF8_LCASE`), asserting no
code-gen'd `SortAggregateExec` appears in the plan while the result stays
correct - that pins the type-based fallback to the interpreted path, which is
the half of the gate a future type-support change is most likely to break.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/SortAggregateBenchmark.scala:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.benchmark
+
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Benchmark to measure performance for sort-based aggregate, focusing on the
whole-stage
+ * code-gen path. Hash-map based aggregation is intentionally excluded; both
+ * `spark.sql.execution.useHashAggregateExec` and
`spark.sql.execution.useObjectHashAggregateExec`
+ * are disabled so the planner always picks
+ * [[org.apache.spark.sql.execution.aggregate.SortAggregateExec]].
+ *
+ * To run this benchmark:
+ * {{{
+ * 1. without sbt: bin/spark-submit --class <this class>
+ * --jars <spark core test jar>,<spark catalyst test jar> <spark sql test
jar>
+ * 2. build/sbt "sql/Test/runMain <this class>"
+ * 3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt
"sql/Test/runMain <this class>"
+ * Results will be written to
"benchmarks/SortAggregateBenchmark-results.txt".
+ * }}}
+ */
+object SortAggregateBenchmark extends SqlBasedBenchmark {
+
+ // Force the planner to pick SortAggregateExec by disabling both hash-based
aggregate operators.
+ private val forceSortAggregate: Map[String, String] = Map(
+ SQLConf.USE_HASH_AGG.key -> "false",
+ SQLConf.USE_OBJECT_HASH_AGG.key -> "false")
+
+ /**
+ * Adds the two cases we care about for a sort aggregate. Whole-stage
code-gen stays enabled in
+ * both so the child pipeline (scan, sort) is code-gen'd either way; only
the sort aggregate's own
+ * code-gen is toggled via `ENABLE_SORT_AGGREGATE_CODEGEN`, which isolates
its contribution:
+ * - code-gen off: sort aggregate falls back to the interpreted
`SortBasedAggregationIterator`;
+ * - code-gen on: code-gen'd `SortAggregateExec`.
+ */
+ private def addGroupingKeyCases(benchmark: Benchmark)(f: () => Unit): Unit =
{
+ benchmark.addCase("codegen = F", numIters = 2) { _ =>
+ withSQLConf(
+ (forceSortAggregate ++ Map(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.ENABLE_SORT_AGGREGATE_CODEGEN.key -> "false")).toSeq: _*) {
+ f()
+ }
+ }
+
+ benchmark.addCase("codegen = T", numIters = 5) { _ =>
+ withSQLConf(
+ (forceSortAggregate ++ Map(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.ENABLE_SORT_AGGREGATE_CODEGEN.key -> "true")).toSeq: _*) {
+ f()
+ }
+ }
+ }
+
+ override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
+ runBenchmark("sort aggregate without grouping") {
+ val N = 500L << 20
+ val benchmark = new Benchmark("sort agg w/o group", N, output = output)
+ addGroupingKeyCases(benchmark) { () =>
+ spark.range(N).selectExpr("sum(id)").noop()
+ }
+ benchmark.run()
+ }
+
+ runBenchmark("sort aggregate with linear keys") {
+ val N = 20 << 22
+ val benchmark = new Benchmark("sort agg w linear keys", N, output =
output)
+ // The child of a sort aggregate must be sorted by the grouping keys.
Sorting the whole input
+ // dominates, so pre-sort once into a temp view and aggregate over the
already-ordered rows.
Review Comment:
Nit: "pre-sort once into a temp view" reads as if the sort cost is paid
once, but a temp view is lazy - the sort re-executes on every iteration (and
without the explicit `sortWithinPartitions` the planner would insert an
equivalent `SortExec` below the aggregate anyway). What the setup actually
achieves is that both cases pay the same sort cost, isolating the aggregate's
own contribution - which also means the reported 1.2~1.3x understates the
aggregate-only speedup. Suggest rewording the comment along those lines.
--
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]
