ulysses-you commented on code in PR #57276: URL: https://github.com/apache/spark/pull/57276#discussion_r3600016095
########## 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: You are right, thanks for catching this. `max` over `Long` was planned as a `HashAggregateExec`, so the sort-aggregate branch was never exercised. Fixed in ada1570: the new `Combine adjacent sort aggregate` test uses `max` over a string column, which produces a non-mutable aggregation buffer and is not a `TypedImperativeAggregate`, so it is genuinely planned as `SortAggregateExec`. It now asserts the concrete aggregate type (and `Complete` mode) both before and after the rule. ########## 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: Agreed -- the existing case was short-circuited by the inserted Exchange and never reached the grouping-equality check. Fixed in ada1570 with a new `Do not combine adjacent aggregates with mismatched grouping expressions` test. It builds a genuinely adjacent partial/final pair (rule disabled so the raw pair survives), then drives `CombineAdjacentAggregation` directly: it first asserts the pair combines with matching grouping keys, then rewrites the grouping expressions to mismatch and asserts the rule refuses to combine. The tags (including `logicalLink`) are copied onto the rewritten node so `isPartialAgg` reaches the canonicalized grouping check rather than short-circuiting on a missing `logicalLink`. -- 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]
