sunchao commented on code in PR #57276: URL: https://github.com/apache/spark/pull/57276#discussion_r3597481632
########## sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala: ########## @@ -0,0 +1,214 @@ +/* + * 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 + +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +class CombineAdjacentAggregationSuite extends QueryTest + with SharedSparkSession + with AdaptiveSparkPlanHelper { + + private def numAggregates(query: String): Int = { + val df = sql(query) + df.collect() + collect(df.queryExecution.executedPlan) { + case agg: BaseAggregateExec => agg + }.size + } + + private def checkNumAggregation( + query: String, + numAggWithDisabled: Int, + numAggWithEnabled: Int): Unit = { + var expectedResult: Array[Row] = null + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") { + val df = sql(query) + expectedResult = df.collect() + assert(collect(df.queryExecution.executedPlan) { + case agg: BaseAggregateExec => agg + }.size == numAggWithDisabled) + } + + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expectedResult) + assert(collect(df.queryExecution.executedPlan) { + case agg: BaseAggregateExec => agg + }.size == numAggWithEnabled) + } + } + + test("Test combine adjacent aggregation") { + withTempView("t") { + spark.range(20).selectExpr(s"id % 3 as k", "id % 7 as v") + .createOrReplaceTempView("t") + + // do not combine if no adjacent aggregation + checkNumAggregation( + "SELECT k, count(*) FROM t GROUP BY k", + numAggWithDisabled = 2, + numAggWithEnabled = 2) + + // combine adjacent hash aggregation + checkNumAggregation( + "SELECT k, count(*) FROM (SELECT /*+ repartition(k) */ * FROM t) GROUP BY k", + numAggWithDisabled = 2, + numAggWithEnabled = 1) + + // combine adjacent sort aggregate + checkNumAggregation( + "SELECT v, max(k) FROM (SELECT /*+ repartition(v) */ * FROM t) GROUP BY v", Review Comment: [P2] Exercise an actual `SortAggregateExec` here This query uses `max` over `Long` values, which Spark plans as two `HashAggregateExec` nodes. Since `checkNumAggregation` only counts the shared `BaseAggregateExec` type, this test passes without exercising the new `SortAggregateExec` rewrite branch. Please force a real sort aggregate and assert the concrete aggregate type before and after the rule so regressions in that branch are caught. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala: ########## @@ -0,0 +1,214 @@ +/* + * 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 + +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +class CombineAdjacentAggregationSuite extends QueryTest + with SharedSparkSession + with AdaptiveSparkPlanHelper { + + private def numAggregates(query: String): Int = { + val df = sql(query) + df.collect() + collect(df.queryExecution.executedPlan) { + case agg: BaseAggregateExec => agg + }.size + } + + private def checkNumAggregation( + query: String, + numAggWithDisabled: Int, + numAggWithEnabled: Int): Unit = { + var expectedResult: Array[Row] = null + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") { + val df = sql(query) + expectedResult = df.collect() + assert(collect(df.queryExecution.executedPlan) { + case agg: BaseAggregateExec => agg + }.size == numAggWithDisabled) + } + + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expectedResult) + assert(collect(df.queryExecution.executedPlan) { + case agg: BaseAggregateExec => agg + }.size == numAggWithEnabled) + } + } + + test("Test combine adjacent aggregation") { + withTempView("t") { + spark.range(20).selectExpr(s"id % 3 as k", "id % 7 as v") + .createOrReplaceTempView("t") + + // do not combine if no adjacent aggregation + checkNumAggregation( + "SELECT k, count(*) FROM t GROUP BY k", + numAggWithDisabled = 2, + numAggWithEnabled = 2) + + // combine adjacent hash aggregation + checkNumAggregation( + "SELECT k, count(*) FROM (SELECT /*+ repartition(k) */ * FROM t) GROUP BY k", + numAggWithDisabled = 2, + numAggWithEnabled = 1) + + // combine adjacent sort aggregate + checkNumAggregation( + "SELECT v, max(k) FROM (SELECT /*+ repartition(v) */ * FROM t) GROUP BY v", + numAggWithDisabled = 2, + numAggWithEnabled = 1) + + // combine adjacent object hash aggregate + checkNumAggregation( + "SELECT k, collect_set(v) FROM (SELECT /*+ repartition(k) */ * FROM t) GROUP BY k", + numAggWithDisabled = 2, + numAggWithEnabled = 1) + + // do not combine adjacent hash aggregation + checkNumAggregation( + "SELECT k, count(distinct v) FROM t GROUP BY k", + numAggWithDisabled = 4, + numAggWithEnabled = 4) + + // combine adjacent hash aggregation + checkNumAggregation( + "SELECT k, count(distinct v) FROM (SELECT /*+ repartition(k) */ * FROM t) GROUP BY k", + numAggWithDisabled = 4, + numAggWithEnabled = 2) + } + } + + test("Combined aggregate reads original input with a zero buffer offset") { + // When adjacent aggregates are combined into a single `Complete` mode aggregate, its child + // becomes the partial aggregate's child, so it reads the original input rather than a row of + // `[groupingKeys, aggregationBuffers]`. `initialInputBufferOffset` must therefore be reset to + // 0. Use multiple grouping keys and multiple aggregate functions so that a stale, non-zero + // offset would bind the aggregate functions against the wrong input columns. + withTempView("t") { + spark.range(60).selectExpr("id % 3 as k1", "id % 5 as k2", "id % 7 as v") + .createOrReplaceTempView("t") + + // hash aggregate with declarative aggregate functions + checkNumAggregation( + """SELECT k1, k2, sum(v), count(v), avg(v), max(v), min(v) + |FROM (SELECT /*+ repartition(k1, k2) */ * FROM t) GROUP BY k1, k2""".stripMargin, + numAggWithDisabled = 2, + numAggWithEnabled = 1) + + // object hash aggregate with imperative aggregate functions + checkNumAggregation( + """SELECT k1, k2, sort_array(collect_list(v)), count(v) + |FROM (SELECT /*+ repartition(k1, k2) */ * FROM t) GROUP BY k1, k2""".stripMargin, + numAggWithDisabled = 2, + numAggWithEnabled = 1) + } + } + + test("Do not combine when a shuffle sits between the partial and final aggregate") { + withTempView("t") { + spark.range(20).selectExpr("id % 3 as k", "id % 7 as v") + .createOrReplaceTempView("t") + + // The input is repartitioned by `k`, but the query groups by `k + 1`, so the partial + // aggregate's output is not partitioned the way the final aggregate requires. + // `EnsureRequirements` therefore inserts an Exchange between the two aggregates, leaving them + // non-adjacent, and the rule (which only matches a final aggregate whose child is the partial + // aggregate) must not combine them. + checkNumAggregation( + "SELECT k + 1, count(*) FROM (SELECT /*+ repartition(k) */ * FROM t) GROUP BY k + 1", Review Comment: [P2] Cover the grouping-expression guard with adjacent aggregates This case stops the rewrite because `EnsureRequirements` inserts an `Exchange`, so the partial and final aggregates are no longer adjacent. It therefore does not exercise the new canonicalized grouping-expression equality check in `isPartialAgg`. Please add a direct rule/physical-plan case where the two aggregates remain adjacent but use mismatched grouping expressions; otherwise a regression or removal of that safety check would remain invisible. -- 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]
