This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new f88a554fafde [SPARK-58210][SQL][FOLLOWUP] Extend
CombineAdjacentAggregation to partial merge
f88a554fafde is described below
commit f88a554fafde5eec35c06af4c24e3e1df6cda313
Author: Wenchen Fan <[email protected]>
AuthorDate: Tue Aug 11 10:34:25 2026 +0800
[SPARK-58210][SQL][FOLLOWUP] Extend CombineAdjacentAggregation to partial
merge
### What changes were proposed in this pull request?
This is a follow-up to #57363. It extends `CombineAdjacentAggregation` to
combine adjacent
`PartialMerge` and `Final` hash aggregates into one `Final` aggregate. This
pattern can be
produced by physical-plan extensions that add an aggregation stage; Spark's
built-in query
planning does not currently produce it.
The combined aggregate preserves the upper `Final` aggregate's child
distribution requirement.
Because the lower aggregate is removed, the combined aggregate takes its
streaming and
shuffle-partition metadata from that lower aggregate and reads its original
child. Filtered
`Partial` aggregation remains supported, while a `PartialMerge` pair with
filters is not combined.
### Why are the changes needed?
`CombineAdjacentAggregation` currently handles only `Partial` followed by
`Final`. A compatible
adjacent `PartialMerge` and `Final` hash aggregate pair can also be
collapsed safely, avoiding an
unnecessary hash aggregation stage.
The upper aggregate's distribution requirement must remain after the rule
runs so that later AQE
optimization does not split partitions that the combined aggregate requires
to remain clustered.
The lower aggregate's streaming and shuffle-partition metadata remain
child-facing properties and
therefore move to the combined aggregate when the lower node is removed.
### Does this PR introduce _any_ user-facing change?
Yes. Spark may produce a single final hash aggregate instead of adjacent
partial-merge and final
hash aggregates when their grouping and logical lineage are compatible.
Query results are
unchanged.
### How was this patch tested?
Added coverage to `CombineAdjacentAggregationSuite` for:
- combining an executable `PartialMerge` and `Final` hash aggregate
pipeline and comparing the
uncombined and combined results;
- preserving the upper distribution requirement and lower streaming,
shuffle-partition,
grouping, child, and buffer-offset metadata;
- refusing to combine a filtered `PartialMerge` pair;
- retaining the distribution requirement under AQE and verifying that
shuffle partition specs
are present but are not partial reducer specs.
Ran:
```
build/sbt 'sql/testOnly
org.apache.spark.sql.execution.CombineAdjacentAggregationSuite' sql/scalastyle
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57859 from cloud-fan/cloud-fan/consolidate-adjacent-aggregation.
Authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../sql/execution/CombineAdjacentAggregation.scala | 64 +++++++++++++---
.../CombineAdjacentAggregationSuite.scala | 88 +++++++++++++++++++++-
2 files changed, 140 insertions(+), 12 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/CombineAdjacentAggregation.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/CombineAdjacentAggregation.scala
index 00643e1638ff..f9d841615d42 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/CombineAdjacentAggregation.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/CombineAdjacentAggregation.scala
@@ -17,13 +17,15 @@
package org.apache.spark.sql.execution
-import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Final,
Partial}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
Complete, Final, Partial, PartialMerge}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec,
HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.internal.SQLConf
/**
- * This rule combines adjacent aggregation with `Partial` and `Final` to
`Complete` mode.
+ * This rule combines adjacent aggregation with `Partial` and `Final` to
`Complete` mode. For
+ * [[HashAggregateExec]], it also combines `PartialMerge` and `Final` to
`Final` mode. The latter
+ * can be produced by physical plan extensions that add an extra aggregation
stage.
* Example for hash aggregate:
* HashAggregate (Final) HashAggregate (Complete)
* | |
@@ -43,19 +45,20 @@ import org.apache.spark.sql.internal.SQLConf
* It supports [[HashAggregateExec]], [[SortAggregateExec]] and
[[ObjectHashAggregateExec]].
*/
object CombineAdjacentAggregation extends Rule[SparkPlan] {
+ private case class CombinedAggregate(
+ aggregateExpressions: Seq[AggregateExpression],
+ initialInputBufferOffset: Int)
+
override def apply(plan: SparkPlan): SparkPlan = {
if (!conf.getConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED)) {
return plan
}
plan.transformDown {
- case finalAgg @ HashAggregateExec(_, _, _, _, _, _, _, _, partialAgg:
HashAggregateExec)
- if isPartialAgg(partialAgg, finalAgg) =>
- finalAgg.copy(
- groupingExpressions = partialAgg.groupingExpressions,
- aggregateExpressions =
partialAgg.aggregateExpressions.map(_.copy(mode = Complete)),
- initialInputBufferOffset = 0,
- child = partialAgg.child)
+ case finalAgg @ HashAggregateExec(_, _, _, _, _, _, _, _, partialAgg:
HashAggregateExec) =>
+ combinedAggregate(partialAgg, finalAgg)
+ .map(combineHashAggregates(partialAgg, finalAgg, _))
+ .getOrElse(finalAgg)
case finalAgg @ SortAggregateExec(_, _, _, _, _, _, _, _, partialAgg:
SortAggregateExec)
if isPartialAgg(partialAgg, finalAgg) =>
@@ -76,6 +79,41 @@ object CombineAdjacentAggregation extends Rule[SparkPlan] {
}
}
+ private def combineHashAggregates(
+ partialAgg: HashAggregateExec,
+ finalAgg: HashAggregateExec,
+ combined: CombinedAggregate): HashAggregateExec = {
+ // Keep the final aggregate's distribution requirement because the rule
runs after
+ // EnsureRequirements. The other child-facing metadata comes from the
removed aggregate.
+ finalAgg.copy(
+ isStreaming = partialAgg.isStreaming,
+ numShufflePartitions = partialAgg.numShufflePartitions,
+ groupingExpressions = partialAgg.groupingExpressions,
+ aggregateExpressions = combined.aggregateExpressions,
+ initialInputBufferOffset = combined.initialInputBufferOffset,
+ child = partialAgg.child)
+ }
+
+ private def combinedAggregate(
+ partialAgg: HashAggregateExec,
+ finalAgg: HashAggregateExec): Option[CombinedAggregate] = {
+ if (!isCompatibleAggregates(partialAgg, finalAgg)) {
+ None
+ } else if (partialAgg.aggregateExpressions.forall(_.mode == Partial)) {
+ Some(CombinedAggregate(
+ partialAgg.aggregateExpressions.map(_.copy(mode = Complete)),
+ initialInputBufferOffset = 0))
+ } else if (partialAgg.aggregateExpressions.forall(_.mode == PartialMerge)
&&
+ partialAgg.aggregateExpressions.forall(_.filter.isEmpty) &&
+ finalAgg.aggregateExpressions.forall(_.filter.isEmpty)) {
+ Some(CombinedAggregate(
+ finalAgg.aggregateExpressions,
+ partialAgg.initialInputBufferOffset))
+ } else {
+ None
+ }
+ }
+
/**
* Check if `partialAgg` is the partial aggregate of `finalAgg`.
*/
@@ -83,7 +121,13 @@ object CombineAdjacentAggregation extends Rule[SparkPlan] {
partialAgg: BaseAggregateExec,
finalAgg: BaseAggregateExec): Boolean = {
partialAgg.aggregateExpressions.forall(_.mode == Partial) &&
- finalAgg.aggregateExpressions.forall(_.mode == Final) &&
+ isCompatibleAggregates(partialAgg, finalAgg)
+ }
+
+ private def isCompatibleAggregates(
+ partialAgg: BaseAggregateExec,
+ finalAgg: BaseAggregateExec): Boolean = {
+ finalAgg.aggregateExpressions.forall(_.mode == Final) &&
partialAgg.groupingExpressions.map(_.canonicalized) ==
finalAgg.groupingExpressions.map(_.canonicalized) &&
finalAgg.logicalLink.isDefined &&
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala
index 17675705060d..bd970f1c34d8 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala
@@ -18,8 +18,9 @@
package org.apache.spark.sql.execution
import org.apache.spark.sql.{QueryTest, Row}
-import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Final}
-import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Final,
PartialMerge}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper,
AQEShuffleReadExec}
import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec,
HashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
@@ -191,6 +192,89 @@ class CombineAdjacentAggregationSuite extends QueryTest
}
}
+ test("Combine adjacent partial merge and final hash aggregates") {
+ withTempView("t") {
+ spark.range(20).selectExpr("id % 3 as k", "id % 7 as
v").createOrReplaceTempView("t")
+ val query = "SELECT k, count(*) FROM (SELECT /*+ repartition(k) */ *
FROM t) GROUP BY k"
+ val finalAgg = withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") {
+ collect(sql(query).queryExecution.executedPlan) {
+ case agg: HashAggregateExec if
agg.child.isInstanceOf[HashAggregateExec] => agg
+ }.head
+ }
+ val partialAgg = finalAgg.child.asInstanceOf[HashAggregateExec]
+ val partialMergeInput = InputAdapter(partialAgg)
+ val partialMergeAgg = partialAgg.copy(
+ requiredChildDistributionExpressions =
Some(partialAgg.groupingExpressions),
+ isStreaming = true,
+ numShufflePartitions = Some(5),
+ initialInputBufferOffset = finalAgg.initialInputBufferOffset,
+ aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode
= PartialMerge)),
+ child = partialMergeInput)
+ partialMergeAgg.copyTagsFrom(partialAgg)
+ val finalOverPartialMerge = finalAgg.copy(child = partialMergeAgg)
+ finalOverPartialMerge.copyTagsFrom(finalAgg)
+
+ val combined = CombineAdjacentAggregation(finalOverPartialMerge)
+ .asInstanceOf[HashAggregateExec]
+ assert(combined.aggregateExpressions.forall(_.mode == Final))
+ assert(combined.requiredChildDistributionExpressions ==
+ finalAgg.requiredChildDistributionExpressions)
+ assert(combined.isStreaming == partialMergeAgg.isStreaming)
+ assert(combined.numShufflePartitions ==
partialMergeAgg.numShufflePartitions)
+ assert(combined.groupingExpressions ==
partialMergeAgg.groupingExpressions)
+ assert(combined.aggregateAttributes == finalAgg.aggregateAttributes)
+ assert(combined.resultExpressions == finalAgg.resultExpressions)
+ assert(combined.initialInputBufferOffset ==
partialMergeAgg.initialInputBufferOffset)
+ assert(combined.child == partialMergeInput)
+ assert(finalOverPartialMerge.executeCollect().map(_.copy()).toSet ==
+ combined.executeCollect().map(_.copy()).toSet)
+
+ val filteredPartialMerge = partialMergeAgg.copy(
+ aggregateExpressions =
partialMergeAgg.aggregateExpressions.zipWithIndex.map {
+ case (agg, 0) => agg.copy(filter = Some(Literal.TrueLiteral))
+ case (agg, _) => agg
+ })
+ filteredPartialMerge.copyTagsFrom(partialAgg)
+ val finalOverFilteredPartialMerge = finalAgg.copy(child =
filteredPartialMerge)
+ finalOverFilteredPartialMerge.copyTagsFrom(finalAgg)
+
assert(collect(CombineAdjacentAggregation(finalOverFilteredPartialMerge)) {
+ case agg: HashAggregateExec => agg
+ }.size == 3)
+ }
+ }
+
+ test("Combined aggregate retains its distribution requirement under AQE") {
+ import testImplicits._
+
+ withTempView("t") {
+ spark.sparkContext.parallelize(
+ (1 to 10).map(i => (if (i > 4) 5 else i, i.toString)), 3)
+ .toDF("k", "v").createOrReplaceTempView("t")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key
-> "true",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "5",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "150") {
+ val df = sql(
+ "SELECT k, count(*) FROM (SELECT /*+ REBALANCE(k) */ * FROM t) GROUP
BY k")
+ checkAnswer(df, Seq(Row(1, 1), Row(2, 1), Row(3, 1), Row(4, 1), Row(5,
6)))
+ assert(collect(df.queryExecution.executedPlan) {
+ case agg: HashAggregateExec => agg
+ }.size == 1)
+ val partitionSpecs = collect(df.queryExecution.executedPlan) {
+ case read: AQEShuffleReadExec => read
+ }.flatMap(_.partitionSpecs)
+ assert(partitionSpecs.nonEmpty)
+
assert(partitionSpecs.forall(!_.isInstanceOf[PartialReducerPartitionSpec]))
+ }
+ }
+ }
+
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")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
