You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "xumingming (via GitHub)" <gi...@apache.org> on 2026/07/03 08:05:17 UTC

[PR] [SPARK-57928][SQL] Collapse redundant Partial+Final aggregate pair when shuffle is skipped [spark]

xumingming opened a new pull request, #56987:
URL: https://github.com/apache/spark/pull/56987

   
   ### What changes were proposed in this pull request?
   
   When EnsureRequirements skips the shuffle between a two-phase aggregate (Partial + Final) because the child's outputPartitioning already satisfies ClusteredDistribution(groupingKeys), the Final agg is completely redundant: it receives data already aggregated within each partition and does no cross-partition merging.
   
   The new RemoveRedundantAggregates physical rule detects this pattern by matching a Final BaseAggregateExec directly over a Partial BaseAggregateExec of the same concrete type (Hash/ObjectHash/SortAgg) with no Exchange in between, and replaces both with a single Complete-mode node using case class copy(). The rule is registered in QueryExecution.preparations and in AdaptiveSparkPlanExec.queryStagePreparationRules (both after EnsureRequirements), so it fires in both non-AQE and AQE execution paths.
   
   Controlled by spark.sql.execution.removeRedundantAggregates (internal, default true).
   
   ### Why are the changes needed?
   
   A representative pattern is an aggregation over a join where the join keys match the GROUP BY columns:
   
   ```sql
   SELECT t1.user_id, t1.region, SUM(t1.amount) AS total
   FROM events t1
   JOIN users t2 ON t1.user_id = t2.user_id AND t1.region = t2.region
   GROUP BY t1.user_id, t1.region
   ```
   
   When the join is a SortMergeJoin on (user_id, region), the output is already hash-partitioned on those keys. EnsureRequirements skips the shuffle before the Final aggregate, but leaves both Partial and Final nodes in the plan. On large inputs with high row counts per partition key, the redundant aggregate pair causes severe disk spill and widespread fallback to sort-based aggregation.
   
   ### Does this PR introduce _any_ user-facing change?
   
   No.
   
   ### How was this patch tested?
   
   Added tests.
   
   
   ### Was this patch authored or co-authored using generative AI tooling?
   
   No.
   


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-57928][SQL] Collapse redundant Partial+Final aggregate pair when shuffle is skipped [spark]

Posted by "xumingming (via GitHub)" <gi...@apache.org>.
xumingming commented on PR #56987:
URL: https://github.com/apache/spark/pull/56987#issuecomment-5181811901

   Another PR for the same purpose has already been merged, closing it.


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-57928][SQL] Collapse redundant Partial+Final aggregate pair when shuffle is skipped [spark]

Posted by "xumingming (via GitHub)" <gi...@apache.org>.
xumingming commented on PR #56987:
URL: https://github.com/apache/spark/pull/56987#issuecomment-4905574265

   @yadavay-amzn thanks for the review.
   
   I added the two coverage tests:
   
   1. **Strict-subset child partitioning** (`GROUP BY a, b` over a join on `a` only) — confirms that `HashPartitioning(a)` satisfying `ClusteredDistribution(a, b)` is enough for the collapse.
   2. **Global aggregate over a single-partition child** — confirms the `Some(Nil)` / `AllTuples` path.
   
   I did not add the computed-grouping-key case. When I tried the suggested shape (`GROUP BY a + b` over a join on `a + b`), Spark inserts a `Project` that aliases the computed expression before the partial aggregate, so the partial aggregate groups on the aliased attribute rather than the raw expression. Because the join output is partitioned on the raw expression and the aggregate groups on the alias, `EnsureRequirements` still inserts a shuffle, and the rule does not fire. So that path does not actually provide additional rule coverage under the current planner.
   
   Both new tests pass in non-AQE and AQE modes, and `sql/Test/scalastyle` is clean. Please take another look.


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-57928][SQL] Collapse redundant Partial+Final aggregate pair when shuffle is skipped [spark]

Posted by "xumingming (via GitHub)" <gi...@apache.org>.
xumingming commented on PR #56987:
URL: https://github.com/apache/spark/pull/56987#issuecomment-4991888545

   @yadavay-amzn Do you have further comments?


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-57928][SQL] Collapse redundant Partial+Final aggregate pair when shuffle is skipped [spark]

Posted by "xumingming (via GitHub)" <gi...@apache.org>.
xumingming closed pull request #56987: [SPARK-57928][SQL] Collapse redundant Partial+Final aggregate pair when shuffle is skipped
URL: https://github.com/apache/spark/pull/56987


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-57928][SQL] Collapse redundant Partial+Final aggregate pair when shuffle is skipped [spark]

Posted by "xumingming (via GitHub)" <gi...@apache.org>.
xumingming commented on PR #56987:
URL: https://github.com/apache/spark/pull/56987#issuecomment-4884553198

   The golden files changes are due to some partial-final aggs collapsed. Take q23a as an example, before the change, the agg related part are:
   
   ```
            :     :     :                 +- * Filter (22)
            :     :     :                    +- * HashAggregate (21)
            :     :     :                       +- * HashAggregate (20)
            :     :     :                          +- * Project (19)
            :     :     :                             +- * SortMergeJoin Inner (18)
   
   
   (20) HashAggregate [codegen id : 8]
   Input [3]: [d_date#11, i_item_sk#12, _groupingexpression#14]
   Keys [3]: [_groupingexpression#14, i_item_sk#12, d_date#11]
   Functions [1]: [partial_count(1)]
   Aggregate Attributes [1]: [count#15]
   Results [4]: [_groupingexpression#14, i_item_sk#12, d_date#11, count#16]
   
   (21) HashAggregate [codegen id : 8]
   Input [4]: [_groupingexpression#14, i_item_sk#12, d_date#11, count#16]
   Keys [3]: [_groupingexpression#14, i_item_sk#12, d_date#11]
   Functions [1]: [count(1)]
   Aggregate Attributes [1]: [count(1)#17]
   Results [2]: [i_item_sk#12 AS item_sk#18, count(1)#17 AS cnt#19]
   ```
   
   `(20) HashAggregate` and `(21) HashAggregate` are partial-final agg pairs that does not have exchange operator in between, so they are collapsed to a single complete mode agg:
   
   ```
            :     :     :                 +- * Filter (21)
            :     :     :                    +- * HashAggregate (20)
            :     :     :                       +- * Project (19)
            :     :     :                          +- * SortMergeJoin Inner (18)
   
   
   (20) HashAggregate [codegen id : 8]
   Input [3]: [d_date#11, i_item_sk#12, _groupingexpression#14]
   Keys [3]: [_groupingexpression#14, i_item_sk#12, d_date#11]
   Functions [1]: [count(1)]
   Aggregate Attributes [1]: [count(1)#15]
   Results [2]: [i_item_sk#12 AS item_sk#16, count(1)#15 AS cnt#17]
   ```
   
   Since there is less operators now, the operator numbers change, then a lot of diffs in the golden files.


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-57928][SQL] Collapse redundant Partial+Final aggregate pair when shuffle is skipped [spark]

Posted by "yadavay-amzn (via GitHub)" <gi...@apache.org>.
yadavay-amzn commented on code in PR #56987:
URL: https://github.com/apache/spark/pull/56987#discussion_r3526262604


##########
sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantAggregatesSuite.scala:
##########
@@ -0,0 +1,239 @@
+/*
+ * 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
+import org.apache.spark.sql.catalyst.expressions.aggregate.Complete
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper,
+  DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec,
+  ObjectHashAggregateExec, SortAggregateExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+abstract class RemoveRedundantAggregatesSuiteBase
+    extends SharedSparkSession
+    with AdaptiveSparkPlanHelper {
+
+  private def aggNodes(df: DataFrame): Seq[BaseAggregateExec] = {
+    val plan = df.queryExecution.executedPlan
+    collectWithSubqueries(plan) { case a: BaseAggregateExec => a }
+  }
+
+  private def isComplete(a: BaseAggregateExec): Boolean =
+    // Grouping-only aggregates have empty aggregateExpressions; treat them as complete
+    // since there are no function modes to check.
+    a.aggregateExpressions.isEmpty ||
+      a.aggregateExpressions.forall(_.mode == Complete)
+
+  // Run query with the rule enabled, assert it collapsed to a single Complete agg, then
+  // re-run with the rule disabled and use that as the correctness reference.
+  private def checkCollapsed(query: String): DataFrame = {
+    val df = sql(query)
+    val aggs = aggNodes(df)
+    assert(aggs.length == 1, s"Expected 1 agg node but got ${aggs.length}: $aggs")
+    assert(isComplete(aggs.head), s"Expected Complete mode but got: ${aggs.head}")
+    val disabledDf = withSQLConf(SQLConf.REMOVE_REDUNDANT_AGGREGATES_ENABLED.key -> "false") {
+      sql(query)
+    }
+    assert(df.schema == disabledDf.schema,
+      "collapsed node output schema must match the original final agg schema")
+    checkAnswer(df, disabledDf.collect())
+    df
+  }
+
+  // Run query with the rule enabled, assert it kept two agg nodes, then re-run with the
+  // rule disabled and use that as the correctness reference.
+  private def checkNotCollapsed(query: String): DataFrame = {
+    val df = sql(query)
+    val aggs = aggNodes(df)
+    assert(aggs.length == 2, s"Expected 2 agg nodes but got ${aggs.length}: $aggs")
+    val disabledDf = withSQLConf(SQLConf.REMOVE_REDUNDANT_AGGREGATES_ENABLED.key -> "false") {
+      sql(query)
+    }
+    assert(df.schema == disabledDf.schema,
+      "non-collapsed node output schema must match the disabled-rule schema")
+    checkAnswer(df, disabledDf.collect())
+    df
+  }
+
+  // Shared test data: t1 is a fact table with many rows per (a, b) key;
+  // t2 is a lookup table with exactly one row per (a, b) key. Joining t1 with t2
+  // on (a, b) and then aggregating by (a, b) tests that the collapsed Complete-mode
+  // node accumulates all raw t1 rows per key -- a bug in buffer initialization would
+  // produce wrong aggregate values since each key has multiple input rows.
+  //
+  // t1: a in [0,4], b in [0,9] (50 distinct pairs, 20 rows each = 1000 rows total)
+  // t2: exactly the same 50 pairs, one row each -- a cross join of range(5) x range(10)
+  //     ensures no duplicates and Long-typed keys matching t1 without implicit casts.
+  private def createFactAndLookup(): Unit = {
+    spark.range(0, 1000, 1, 4)
+      .selectExpr("(id % 5) as a", "(id % 10) as b", "cast(id as double) as v")
+      .createOrReplaceTempView("t1")
+    sql("SELECT a.id as a, b.id as b FROM range(5) a CROSS JOIN range(10) b")
+      .createOrReplaceTempView("t2")
+  }
+
+  test("HashAgg: grouping only, no aggregate functions") {
+    withTempView("t1", "t2") {
+      createFactAndLookup()
+      checkCollapsed(
+        """SELECT /*+ MERGE(t1, t2) */ t1.a, t1.b
+          |FROM t1 JOIN t2 ON t1.a = t2.a AND t1.b = t2.b
+          |GROUP BY t1.a, t1.b
+          |""".stripMargin)
+    }
+  }
+
+  test("HashAgg: with SUM aggregate function") {
+    withTempView("t1", "t2") {
+      // Each of the 50 (a, b) keys matches 20 t1 rows. sum(v) accumulates 20 raw
+      // values per key. If Complete mode used mergeExpressions (the Final path) it
+      // would treat raw doubles as pre-aggregated buffers and produce wrong sums.
+      createFactAndLookup()
+      val df = checkCollapsed(
+        """SELECT /*+ MERGE(t1, t2) */ t1.a, t1.b, sum(t1.v) as s
+          |FROM t1 JOIN t2 ON t1.a = t2.a AND t1.b = t2.b
+          |GROUP BY t1.a, t1.b
+          |""".stripMargin)
+      assert(aggNodes(df).head.isInstanceOf[HashAggregateExec])
+    }
+  }

Review Comment:
   Non-blocking: every collapse test here joins on the *full* grouping key, so the child ends up clustered on exactly the grouping keys. The rule's motivating case (per the PR description) is really the *subset* case - the child partitioned on a strict subset of the GROUP BY columns (e.g. `GROUP BY a, b` over a join on `a` only), where `HashPartitioning(a)` satisfies `ClusteredDistribution(a, b)` by subset and EnsureRequirements still skips the shuffle. That path is currently only exercised by the TPC-DS golden diffs; a `checkCollapsed` case with a strict-subset partitioning would give it a direct `checkAnswer` assertion.



##########
sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantAggregatesSuite.scala:
##########
@@ -0,0 +1,239 @@
+/*
+ * 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
+import org.apache.spark.sql.catalyst.expressions.aggregate.Complete
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper,
+  DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec,
+  ObjectHashAggregateExec, SortAggregateExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+abstract class RemoveRedundantAggregatesSuiteBase
+    extends SharedSparkSession
+    with AdaptiveSparkPlanHelper {
+
+  private def aggNodes(df: DataFrame): Seq[BaseAggregateExec] = {
+    val plan = df.queryExecution.executedPlan
+    collectWithSubqueries(plan) { case a: BaseAggregateExec => a }
+  }
+
+  private def isComplete(a: BaseAggregateExec): Boolean =
+    // Grouping-only aggregates have empty aggregateExpressions; treat them as complete
+    // since there are no function modes to check.
+    a.aggregateExpressions.isEmpty ||
+      a.aggregateExpressions.forall(_.mode == Complete)
+
+  // Run query with the rule enabled, assert it collapsed to a single Complete agg, then
+  // re-run with the rule disabled and use that as the correctness reference.
+  private def checkCollapsed(query: String): DataFrame = {
+    val df = sql(query)
+    val aggs = aggNodes(df)
+    assert(aggs.length == 1, s"Expected 1 agg node but got ${aggs.length}: $aggs")
+    assert(isComplete(aggs.head), s"Expected Complete mode but got: ${aggs.head}")
+    val disabledDf = withSQLConf(SQLConf.REMOVE_REDUNDANT_AGGREGATES_ENABLED.key -> "false") {
+      sql(query)
+    }
+    assert(df.schema == disabledDf.schema,
+      "collapsed node output schema must match the original final agg schema")
+    checkAnswer(df, disabledDf.collect())
+    df
+  }
+
+  // Run query with the rule enabled, assert it kept two agg nodes, then re-run with the
+  // rule disabled and use that as the correctness reference.
+  private def checkNotCollapsed(query: String): DataFrame = {
+    val df = sql(query)
+    val aggs = aggNodes(df)
+    assert(aggs.length == 2, s"Expected 2 agg nodes but got ${aggs.length}: $aggs")
+    val disabledDf = withSQLConf(SQLConf.REMOVE_REDUNDANT_AGGREGATES_ENABLED.key -> "false") {
+      sql(query)
+    }
+    assert(df.schema == disabledDf.schema,
+      "non-collapsed node output schema must match the disabled-rule schema")
+    checkAnswer(df, disabledDf.collect())
+    df
+  }
+
+  // Shared test data: t1 is a fact table with many rows per (a, b) key;
+  // t2 is a lookup table with exactly one row per (a, b) key. Joining t1 with t2
+  // on (a, b) and then aggregating by (a, b) tests that the collapsed Complete-mode
+  // node accumulates all raw t1 rows per key -- a bug in buffer initialization would
+  // produce wrong aggregate values since each key has multiple input rows.
+  //
+  // t1: a in [0,4], b in [0,9] (50 distinct pairs, 20 rows each = 1000 rows total)
+  // t2: exactly the same 50 pairs, one row each -- a cross join of range(5) x range(10)
+  //     ensures no duplicates and Long-typed keys matching t1 without implicit casts.
+  private def createFactAndLookup(): Unit = {
+    spark.range(0, 1000, 1, 4)
+      .selectExpr("(id % 5) as a", "(id % 10) as b", "cast(id as double) as v")
+      .createOrReplaceTempView("t1")
+    sql("SELECT a.id as a, b.id as b FROM range(5) a CROSS JOIN range(10) b")
+      .createOrReplaceTempView("t2")
+  }
+
+  test("HashAgg: grouping only, no aggregate functions") {
+    withTempView("t1", "t2") {
+      createFactAndLookup()
+      checkCollapsed(
+        """SELECT /*+ MERGE(t1, t2) */ t1.a, t1.b
+          |FROM t1 JOIN t2 ON t1.a = t2.a AND t1.b = t2.b
+          |GROUP BY t1.a, t1.b
+          |""".stripMargin)
+    }
+  }
+
+  test("HashAgg: with SUM aggregate function") {
+    withTempView("t1", "t2") {
+      // Each of the 50 (a, b) keys matches 20 t1 rows. sum(v) accumulates 20 raw
+      // values per key. If Complete mode used mergeExpressions (the Final path) it
+      // would treat raw doubles as pre-aggregated buffers and produce wrong sums.
+      createFactAndLookup()
+      val df = checkCollapsed(
+        """SELECT /*+ MERGE(t1, t2) */ t1.a, t1.b, sum(t1.v) as s
+          |FROM t1 JOIN t2 ON t1.a = t2.a AND t1.b = t2.b
+          |GROUP BY t1.a, t1.b
+          |""".stripMargin)
+      assert(aggNodes(df).head.isInstanceOf[HashAggregateExec])
+    }
+  }
+
+  test("HashAgg: with AVG (multi-buffer aggregate)") {
+    withTempView("t1", "t2") {
+      // AVG maintains two buffer slots (sum + count). With 20 raw rows per key,
+      // a wrong buffer initialisation would produce an incorrect average.
+      createFactAndLookup()
+      checkCollapsed(
+        """SELECT /*+ MERGE(t1, t2) */ t1.a, t1.b, avg(t1.v) as av
+          |FROM t1 JOIN t2 ON t1.a = t2.a AND t1.b = t2.b
+          |GROUP BY t1.a, t1.b
+          |""".stripMargin)
+    }
+  }
+
+  test("ObjectHashAgg: collect_list collapses to Complete mode") {
+    withTempView("t1", "t2") {
+      // collect_list must accumulate all 20 matching t1.v values per key.
+      // A wrong buffer initialisation would produce an empty or truncated list.
+      createFactAndLookup()
+      val df = checkCollapsed(
+        """SELECT /*+ MERGE(t1, t2) */ t1.a, t1.b,
+          |  sort_array(collect_list(t1.v)) as vals
+          |FROM t1 JOIN t2 ON t1.a = t2.a AND t1.b = t2.b
+          |GROUP BY t1.a, t1.b
+          |""".stripMargin)
+      // collect_list is a TypedImperativeAggregate -> ObjectHashAggregateExec.
+      assert(aggNodes(df).head.isInstanceOf[ObjectHashAggregateExec])
+    }
+  }
+
+  test("SortAgg: collapses to Complete mode when shuffle is skipped") {
+    withSQLConf(
+        SQLConf.USE_OBJECT_HASH_AGG.key -> "false",
+        SQLConf.PREFER_SORTMERGEJOIN.key -> "true") {
+      withTempView("t1", "t2") {
+        // ObjectHash disabled -> SortAggregateExec. The join on exactly the grouping
+        // keys makes EnsureRequirements skip the shuffle, triggering the rule.
+        // 20 raw rows per key means a wrong buffer path would produce wrong lists.
+        createFactAndLookup()
+        val df = checkCollapsed(
+          """SELECT /*+ MERGE(t1, t2) */ t1.a, t1.b,
+            |  sort_array(collect_list(t1.v)) as vals
+            |FROM t1 JOIN t2 ON t1.a = t2.a AND t1.b = t2.b
+            |GROUP BY t1.a, t1.b
+            |""".stripMargin)
+        // ObjectHash disabled -> falls through to SortAggregateExec.
+        assert(aggNodes(df).head.isInstanceOf[SortAggregateExec])
+      }
+    }
+  }
+
+  test("Exchange between partial and final prevents collapsing") {
+    withTempView("t") {
+      // A plain range scan gives no special partitioning, so EnsureRequirements inserts an
+      // Exchange between Partial and Final. The rule must NOT fire.
+      spark.range(0, 100, 1, 4)
+        .selectExpr("id % 10 as k", "cast(id as double) as v")
+        .createOrReplaceTempView("t")
+      checkNotCollapsed("SELECT k, sum(v) as s FROM t GROUP BY k")
+    }
+  }
+
+  test("distinct aggregate: correctness preserved") {
+    withTempView("t1", "t2") {
+      // COUNT(DISTINCT v) plan is not directly targeted by this rule (no adjacent

Review Comment:
   Nit: two more shapes worth a quick `checkCollapsed` since each hits a distinct branch of the required-distribution logic and neither is currently exercised - a global aggregate with a single-partition child (no GROUP BY -> `Some(Nil)` -> `AllTuples`), and a computed grouping key (`GROUP BY a + b`, where the collapsed node requires `ClusteredDistribution(Seq(a + b))` rather than a bare attribute). Both are safe as far as I traced, just untested.



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org