(spark) branch master updated: [SPARK-57996][SQL] CoalesceShufflePartitions should coalesce partitioning-aware UnionExec children as a single group

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon-hyun 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 217fb4e43ddf [SPARK-57996][SQL] CoalesceShufflePartitions should coalesce partitioning-aware UnionExec children as a single group
217fb4e43ddf is described below

commit 217fb4e43ddfb0e93903e7081a8fa3e8bef68e85
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Tue Jul 7 21:47:37 2026 -0700

    [SPARK-57996][SQL] CoalesceShufflePartitions should coalesce partitioning-aware UnionExec children as a single group
    
    ### What changes were proposed in this pull request?
    
    `CoalesceShufflePartitions.collectCoalesceGroups` treats every `UnionExec` as a node whose children do not need compatible partitioning, so each union child becomes an independent coalesce group and may be coalesced to a different number of partitions.
    
    `UnionExec.outputPartitioning` can report a real partitioning (`HashPartitioning` / `KeyedPartitioning` / `SinglePartition`) when its children are compatibly partitioned. When it does, downstream operators rely on that partitioning, so the children must stay co-partitioned. This PR makes `childrenNeedCompatiblePartitioning` return `true` for a `UnionExec` that is not a plain union (i.e. reports a real partitioning), so its shuffle stages are coalesced together as a single group. `Unio [...]
    
    ### Why are the changes needed?
    
    When a partitioning-aware union's children are coalesced independently, they end up with different partition counts and are no longer co-partitioned. This does not produce a wrong result: `AdaptiveSparkPlanExec` runs `ValidateRequirements` after each `AQEShuffleReadRule` and reverts the entire `CoalesceShufflePartitions` application when the co-partitioning check fails. The net effect is that shuffle partition coalescing is silently lost whenever a union reports a real partitioning.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No. Results are unchanged; the executed plan may now coalesce shuffle partitions that were previously left un-coalesced.
    
    ### How was this patch tested?
    
    New test in `AdaptiveQueryExecSuite`. Without the fix the coalescing is reverted (no `AQEShuffleReadExec`); with the fix both union children are coalesced to the same number of partitions.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Opus 4.8
    
    🤖 Generated with [Claude Code](https://claude.com/claude-code)
    
    Closes #57077 from ulysses-you/SPARK-57996.
    
    Authored-by: Xiduo You <ul...@gmail.com>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 .../adaptive/CoalesceShufflePartitions.scala       |  6 ++-
 .../sql/execution/basicPhysicalOperators.scala     |  2 +-
 .../adaptive/AdaptiveQueryExecSuite.scala          | 47 ++++++++++++++++++++++
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
index d44667ea2fa3..32021da21171 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
@@ -182,7 +182,11 @@ case class CoalesceShufflePartitions(session: SparkSession) extends AQEShuffleRe
 
   private def childrenNeedCompatiblePartitioning(p: SparkPlan): Boolean = p match {
     // TODO: match more plan nodes here.
-    case _: UnionExec => false
+    // A UnionExec that reports a real outputPartitioning (i.e. it is not a plain union) has that
+    // partitioning relied on by downstream operators, so its children must remain co-partitioned
+    // -- coalesce them as a single group. When it reports UnknownPartitioning (config off, or
+    // children not compatibly partitioned), the children may still be coalesced independently.
+    case u: UnionExec => !u.isPlainUnion
     case _: CartesianProductExec => false
     case _: BroadcastHashJoinExec => false
     case _: BroadcastNestedLoopJoinExec => false
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
index 4e25bf8bef40..79cca1c4cbd8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
@@ -987,7 +987,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup
   // codegen is disabled for it (`supportCodegenFailureReason` reports "partitioning-aware"):
   // the per-partition key descriptor is consumed by a downstream `GroupPartitionsExec`, and
   // keeping these unions out of whole-stage codegen matches the `HashPartitioning` union case.
-  private def isPlainUnion: Boolean = outputPartitioning.isInstanceOf[UnknownPartitioning]
+  private[sql] def isPlainUnion: Boolean = outputPartitioning.isInstanceOf[UnknownPartitioning]
 
   // Per-child projection from the child's output to the union's output. The wrapped
   // child is always the source `Attribute` (deterministic by construction); the Alias
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
index 89d0316a0c59..8cf6fbf921da 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -3599,6 +3599,53 @@ class AdaptiveQueryExecSuite
     }
   }
 
+  test("SPARK-57996: CoalesceShufflePartitions should coalesce partitioning-aware UnionExec " +
+    "children as a single group") {
+    // A UNION ALL of two aggregates on the same key, feeding a downstream aggregate on that key.
+    // With UNION_OUTPUT_PARTITIONING on, the union reports the shared HashPartitioning that the
+    // outer aggregate relies on. The two children have very different data sizes, so coalescing
+    // them independently would produce different partition counts, breaking co-partitioning and
+    // causing AQE to revert the coalescing. They must be coalesced together as a single group.
+    withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+      SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+      SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "1048576",
+      SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "2",
+      SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_SIZE.key -> "1",
+      SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
+      withTable("t1", "t2") {
+        sql("CREATE TABLE t1 USING parquet AS SELECT id AS c1, uuid() AS c2 FROM range(10)")
+        sql("CREATE TABLE t2 USING parquet AS SELECT id AS c1, uuid() AS c2 FROM range(100)")
+
+        val query =
+          """
+            |SELECT c1, c2 FROM (
+            |  SELECT c1, c2 FROM t1 GROUP BY c1, c2
+            |  UNION ALL
+            |  SELECT c1, c2 FROM t2 GROUP BY c1, c2
+            |) GROUP BY c1, c2
+            |""".stripMargin
+
+        val correctResults = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") {
+          sql(query).collect()
+        }
+
+        withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") {
+          val df = sql(query)
+          df.collect()
+          val reads = collect(df.queryExecution.executedPlan) {
+            case r: AQEShuffleReadExec => r
+          }
+          // Both union children are coalesced (the plan is not reverted) and to the same number
+          // of partitions, so they remain co-partitioned.
+          assert(reads.size === 2)
+          assert(reads.forall(_.hasCoalescedPartition))
+          assert(reads.map(_.partitionSpecs.length).distinct.length === 1)
+          checkAnswer(df, correctResults)
+        }
+      }
+    }
+  }
+
   test("SPARK-44065: Optimize BroadcastHashJoin skew") {
     withSQLConf(
       SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",


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