From c4e786f08de7fe7b32467c7bdd0b7355da4c2faa Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Thu, 17 Sep 2026 17:28:40 +0800 Subject: [PATCH 1/2] [SPARK-59564][SQL] Combine adjacent aggregation across a GroupPartitionsExec ### What changes were proposed in this pull request? `CombineAdjacentAggregation` now looks through the `GroupPartitionsExec` and the local sorts `EnsureRequirements` may have put between the partial and the final aggregate, so the pair is combined even when the final aggregate's distribution was satisfied without a shuffle. The grouping is re-parented onto the aggregate's child with the key positions it projects moved into that child's key space, which is what `GroupPartitionsExec.withKeyPositionsFor` answers. It was planned against the aggregate, whose partitioning is the child's projected down to what the aggregate's output keeps, so those positions name a key space the child being handed need not share. A sort crossed above the aggregate orders the rows the combined aggregate reads, by the grouping the two aggregates share, so the sort feeding the partial aggregate goes with it. Where the partial aggregate holds no sort of its own, the pair is left alone instead, being then the only cardinality reducer before that sort. With no sort crossed at all, the sort below the aggregate stays below it: the aggregate reads what it read, and the ordering it claims is the one it had. ### Why are the changes needed? The rule only matched a strictly adjacent pair, so it missed the shape where `EnsureRequirements` satisfies the final aggregate's `ClusteredDistribution` with a `GroupPartitionsExec` rather than a shuffle. Folding the pair there removes an aggregation pass, and where a sort fed the partial aggregate, that sort as well. ### Does this PR introduce _any_ user-facing change? Yes, the plan changes for an aggregate pair whose child needs a `GroupPartitionsExec`, e.g. for `GROUP BY id, name` over a v2 table partitioned by `(id, name)` whose keys repeat across splits: Before: ``` HashAggregate (Final) +- GroupPartitions +- HashAggregate (Partial) +- BatchScan ``` After: ``` HashAggregate (Complete) +- GroupPartitions +- BatchScan ``` A grouping on part of the partition keys is covered too: the positions are moved onto the scan's key space, which is a lookup, the projected expressions being the child's own. Where a sort sits between the pair it goes with the partial aggregate, so two sorts become one. Where the partial aggregate holds no sort of its own and a sort above orders the combined aggregate's rows, the pair is left alone and the plan is unchanged. That last path is taken by more shapes than a source that declares its ordering: with `spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled` the ordering derived from the partition keys already satisfies the partial aggregate's, so it holds no sort of its own to give up. The bail is load-bearing for a grouping that derives its ordering, not only for one whose source declares it. Results are unchanged. ### How was this patch tested? New tests in `KeyGroupedPartitioningSuite`: the pair combined across the grouping, including the object-hash pair and the AQE path; the sort pair with the sort in between; a grouping on part of the partition keys combined with its positions moved; a sort pair whose grouping covers part of the partition keys; a grouping the scan reports narrowed itself; the pair kept where the source already orders the aggregate's input; and a comparison against the plan `bypassPartialAggregation` builds for the same query, whose grouping projects what the fold's does. Each asserts the aggregates left and their mode, the grouping's positions where they matter, and `ValidateRequirements.validate(plan)`, and compares the answer against the rule-off run. The two sort tests' rule-off arms pin the two sorts the fold takes down to one. Verified the tests discriminate by mutation: keeping the positions unmoved fails both narrowed-grouping tests with three rows instead of two, and dropping the sort bail fails the ordered-source test with one aggregate instead of two. Ran `sql/Test/compile`, `KeyGroupedPartitioningSuite` (191), and `AdaptivePartialAggregationSuite`, `CombineAdjacentAggregationSuite`, `DataFrameAggregateSuite`, `EnsureRequirementsSuite`, `GroupPartitionsExecSuite`, `PushDownLocalSortSuite`, `RemoveRedundantSortsSuite`, `RemoveRedundantWindowGroupLimitsSuite`, `ReplaceHashWithSortAggSuite` and `SQLMetricsSuite` (603 in total, no failures). ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (deepseek-flash) Assisted-by: Claude Code (deepseek-flash) --- .../CombineAdjacentAggregation.scala | 104 ++++- .../datasources/v2/GroupPartitionsExec.scala | 37 ++ .../KeyGroupedPartitioningSuite.scala | 424 +++++++++++++++++- 3 files changed, 533 insertions(+), 32 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 f9d841615d427..31c308f44192f 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 @@ -20,6 +20,7 @@ package org.apache.spark.sql.execution 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.execution.datasources.v2.GroupPartitionsExec import org.apache.spark.sql.internal.SQLConf /** @@ -43,6 +44,10 @@ import org.apache.spark.sql.internal.SQLConf * Exchange * * It supports [[HashAggregateExec]], [[SortAggregateExec]] and [[ObjectHashAggregateExec]]. + * + * A [[GroupPartitionsExec]] and the local sorts `EnsureRequirements` put between the two aggregates + * are looked through, so the pair is combined even when the final aggregate's distribution was + * satisfied without a shuffle. See `detachAggregate`. */ object CombineAdjacentAggregation extends Rule[SparkPlan] { private case class CombinedAggregate( @@ -55,34 +60,91 @@ object CombineAdjacentAggregation extends Rule[SparkPlan] { } plan.transformDown { - case finalAgg @ HashAggregateExec(_, _, _, _, _, _, _, _, partialAgg: HashAggregateExec) => - combinedAggregate(partialAgg, finalAgg) - .map(combineHashAggregates(partialAgg, finalAgg, _)) - .getOrElse(finalAgg) + case finalAgg: HashAggregateExec => + detachAggregate(finalAgg.child, hasUpperSort = false) match { + case Some((partialAgg: HashAggregateExec, child)) => + combinedAggregate(partialAgg, finalAgg) + .map(combineHashAggregates(partialAgg, finalAgg, _, child)) + .getOrElse(finalAgg) + case _ => finalAgg + } - case finalAgg @ SortAggregateExec(_, _, _, _, _, _, _, _, partialAgg: SortAggregateExec) - if isPartialAgg(partialAgg, finalAgg) => - finalAgg.copy( - groupingExpressions = partialAgg.groupingExpressions, - aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = Complete)), - initialInputBufferOffset = 0, - child = partialAgg.child) + case finalAgg: SortAggregateExec => + detachAggregate(finalAgg.child, hasUpperSort = false) match { + case Some((partialAgg: SortAggregateExec, child)) if isPartialAgg(partialAgg, finalAgg) => + finalAgg.copy( + groupingExpressions = partialAgg.groupingExpressions, + aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = Complete)), + initialInputBufferOffset = 0, + child = child) + case _ => finalAgg + } - case finalAgg @ ObjectHashAggregateExec(_, _, _, _, _, _, _, _, - partialAgg: ObjectHashAggregateExec) - if isPartialAgg(partialAgg, finalAgg) => - finalAgg.copy( - groupingExpressions = partialAgg.groupingExpressions, - aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = Complete)), - initialInputBufferOffset = 0, - child = partialAgg.child) + case finalAgg: ObjectHashAggregateExec => + detachAggregate(finalAgg.child, hasUpperSort = false) match { + case Some((partialAgg: ObjectHashAggregateExec, child)) + if isPartialAgg(partialAgg, finalAgg) => + finalAgg.copy( + groupingExpressions = partialAgg.groupingExpressions, + aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = Complete)), + initialInputBufferOffset = 0, + child = child) + case _ => finalAgg + } } } + /** + * Detaches the aggregate at the bottom of the chain `plan` starts and hands it back together with + * the subtree to leave where it was, or `None` when the chain bottoms out at no aggregate, or at + * one that cannot leave. The chain's `GroupPartitionsExec` and local sorts are crossed in place, + * so the combined aggregate reads whatever ends up on top of them. + * + * A sort crossed above the aggregate orders the rows the combined aggregate reads, by the + * grouping the two aggregates share, so the sort the aggregate reads goes with it: that crossed + * sort is what orders those rows. Where the aggregate holds no sort of its own, it stays, being + * the only cardinality reducer before that sort. With no sort crossed at all, the sort below the + * aggregate stays below it: the aggregate reads what it read, and the ordering it claims is the + * one it had. + * + * @param hasUpperSort whether a local sort has been crossed above `plan`, which is what makes the + * sort below the aggregate dead. + */ + private def detachAggregate( + plan: SparkPlan, + hasUpperSort: Boolean): Option[(BaseAggregateExec, SparkPlan)] = plan match { + case aggregate: BaseAggregateExec => + if (!hasUpperSort) { + Some((aggregate, aggregate.child)) + } else { + aggregate.child match { + case sort: SortExec if !sort.global => Some((aggregate, sort.child)) + case _ => None + } + } + + case group: GroupPartitionsExec => + detachAggregate(group.child, hasUpperSort) match { + case Some((aggregate, child)) => + group.withKeyPositionsFor(child).map(regrouped => (aggregate, regrouped)) + case _ => None + } + + case sort: SortExec if !sort.global => + detachAggregate(sort.child, hasUpperSort = true) match { + case Some((aggregate, child)) => + Some((aggregate, sort.withNewChildren(Seq(child)))) + case _ => None + } + + case _ => None + } + private def combineHashAggregates( partialAgg: HashAggregateExec, finalAgg: HashAggregateExec, - combined: CombinedAggregate): HashAggregateExec = { + combined: CombinedAggregate, + child: SparkPlan): 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( @@ -91,7 +153,7 @@ object CombineAdjacentAggregation extends Rule[SparkPlan] { groupingExpressions = partialAgg.groupingExpressions, aggregateExpressions = combined.aggregateExpressions, initialInputBufferOffset = combined.initialInputBufferOffset, - child = partialAgg.child) + child = child) } private def combinedAggregate( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala index b32e842b35bf5..063eb71268b11 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala @@ -405,6 +405,43 @@ case class GroupPartitionsExec( } } + /** + * This node reading `newChild`, with the key positions it projects moved to where the expressions + * they name sit there, or `None` when it may not be re-parented onto it. It is named for what it + * answers rather than for the node it returns, which is the node it is called on. + * + * `EnsureRequirements` computed `joinKeyPositions` against the child this node was planned for, + * whose partitioning is that child's projected down to the positions the operator above keeps, so + * the positions name a key space `newChild` need not share. The projected expressions are the + * planned child's own (`KeyedPartitioning.project` builds them that way), which makes moving them + * a lookup; a child that does not hold one is turned away. + * + * Moving them is all this does: what the operator above reads, the ordering it was planned + * against included, is the caller's to hold, since nothing here knows what that operator + * requires. + */ + def withKeyPositionsFor(newChild: SparkPlan): Option[GroupPartitionsExec] = { + // The member of each child's partitioning this reads has to be the one `grouping` reads, since + // the positions are only meaningful for that member. `representativeOf` answers as the lookup + // there does: the first keyed member, nested collections included. + val childKeyed = PartitioningCollection.representativeOf(child.outputPartitioning) + val newChildKeyed = PartitioningCollection.representativeOf(newChild.outputPartitioning) + (childKeyed, newChildKeyed) match { + case (Some(childKp), Some(newChildKp)) => + val plannedInNewChild = childKp.expressions.map(newChildKp.expressions.indexOf) + if (plannedInNewChild.exists(_ < 0)) { + return None + } + val positions = joinKeyPositions.fold(plannedInNewChild)(_.map(plannedInNewChild)) + val regrouped = copy( + child = newChild, + joinKeyPositions = Option.when(positions != newChildKp.expressions.indices)(positions)) + regrouped.copyTagsFrom(this) + Some(regrouped) + case _ => None + } + } + override def simpleString(maxFields: Int): String = { s"$nodeName${planSummaryParts(maxFields).map(" " + _).mkString("")}" } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index c49b3c3e92104..98ef135392bc9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -24,6 +24,7 @@ import org.apache.spark.rdd.SortedMergeCoalescedRDD import org.apache.spark.sql.{DataFrame, ExplainSuiteHelper, Row} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, ExprId, Literal, TransformExpression} +import org.apache.spark.sql.catalyst.expressions.aggregate.Complete import org.apache.spark.sql.catalyst.plans.{Cross, ExistenceJoin, Inner, JoinType, LeftAnti, LeftSemi, LeftSingle} import org.apache.spark.sql.catalyst.plans.physical import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning @@ -44,6 +45,7 @@ import org.apache.spark.sql.execution.{ SparkPlan, UnionExec, WholeStageCodegenExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, GroupPartitionsExec} import org.apache.spark.sql.execution.exchange.{EnsureRequirements, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike, ValidateRequirements} import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, ShuffledJoin, SortMergeJoinExec} @@ -883,6 +885,14 @@ class KeyGroupedPartitioningSuite } } + /** The first node under `plan` that is not a codegen wrapper or the scan's projection. */ + protected def unwrapWrappers(plan: SparkPlan): SparkPlan = plan match { + case w: WholeStageCodegenExec => unwrapWrappers(w.child) + case i: InputAdapter => unwrapWrappers(i.child) + case p: ProjectExec => unwrapWrappers(p.child) + case other => other + } + /** Every `KeyedPartitioning` these nodes report, flattening partitioning collections. */ protected def keyedPartitioningsOf( nodes: Seq[SparkPlan]): Seq[physical.KeyedPartitioning] = { @@ -6391,13 +6401,6 @@ class KeyGroupedPartitioningSuite // A stage boundary below the sort shows up as codegen wrappers around the grouping, and the // scan's projection sits between a limit node and the scan it reads. - def unwrap(plan: SparkPlan): SparkPlan = plan match { - case w: WholeStageCodegenExec => unwrap(w.child) - case i: InputAdapter => unwrap(i.child) - case p: ProjectExec => unwrap(p.child) - case other => other - } - // One limit node is left, with one local sort in the whole plan and no global one. def assertFinalAloneWithOneLocalSort(plan: SparkPlan): Unit = { val limits = limitModes(plan) @@ -6436,7 +6439,7 @@ class KeyGroupedPartitioningSuite assertFinalAloneWithOneLocalSort(byKeyPlan) finalChild(byKeyPlan) match { case g: GroupPartitionsExec => - assert(unwrap(g.child).isInstanceOf[SortExec], + assert(unwrapWrappers(g.child).isInstanceOf[SortExec], s"expected the sort that fed the partial node below the grouping:\n$byKeyPlan") case other => fail(s"expected the final limit to read the grouping, got $other:\n$byKeyPlan") @@ -6449,7 +6452,7 @@ class KeyGroupedPartitioningSuite checkAnswer(byPrice, Seq(Row(1L, "aa", 10.0f), Row(2L, "cc", 30.0f))) val byPricePlan = byPrice.queryExecution.executedPlan assertFinalAloneWithOneLocalSort(byPricePlan) - assert(unwrap(finalChild(byPricePlan)).isInstanceOf[SortExec], + assert(unwrapWrappers(finalChild(byPricePlan)).isInstanceOf[SortExec], s"expected the sort above the grouping to be the one left:\n$byPricePlan") // A source that reports the ordering the window needs leaves the partial node without a sort @@ -6463,14 +6466,413 @@ class KeyGroupedPartitioningSuite assert(limitModes(reportedPlan) == Seq(Final, Partial), s"expected both limit nodes, got ${limitModes(reportedPlan)}:\n$reportedPlan") assert(collectFirst(reportedPlan) { - case w: WindowGroupLimitExec if w.mode == Partial => unwrap(w.child) + case w: WindowGroupLimitExec if w.mode == Partial => unwrapWrappers(w.child) }.exists(_.isInstanceOf[BatchScanExec]), s"expected the partial node to read the scan, with no sort of its own:\n$reportedPlan") - assert(unwrap(finalChild(reportedPlan)).isInstanceOf[SortExec], + assert(unwrapWrappers(finalChild(reportedPlan)).isInstanceOf[SortExec], s"expected the sort above the grouping to be the one left:\n$reportedPlan") } } + test("SPARK-59564: combine adjacent aggregates across a GroupPartitionsExec") { + // (1, 'aa') is stored in two splits, so the table's reported KeyedPartitioning is not grouped + // and EnsureRequirements coalesces the two splits with a GroupPartitionsExec to satisfy the + // final aggregate's clustered distribution. The partial and final aggregates are therefore not + // adjacent, and the rule has to look through the grouping to reach the pair. + val partitions = Array(identity("id"), identity("name")) + createTable(items, itemsColumns, partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " + + "(1, 'aa', 20.0, cast('2020-01-01' as timestamp)), " + + "(2, 'bb', 30.0, cast('2020-01-01' as timestamp))") + + // Grouping on the partition keys themselves keeps the partial aggregate from projecting any of + // them away, which is what lets the grouping be re-parented onto the scan it was reading. + val query = s"SELECT id, name, count(*) FROM testcat.ns.$items GROUP BY id, name" + val expected = Seq(Row(1L, "aa", 2L), Row(2L, "bb", 1L)) + + def aggregates(plan: SparkPlan): Seq[BaseAggregateExec] = + collect(plan) { case agg: BaseAggregateExec => agg } + + // The same pair planned as object-hash aggregates, whose answer is order-insensitive so the two + // plans can be compared. + val objectHashQuery = + s"SELECT id, name, sort_array(collect_set(price)) FROM testcat.ns.$items GROUP BY id, name" + + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val objectHashExpected = withSQLConf( + SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") { + val plan = sql(query).queryExecution.executedPlan + val aggs = aggregates(plan) + assert(aggs.size == 2, s"expected the pair to be planned separately:\n$plan") + assert(collectAllGroupPartitions(plan).nonEmpty, + s"the grouping is what makes the pair non-adjacent:\n$plan") + sql(objectHashQuery).collect() + } + + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + val aggs = aggregates(plan) + assert(aggs.size == 1, s"expected one combined aggregate, got ${aggs.size}:\n$plan") + assert(aggs.head.aggregateExpressions.forall(_.mode == Complete), + s"expected the combined aggregate to be complete:\n$plan") + // The grouping stays, and reads the scan the partial aggregate was reading: the combine + // only drops the partial aggregate, which was the grouping's child. + val grouping = collectAllGroupPartitions(plan) + assert(grouping.size == 1, s"expected the grouping to stay, got ${grouping.size}:\n$plan") + assert(unwrapWrappers(grouping.head.child).isInstanceOf[BatchScanExec], + s"expected the grouping to read the scan:\n$plan") + assert(ValidateRequirements.validate(plan), s"the combined plan has to hold up:\n$plan") + + val objectHash = sql(objectHashQuery) + checkAnswer(objectHash, objectHashExpected) + val objectHashPlan = objectHash.queryExecution.executedPlan + val objectHashAggs = aggregates(objectHashPlan) + assert(objectHashAggs.size == 1 && + objectHashAggs.head.isInstanceOf[ObjectHashAggregateExec] && + objectHashAggs.head.aggregateExpressions.forall(_.mode == Complete), + s"expected one combined object hash aggregate in complete mode:\n$objectHashPlan") + assert(unwrapWrappers(collectAllGroupPartitions(objectHashPlan).head.child) + .isInstanceOf[BatchScanExec], + s"expected the grouping to read the scan:\n$objectHashPlan") + } + + // AQE runs the rule from its stage-preparation rules, on a plan whose grouping was inserted + // by the stage-preparation `EnsureRequirements` rather than by the initial planning pass, and + // re-runs `EnsureRequirements` over the folded plan. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + val aggs = aggregates(plan) + assert(aggs.size == 1, s"expected one combined aggregate, got ${aggs.size}:\n$plan") + assert(aggs.head.aggregateExpressions.forall(_.mode == Complete), + s"expected the combined aggregate to be complete:\n$plan") + val grouping = collectAllGroupPartitions(plan) + assert(grouping.size == 1, s"expected the grouping to stay, got ${grouping.size}:\n$plan") + assert(unwrapWrappers(grouping.head.child).isInstanceOf[BatchScanExec], + s"expected the grouping to read the scan:\n$plan") + assert(ValidateRequirements.validate(plan), s"the combined plan has to hold up:\n$plan") + } + } + } + + test("SPARK-59564: combine across a grouping the partial aggregate narrowed") { + // (1, 'aa') and (2, 'aa') are two splits sharing `name`. Grouping by `name` alone makes the + // partial aggregate project `id` away, collapsing the scan's KP([id, name]) to + // KP([name], isCollapsed = true), and the grouping coalesces the two splits on that narrowed + // key. That key sits at position 0 of the narrowing but at position 1 of the scan's, so + // re-parenting the grouping has to translate the positions it projects: keeping them would have + // it group by `id`, where the two 'aa' rows land in different partitions and the final + // aggregate returns a row per partition, a wrong answer rather than a slower plan. + // + // `max(id)` is what makes the shape: it keeps `id` in the scan's output, so the scan reports + // the full KP([id, name]) and the narrowing happens at the partial aggregate, the node the fold + // takes away. The narrowed key being collapsed is what `allowKeysSubsetOfPartitionKeys` is + // needed for, the same way the SPARK-46367 test does. + val partitions = Array(identity("id"), identity("name")) + createTable(items, itemsColumns, partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " + + "(2, 'aa', 20.0, cast('2020-01-01' as timestamp)), " + + "(3, 'cc', 30.0, cast('2020-01-01' as timestamp))") + + val query = s"SELECT name, max(id), count(*) FROM testcat.ns.$items GROUP BY name" + val expected = Seq(Row("aa", 2L, 2L), Row("cc", 3L, 1L)) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") { + val df = sql(query) + checkAnswer(df, expected) + assert(collectAllGroupPartitions(df.queryExecution.executedPlan).nonEmpty, + s"expected the grouping the pair is separated by:\n${df.queryExecution.executedPlan}") + } + + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 1, s"expected one combined aggregate, got ${aggs.size}:\n$plan") + assert(aggs.head.aggregateExpressions.forall(_.mode == Complete), + s"expected the combined aggregate to be complete:\n$plan") + // The grouping stays, reading the scan with `name` translated to the position it holds + // there, which is what keeps the coalescing the same one it did before. + val grouping = collectAllGroupPartitions(plan) + assert(grouping.size == 1, s"expected the grouping to stay, got ${grouping.size}:\n$plan") + assert(grouping.head.joinKeyPositions == Some(Seq(1)), + s"expected `name` translated to position 1 of the scan:\n$plan") + assert(unwrapWrappers(grouping.head.child).isInstanceOf[BatchScanExec], + s"expected the grouping to read the scan:\n$plan") + assert(ValidateRequirements.validate(plan), s"the combined plan has to hold up:\n$plan") + } + } + } + + test("SPARK-59564: combine where the scan reports the narrowed partitioning itself") { + // Nothing references `id`, so the pruning takes it out of the scan's output and the scan + // reports the keyed partitioning projected down to `name` on its own: the scan narrows, rather + // than the partial aggregate under it. The grouping was therefore planned against that same + // space, and re-parenting it changes no position, unlike the case above. + val partitions = Array(identity("id"), identity("name")) + createTable(items, itemsColumns, partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " + + "(2, 'aa', 20.0, cast('2020-01-01' as timestamp)), " + + "(3, 'cc', 30.0, cast('2020-01-01' as timestamp))") + + val query = s"SELECT name, count(*) FROM testcat.ns.$items GROUP BY name" + val expected = Seq(Row("aa", 2L), Row("cc", 1L)) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") { + val plan = sql(query).queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 2, s"expected the pair to be planned separately:\n$plan") + // The scan reports `name` alone, and its keyed partitioning with it. + val scan = collect(plan) { case s: BatchScanExec => s }.head + assert(scan.output.map(_.name) == Seq("name"), + s"expected `id` to be pruned out of the scan:\n$plan") + scan.outputPartitioning match { + case kp: KeyedPartitioning => + assert(kp.expressions == scan.output, + s"expected the scan to report the narrowed keyed partitioning:\n$plan") + case other => fail(s"expected a keyed partitioning, got $other:\n$plan") + } + assert(collectAllGroupPartitions(plan).nonEmpty, + s"expected the grouping the pair is separated by:\n$plan") + } + + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 1, s"expected one combined aggregate, got ${aggs.size}:\n$plan") + assert(aggs.head.aggregateExpressions.forall(_.mode == Complete), + s"expected the combined aggregate to be complete:\n$plan") + val grouping = collectAllGroupPartitions(plan) + assert(grouping.size == 1, s"expected the grouping to stay, got ${grouping.size}:\n$plan") + // Nothing to translate: the positions index the space the scan reports already. + assert(grouping.head.joinKeyPositions.isEmpty, + s"expected the grouping to keep the positions it was planned with:\n$plan") + assert(unwrapWrappers(grouping.head.child).isInstanceOf[BatchScanExec], + s"expected the grouping to read the scan:\n$plan") + assert(ValidateRequirements.validate(plan), s"the combined plan has to hold up:\n$plan") + } + } + } + + test("SPARK-59564: the fold projects what the planner projects without a partial aggregate") { + // With `bypassPartialAggregation`, the planner runs one `Complete` aggregate and has + // `EnsureRequirements` satisfy its distribution, which for a grouping on part of the partition + // keys gives a `GroupPartitionsExec` over the scan. That plan is what the fold has to reproduce + // when the partial aggregation is planned and then taken away, so the positions come from the + // planner rather than from this rule's own arithmetic. + val partitions = Array(identity("id"), identity("name")) + createTable(items, itemsColumns, partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " + + "(2, 'aa', 20.0, cast('2020-01-01' as timestamp)), " + + "(3, 'cc', 30.0, cast('2020-01-01' as timestamp))") + + val query = s"SELECT name, max(id), count(*) FROM testcat.ns.$items GROUP BY name" + val expected = Seq(Row("aa", 2L, 2L), Row("cc", 3L, 1L)) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + def projections(plan: SparkPlan): Seq[Option[Seq[Int]]] = + collectAllGroupPartitions(plan).map(_.joinKeyPositions) + + // Plan and collect inside each block: `executedPlan` is lazy, so asking for it outside would + // plan under the default config and the toggle would do nothing. + val planned = withSQLConf(SQLConf.BYPASS_PARTIAL_AGGREGATION.key -> "true") { + val df = sql(query) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 1 && aggs.head.aggregateExpressions.forall(_.mode == Complete), + s"expected the planner to run one complete aggregate:\n$plan") + assert(projections(plan).nonEmpty, + s"expected the planner to satisfy the distribution with a grouping:\n$plan") + projections(plan) + } + + val folded = withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + // The combined answer is asserted by the narrowed-grouping test, which runs the same query. + // This one is about the plan the fold leaves, so that a wrong projection fails here rather + // than being caught by the answer. + val plan = sql(query).queryExecution.executedPlan + assert(collect(plan) { case agg: BaseAggregateExec => agg }.size == 1, + s"expected one combined aggregate:\n$plan") + projections(plan) + } + + assert(folded == planned, + s"expected the fold to project what the planner does without a partial: $planned") + } + } + + test("SPARK-59564: combine adjacent sort aggregates across a GroupPartitionsExec") { + // `max` over a string column cannot be planned as a hash aggregate, so the pair is a pair of + // SortAggregateExecs and each of them needs its input ordered by the grouping keys. The sort + // below the partial aggregate only gave the partial aggregate that ordering, and the sort the + // grouping forces above it orders the rows the combined aggregate reads by the same keys, so + // the sort below goes with the partial aggregate. + val partitions = Array(identity("id"), identity("name")) + createTable(items, itemsColumns, partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " + + "(1, 'aa', 20.0, cast('2020-01-01' as timestamp)), " + + "(2, 'bb', 30.0, cast('2020-01-01' as timestamp))") + + val query = + s"SELECT id, name, max(cast(price as string)) FROM testcat.ns.$items GROUP BY id, name" + val expected = Seq(Row(1L, "aa", "20.0"), Row(2L, "bb", "30.0")) + + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") { + val plan = sql(query).queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 2 && aggs.forall(_.isInstanceOf[SortAggregateExec]), + s"expected a pair of sort aggregates:\n$plan") + assert(collectAllGroupPartitions(plan).nonEmpty, + s"the grouping is what makes the pair non-adjacent:\n$plan") + // Two sorts, one for the partial aggregate and one the final one reads. + assert(collect(plan) { case sort: SortExec => sort }.size == 2, + s"expected two sorts feeding the pair:\n$plan") + } + + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 1, s"expected one combined aggregate, got ${aggs.size}:\n$plan") + assert(aggs.head.aggregateExpressions.forall(_.mode == Complete), + s"expected the combined aggregate to be complete:\n$plan") + val sorts = collect(plan) { case sort: SortExec => sort } + assert(sorts.size == 1 && !sorts.head.global, + s"expected the sort above the grouping to be the one left:\n$plan") + assert(unwrapWrappers(sorts.head.child).isInstanceOf[GroupPartitionsExec], + s"expected the sort to read the grouping:\n$plan") + assert(collectAllGroupPartitions(plan).size == 1, + s"expected the grouping to stay:\n$plan") + assert(ValidateRequirements.validate(plan), s"the combined plan has to hold up:\n$plan") + } + } + } + + test("SPARK-59564: combine adjacent sort aggregates across a narrowed grouping") { + // `max(cast(id as string))` cannot be planned as a hash aggregate, and keeps `id` in the scan's + // output, so the grouping is handed the full keyed partitioning and the positions it projects + // have to be translated onto `name`. The sort feeding the partial aggregate goes with it, as in + // the sort test above. + val partitions = Array(identity("id"), identity("name")) + createTable(items, itemsColumns, partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " + + "(2, 'aa', 20.0, cast('2020-01-01' as timestamp)), " + + "(3, 'cc', 30.0, cast('2020-01-01' as timestamp))") + + val query = s"SELECT name, max(cast(id as string)) FROM testcat.ns.$items GROUP BY name" + val expected = Seq(Row("aa", "2"), Row("cc", "3")) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") { + val plan = sql(query).queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 2 && aggs.forall(_.isInstanceOf[SortAggregateExec]), + s"expected a pair of sort aggregates:\n$plan") + assert(collectAllGroupPartitions(plan).nonEmpty, + s"expected the grouping the pair is separated by:\n$plan") + // Two sorts, one for the partial aggregate and one the final one reads. + assert(collect(plan) { case sort: SortExec => sort }.size == 2, + s"expected two sorts feeding the pair:\n$plan") + } + + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 1, s"expected one combined aggregate, got ${aggs.size}:\n$plan") + assert(aggs.head.aggregateExpressions.forall(_.mode == Complete), + s"expected the combined aggregate to be complete:\n$plan") + val grouping = collectAllGroupPartitions(plan) + assert(grouping.size == 1, s"expected the grouping to stay, got ${grouping.size}:\n$plan") + assert(grouping.head.joinKeyPositions == Some(Seq(1)), + s"expected `name` translated to position 1 of the scan:\n$plan") + // The one sort left is the one the grouping forced above itself. + val sorts = collect(plan) { case sort: SortExec => sort } + assert(sorts.size == 1 && !sorts.head.global, + s"expected one local sort:\n$plan") + assert(unwrapWrappers(sorts.head.child).isInstanceOf[GroupPartitionsExec], + s"expected the sort to read the grouping:\n$plan") + assert(ValidateRequirements.validate(plan), s"the combined plan has to hold up:\n$plan") + } + } + } + + test("SPARK-59564: keep the pair where the source already orders the aggregate's input") { + // The source reports the ordering the sort aggregate needs, so the partial aggregate holds no + // sort of its own. A sort between the pair still lands there, because the grouping does not + // report the ordering it coalesced: removing the partial aggregate would then hand that sort + // the whole scan instead of the partial aggregate's output, which is the trade + // `spark.sql.execution.pushDownLocalSort.throughCardinalityReducer` declines by default. The + // same pair over an unordered source keeps its own sort below the partial aggregate and is the + // shape the sort aggregate test above combines. + val partitions = Array(identity("id"), identity("name")) + val orderedItems = "ordered_aggregate_items" + createTable(orderedItems, itemsColumns, partitions, + ordering = Array( + sort(column("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST), + sort(column("name"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST), + sort(column("price"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST))) + sql(s"INSERT INTO testcat.ns.$orderedItems VALUES " + + "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " + + "(1, 'aa', 20.0, cast('2020-01-01' as timestamp)), " + + "(2, 'bb', 30.0, cast('2020-01-01' as timestamp))") + + val query = + s"SELECT id, name, max(cast(price as string)) FROM testcat.ns.$orderedItems GROUP BY id, name" + val expected = Seq(Row(1L, "aa", "20.0"), Row(2L, "bb", "30.0")) + + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") { + checkAnswer(sql(query), expected) + } + + withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 2, s"expected the pair to be kept, got ${aggs.size}:\n$plan") + // One sort, the one the grouping forced above itself, and the partial aggregate reads the + // scan directly. + val sorts = collect(plan) { case sort: SortExec => sort } + assert(sorts.size == 1 && !sorts.head.global, s"expected one local sort:\n$plan") + assert(unwrapWrappers(sorts.head.child).isInstanceOf[GroupPartitionsExec], + s"expected the sort to read the grouping:\n$plan") + assert(collect(plan) { case agg: BaseAggregateExec => agg }.exists { agg => + unwrapWrappers(agg.child).isInstanceOf[BatchScanExec] + }, s"expected the partial aggregate to read the scan, with no sort of its own:\n$plan") + } + } + } + test("SPARK-59022: keyed shuffle follows the declared partition key order") { val cols = Array( Column.create("id", LongType), From aafc7cb930696898193142ea10ec1ce2a0118e48 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Sun, 20 Sep 2026 18:48:47 +0800 Subject: [PATCH 2/2] [SPARK-59564][SQL] Rebuild the grouping over the new child and hold the derived-ordering bail `withKeyPositionsFor` handed the node to a new child with `copy`, which carries `grouping`, the partitioning it reports and the child partitioning it was decided over. A node reports that partitioning only while its child still reports the one it was planned over, and throws at execution otherwise, so the fold failed on a narrowed grouping instead of folding. Rebuild through `GroupPartitionsExec.apply` with the translated positions, so all three are decided over the new child; an aligned node is turned away, its slot order being the parent distribution's rather than a function of the child. `withKeyPositionsFor`'s doc then credits `AliasAwareOutputExpression.projectKeyedPartitionings` for keeping the projected expressions the child's own: `KeyedPartitioning.project` builds them, but the projection replaces them, and only its no-alias arm keeps the child's expression. A partial aggregate takes that arm, its `resultExpressions` being `groupingAttributes ++ bufferAttributes`. The sort-aggregate test gains an arm for the ordering derived from the partition keys, which leaves the partial aggregate no sort of its own, and asserts what its declared-ordering sibling does: one local sort, reading the grouping, with the partial aggregate on the scan. Assisted-by: deepseek-flash v4.1 --- .../datasources/v2/GroupPartitionsExec.scala | 29 ++++++++++++++----- .../KeyGroupedPartitioningSuite.scala | 22 ++++++++++++++ 2 files changed, 43 insertions(+), 8 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala index 063eb71268b11..4db1d15a36416 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala @@ -406,21 +406,31 @@ case class GroupPartitionsExec( } /** - * This node reading `newChild`, with the key positions it projects moved to where the expressions - * they name sit there, or `None` when it may not be re-parented onto it. It is named for what it - * answers rather than for the node it returns, which is the node it is called on. + * A node over `newChild`, with the key positions this one projects moved to where the expressions + * they name sit there, or `None` when this one may not be rebuilt over it. * * `EnsureRequirements` computed `joinKeyPositions` against the child this node was planned for, * whose partitioning is that child's projected down to the positions the operator above keeps, so * the positions name a key space `newChild` need not share. The projected expressions are the - * planned child's own (`KeyedPartitioning.project` builds them that way), which makes moving them - * a lookup; a child that does not hold one is turned away. + * planned child's own, which makes moving them a lookup; a child that does not hold one is turned + * away. The projection keeps them so (`AliasAwareOutputExpression.projectKeyedPartitionings`): + * with no aliases it has nothing to trade them for, and a partial aggregate has none, its + * `resultExpressions` being `groupingAttributes ++ bufferAttributes`. * - * Moving them is all this does: what the operator above reads, the ordering it was planned + * The node comes back through the factory rather than `copy`: `grouping` and + * `plannedPartitioning` are decided for one child's partitioning, so a node handed another has to + * be decided again over it. + * + * Rebuilding is all this does: what the operator above reads, the ordering it was planned * against included, is the caller's to hold, since nothing here knows what that operator * requires. */ def withKeyPositionsFor(newChild: SparkPlan): Option[GroupPartitionsExec] = { + // An aligned node's slot order comes from the parent distribution rather than the child, so the + // factory cannot re-decide it for `newChild`. + if (expectedKeyCount.isDefined) { + return None + } // The member of each child's partitioning this reads has to be the one `grouping` reads, since // the positions are only meaningful for that member. `representativeOf` answers as the lookup // there does: the first keyed member, nested collections included. @@ -433,9 +443,12 @@ case class GroupPartitionsExec( return None } val positions = joinKeyPositions.fold(plannedInNewChild)(_.map(plannedInNewChild)) - val regrouped = copy( + val regrouped = GroupPartitionsExec( child = newChild, - joinKeyPositions = Option.when(positions != newChildKp.expressions.indices)(positions)) + joinKeyPositions = Option.when(positions != newChildKp.expressions.indices)(positions), + reducers = reducers, + distributePartitions = distributePartitions, + enableSortedMerge = enableSortedMerge) regrouped.copyTagsFrom(this) Some(regrouped) case _ => None diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index 98ef135392bc9..dcbcb1ef1cbd6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -6769,6 +6769,28 @@ class KeyGroupedPartitioningSuite s"expected the grouping to stay:\n$plan") assert(ValidateRequirements.validate(plan), s"the combined plan has to hold up:\n$plan") } + + // The ordering derived from the partition keys satisfies the partial aggregate the way a + // declared one does, so the bail holds for it too. + withSQLConf( + SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "true", + SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") { + val df = sql(query) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + val aggs = collect(plan) { case agg: BaseAggregateExec => agg } + assert(aggs.size == 2, + s"the derived ordering leaves the partial aggregate no sort to give up either:\n$plan") + // One sort, the one the grouping forced above itself, and the partial aggregate reads the + // scan directly. + val sorts = collect(plan) { case sort: SortExec => sort } + assert(sorts.size == 1 && !sorts.head.global, s"expected one local sort:\n$plan") + assert(unwrapWrappers(sorts.head.child).isInstanceOf[GroupPartitionsExec], + s"expected the sort to read the grouping:\n$plan") + assert(aggs.exists { agg => + unwrapWrappers(agg.child).isInstanceOf[BatchScanExec] + }, s"expected the partial aggregate to read the scan, with no sort of its own:\n$plan") + } } }