diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index c7535c3cebe51..e0a140bff0a32 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -1504,6 +1504,13 @@ object SQLConf { .booleanConf .createWithDefault(true) + val REPLACE_HASH_WITH_SORT_AGG_ENABLED = buildConf("spark.sql.execution.replaceHashWithSortAgg") + .internal() + .doc("Whether to replace hash aggregate node with sort aggregate based on children's ordering") + .version("3.3.0") + .booleanConf + .createWithDefault(false) + val STATE_STORE_PROVIDER_CLASS = buildConf("spark.sql.streaming.stateStore.providerClass") .internal() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index bb1b0ca3b645a..26c6904a896a5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -423,6 +423,9 @@ object QueryExecution { PlanSubqueries(sparkSession), RemoveRedundantProjects, EnsureRequirements(), + // `ReplaceHashWithSortAgg` needs to be added after `EnsureRequirements` to guarantee the + // sort order of each node is checked to be valid. + ReplaceHashWithSortAgg, // `RemoveRedundantSorts` needs to be added after `EnsureRequirements` to guarantee the same // number of partitions when instantiating PartitioningCollection. RemoveRedundantSorts, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAgg.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAgg.scala new file mode 100644 index 0000000000000..63ad2d0cafb75 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAgg.scala @@ -0,0 +1,106 @@ +/* + * 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.catalyst.expressions.SortOrder +import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Final, Partial} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.aggregate.HashAggregateExec +import org.apache.spark.sql.internal.SQLConf + +/** + * Replace [[HashAggregateExec]] with [[SortAggregateExec]] in the spark plan if: + * + * 1. The plan is a pair of partial and final [[HashAggregateExec]], and the child of partial + * aggregate satisfies the sort order of corresponding [[SortAggregateExec]]. + * or + * 2. The plan is a [[HashAggregateExec]], and the child satisfies the sort order of + * corresponding [[SortAggregateExec]]. + * + * Examples: + * 1. aggregate after join: + * + * HashAggregate(t1.i, SUM, final) + * | SortAggregate(t1.i, SUM, complete) + * HashAggregate(t1.i, SUM, partial) => | + * | SortMergeJoin(t1.i = t2.j) + * SortMergeJoin(t1.i = t2.j) + * + * 2. aggregate after sort: + * + * HashAggregate(t1.i, SUM, partial) SortAggregate(t1.i, SUM, partial) + * | => | + * Sort(t1.i) Sort(t1.i) + * + * [[HashAggregateExec]] can be replaced when its child satisfies the sort order of + * corresponding [[SortAggregateExec]]. [[SortAggregateExec]] is faster in the sense that + * it does not have hashing overhead of [[HashAggregateExec]]. + */ +object ReplaceHashWithSortAgg extends Rule[SparkPlan] { + def apply(plan: SparkPlan): SparkPlan = { + if (!conf.getConf(SQLConf.REPLACE_HASH_WITH_SORT_AGG_ENABLED)) { + plan + } else { + replaceHashAgg(plan) + } + } + + /** + * Replace [[HashAggregateExec]] with [[SortAggregateExec]]. + */ + private def replaceHashAgg(plan: SparkPlan): SparkPlan = { + plan.transformDown { + case hashAgg: HashAggregateExec if hashAgg.groupingExpressions.nonEmpty => + val sortAgg = hashAgg.toSortAggregate + hashAgg.child match { + case partialAgg: HashAggregateExec if isPartialAgg(partialAgg, hashAgg) => + if (SortOrder.orderingSatisfies( + partialAgg.child.outputOrdering, sortAgg.requiredChildOrdering.head)) { + sortAgg.copy( + aggregateExpressions = sortAgg.aggregateExpressions.map(_.copy(mode = Complete)), + child = partialAgg.child) + } else { + hashAgg + } + case other => + if (SortOrder.orderingSatisfies( + other.outputOrdering, sortAgg.requiredChildOrdering.head)) { + sortAgg + } else { + hashAgg + } + } + case other => other + } + } + + /** + * Check if `partialAgg` to be partial aggregate of `finalAgg`. + */ + private def isPartialAgg(partialAgg: HashAggregateExec, finalAgg: HashAggregateExec): Boolean = { + if (partialAgg.aggregateExpressions.forall(_.mode == Partial) && + finalAgg.aggregateExpressions.forall(_.mode == Final)) { + (finalAgg.logicalLink, partialAgg.logicalLink) match { + case (Some(agg1), Some(agg2)) => agg1.sameResult(agg2) + case _ => false + } + } else { + false + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index a4fe4529ed3d4..2b42804e784ed 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -116,6 +116,7 @@ case class AdaptiveSparkPlanExec( Seq( RemoveRedundantProjects, ensureRequirements, + ReplaceHashWithSortAgg, RemoveRedundantSorts, DisableUnnecessaryBucketedScan, OptimizeSkewedJoin(ensureRequirements) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala index 8545154028602..85e81cb12dca4 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala @@ -1153,6 +1153,15 @@ case class HashAggregateExec( } } + /** + * The corresponding [[SortAggregateExec]] to get same result as this node. + */ + def toSortAggregate: SortAggregateExec = { + SortAggregateExec( + requiredChildDistributionExpressions, groupingExpressions, aggregateExpressions, + aggregateAttributes, initialInputBufferOffset, resultExpressions, child) + } + override protected def withNewChildInternal(newChild: SparkPlan): HashAggregateExec = copy(child = newChild) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala new file mode 100644 index 0000000000000..78765fdf4f757 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala @@ -0,0 +1,131 @@ +/* + * 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.{DataFrame, QueryTest} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} +import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +abstract class ReplaceHashWithSortAggSuiteBase + extends QueryTest + with SharedSparkSession + with AdaptiveSparkPlanHelper { + + private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = { + val plan = df.queryExecution.executedPlan + assert(collectWithSubqueries(plan) { case s: HashAggregateExec => s }.length == hashAggCount) + assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount) + } + + private def checkAggs( + query: String, + enabledHashAggCount: Int, + enabledSortAggCount: Int, + disabledHashAggCount: Int, + disabledSortAggCount: Int): Unit = { + withSQLConf(SQLConf.REPLACE_HASH_WITH_SORT_AGG_ENABLED.key -> "true") { + val df = sql(query) + checkNumAggs(df, enabledHashAggCount, enabledSortAggCount) + val result = df.collect() + withSQLConf(SQLConf.REPLACE_HASH_WITH_SORT_AGG_ENABLED.key -> "false") { + val df = sql(query) + checkNumAggs(df, disabledHashAggCount, disabledSortAggCount) + checkAnswer(df, result) + } + } + } + + test("replace partial hash aggregate with sort aggregate") { + withTempView("t") { + spark.range(100).selectExpr("id as key").repartition(10).createOrReplaceTempView("t") + val query = + """ + |SELECT key, FIRST(key) + |FROM + |( + | SELECT key + | FROM t + | WHERE key > 10 + | SORT BY key + |) + |GROUP BY key + """.stripMargin + checkAggs(query, 1, 1, 2, 0) + } + } + + test("replace partial and final hash aggregate together with sort aggregate") { + withTempView("t1", "t2") { + spark.range(100).selectExpr("id as key").createOrReplaceTempView("t1") + spark.range(50).selectExpr("id as key").createOrReplaceTempView("t2") + val query = + """ + |SELECT key, COUNT(key) + |FROM + |( + | SELECT /*+ SHUFFLE_MERGE(t1) */ t1.key AS key + | FROM t1 + | JOIN t2 + | ON t1.key = t2.key + |) + |GROUP BY key + """.stripMargin + checkAggs(query, 0, 1, 2, 0) + } + } + + test("do not replace hash aggregate if child does not have sort order") { + withTempView("t1", "t2") { + spark.range(100).selectExpr("id as key").createOrReplaceTempView("t1") + spark.range(50).selectExpr("id as key").createOrReplaceTempView("t2") + val query = + """ + |SELECT key, COUNT(key) + |FROM + |( + | SELECT /*+ BROADCAST(t1) */ t1.key AS key + | FROM t1 + | JOIN t2 + | ON t1.key = t2.key + |) + |GROUP BY key + """.stripMargin + checkAggs(query, 2, 0, 2, 0) + } + } + + test("do not replace hash aggregate if there is no group-by column") { + withTempView("t1") { + spark.range(100).selectExpr("id as key").createOrReplaceTempView("t1") + val query = + """ + |SELECT COUNT(key) + |FROM t1 + """.stripMargin + checkAggs(query, 2, 0, 2, 0) + } + } +} + +class ReplaceHashWithSortAggSuite extends ReplaceHashWithSortAggSuiteBase + with DisableAdaptiveExecutionSuite + +class ReplaceHashWithSortAggSuiteAE extends ReplaceHashWithSortAggSuiteBase + with EnableAdaptiveExecutionSuite