ulysses-you commented on code in PR #57153:
URL: https://github.com/apache/spark/pull/57153#discussion_r3583862636
##########
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:
Good point. Added a paragraph to the comment block spelling this out: we
keep the inherited `needCopyResult = false`, but its original justification no
longer applies once the with-keys path streams results mid-scan. It stays safe
for a different reason -- the stop-check discipline guarantees at most one
output row is buffered before it's consumed, so the reused output row is never
aliased, and any in-stage parent that multiplies rows (e.g. a join) declares
`needCopyResult = true` itself. Done in 874cc2a.
##########
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:
Added a "non-binary collated key" test that groups by a `UTF8_LCASE` string
key (not binary-stable), asserting no code-gen'd `SortAggregateExec` appears in
the plan (the interpreted fallback is used) while the collation-aware result
stays correct ('a'/'A' and 'b'/'B' each collapse to one group, 'c' stays
alone). This pins the `isBinaryStable = false` half of the gate. Done in
874cc2a.
##########
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:
Reworded along those lines: the temp view is lazy so the sort re-executes
each iteration (and without the explicit `sortWithinPartitions` the planner
would insert an equivalent `SortExec` below the aggregate anyway); pinning the
sort just makes both cases pay the same sort cost, isolating the aggregate's
own contribution -- which means the reported speedup understates the
aggregate-only gain. Done in 874cc2a.
--
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]
