[SPARK-58317][SQL] Union output partitioning should support PartitioningCollection children

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala

@@ -913,6 +913,15 @@ object PartitioningCollection {
     }
     new PartitioningCollection(partitionings.map(intern))
   }
+
+  /**
+   * Flattens a partitioning into its leaf partitionings: a [[PartitioningCollection]] is
+   * recursively replaced by its members, and any other partitioning yields itself.
+   */
+  def flatten(partitioning: Partitioning): Seq[Partitioning] = partitioning match {
+    case PartitioningCollection(partitionings) => partitionings.flatMap(flatten)
+    case other => other +: Nil
+  }
 }
 
 /**

sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala

@@ -31,7 +31,8 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
   with AliasAwareOutputExpression {
   final override def outputPartitioning: Partitioning = {
     val (keyedPartitionings, otherPartitionings) =
-      flattenPartitioning(child.outputPartitioning).partition(_.isInstanceOf[KeyedPartitioning])
+      PartitioningCollection.flatten(child.outputPartitioning)
+        .partition(_.isInstanceOf[KeyedPartitioning])
 
     val projectedKPs =
       projectKeyedPartitionings(keyedPartitionings.map(_.asInstanceOf[KeyedPartitioning]))
@@ -148,15 +149,6 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
       .map(projectedExprs =>
         new KeyedPartitioning(projectedExprs, sharedKeys, isGrouped, isNarrowed))
   }
-
-  private def flattenPartitioning(partitioning: Partitioning): Seq[Partitioning] = {
-    partitioning match {
-      case PartitioningCollection(childPartitionings) =>
-        childPartitionings.flatMap(flattenPartitioning)
-      case rest =>
-        rest +: Nil
-    }
-  }
 }
 
 trait OrderPreservingUnaryExecNode

sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala

@@ -919,6 +919,9 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup
     Seq(firstPartitioning) ++ convertedOtherPartitionings
   }
 
+  // Compares two leaf partitionings for union pass-through equivalence. Callers pass leaf
+  // partitionings only; a `PartitioningCollection` is flattened to its members by
+  // `outputPartitioning` before reaching here.
   private def comparePartitioning(left: Partitioning, right: Partitioning): Boolean = {
     (left, right) match {
       case (SinglePartition, SinglePartition) => true
@@ -937,47 +940,72 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup
   }
 
   override def outputPartitioning: Partitioning = {
-    if (conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) {
-      val partitionings = prepareOutputPartitioning()
-      if (partitionings.forall(comparePartitioning(_, partitionings.head))) {
-        val partitioner = partitionings.head
-
-        // Take the output attributes of this union and map the partitioner to them.
-        val attributeMap = children.head.output.zip(output).toMap
-        partitioner match {
-          case headKp: KeyedPartitioning =>
-            // A `UnionExec` concatenates its children's partitions in order (one child's
-            // partitions after another's), so the merged `KeyedPartitioning` carries the
-            // concatenation of the children's partition keys, one key per physical output
-            // partition. Children usually hold different key sets, so the merged keys often
-            // contain duplicates and `isGrouped` is false; a downstream `GroupPartitionsExec`
-            // regroups partitions that share a key. The children's expressions have already
-            // been remapped to the first child's attributes by `prepareOutputPartitioning`;
-            // here they are remapped to the union's output attributes.
-            val mergedKeys = partitionings.flatMap {
-              case k: KeyedPartitioning => k.partitionKeys
-              case _ => return super.outputPartitioning
-            }
-            val mergedExpressions = headKp.expressions.map(_.transform {
-              case a: Attribute if attributeMap.contains(a) => attributeMap(a)
-            })
-            val isGrouped = mergedKeys.distinct.size == mergedKeys.size
-            val isNarrowed = partitionings.exists {
-              case k: KeyedPartitioning => k.isNarrowed
-              case _ => false
-            }
-            KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isNarrowed)
-          case e: Expression =>
-            e.transform {
-              case a: Attribute if attributeMap.contains(a) => attributeMap(a)
-            }.asInstanceOf[Partitioning]
-          case _ => partitioner
-        }
+    if (!conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) {
+      return super.outputPartitioning
+    }
+
+    // Children's partitionings with attributes remapped to the first child's attributes.
+    val partitionings = prepareOutputPartitioning()
+    // Map from the first child's attributes to this union's own output attributes.
+    val attributeMap = children.head.output.zip(output).toMap
+    def toUnionOutput(p: Partitioning): Partitioning = p match {
+      case e: Expression =>
+        e.transform {
+          case a: Attribute if attributeMap.contains(a) => attributeMap(a)
+        }.asInstanceOf[Partitioning]
+      case _ => p
+    }
+
+    // Case A: every child is a single `KeyedPartitioning`. A `UnionExec` concatenates its
+    // children's partitions in order (one child's partitions after another's), so the merged
+    // `KeyedPartitioning` carries the concatenation of the children's partition keys, one key
+    // per physical output partition. Children usually hold different key sets, so the merged
+    // keys often contain duplicates and `isGrouped` is false; a downstream `GroupPartitionsExec`
+    // regroups partitions that share a key. This concatenation (numPartitions = sum) is a
+    // distinct physical strategy from the co-located pass-through below (numPartitions = N), so
+    // it is kept as a separate case and never folded into a `PartitioningCollection`.
+    if (partitionings.forall(_.isInstanceOf[KeyedPartitioning])) {
+      val kps = partitionings.map(_.asInstanceOf[KeyedPartitioning])
+      val headKp = kps.head
+      // The `KeyedPartitioning`s must agree on the partition expressions to merge.
+      val compatible = kps.forall(comparePartitioning(_, headKp))
+      if (compatible) {
+        val mergedKeys = kps.flatMap(_.partitionKeys)
+        val mergedExpressions = headKp.expressions.map(_.transform {
+          case a: Attribute if attributeMap.contains(a) => attributeMap(a)
+        })
+        val isGrouped = mergedKeys.distinct.size == mergedKeys.size
+        val isNarrowed = kps.exists(_.isNarrowed)
+        return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isNarrowed)
       } else {
-        super.outputPartitioning
+        return super.outputPartitioning
       }
-    } else {
-      super.outputPartitioning
+    }
+
+    // Case B: treat each child's partitioning as a set of candidate partitionings (a
+    // `PartitioningCollection` flattens to its members; a single partitioning is a one-element
+    // set) and pass through the intersection across all children. Only index-co-locatable
+    // partitionings participate; `KeyedPartitioning` is excluded here because its concatenation
+    // semantics (Case A) are incompatible with the co-located union RDD.
+    val candidateSets = partitionings.map { p =>
+      PartitioningCollection.flatten(p).filter {
+        case _: HashPartitioningLike => true
+        case SinglePartition => true
+        case _ => false
+      }
+    }
+    // Intersect across all children, anchored on the first child's set. Every surviving member
+    // is shared by all children, so their `numPartitions` agree; a `PartitioningCollection`
+    // built from a subset of one child's members therefore keeps its uniform-numPartitions
+    // invariant, and the co-located `doExecute` arm's invariant holds.
+    val head = candidateSets.head
+    val intersection = head.filter { c =>
+      candidateSets.tail.forall(_.exists(comparePartitioning(c, _)))
+    }
+    intersection match {
+      case Seq() => super.outputPartitioning
+      case Seq(p) => toUnionOutput(p)
+      case ps => PartitioningCollection.fromPartitionings(ps.map(toUnionOutput))
     }
   }
 

sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala

@@ -22,7 +22,7 @@ import java.util.Locale
 
 import org.apache.spark.sql.catalyst.optimizer.RemoveNoopUnion
 import org.apache.spark.sql.catalyst.plans.logical.Union
-import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, UnknownPartitioning}
+import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, KeyedPartitioning, PartitioningCollection, UnknownPartitioning}
 import org.apache.spark.sql.connector.catalog.InMemoryCatalog
 import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
@@ -1659,6 +1659,192 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP
     }
   }
 
+  test("SPARK-58317: union partitioning - PartitioningCollection child intersects to single") {
+    withSQLConf(
+        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+        SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+      withTempView("t1", "t2", "t3", "t4") {
+        Seq((1, 2, 4), (1, 3, 5), (2, 2, 3)).toDF("c1", "c2", "c3").createOrReplaceTempView("t1")
+        Seq((1, 9), (2, 9)).toDF("c1", "x").createOrReplaceTempView("t2")
+        Seq((1, 2, 4), (2, 4, 5), (3, 6, 7)).toDF("c1", "c2", "c3").createOrReplaceTempView("t3")
+        Seq((1, 9), (3, 9)).toDF("c1", "y").createOrReplaceTempView("t4")
+
+        // The first branch is an inner shuffled-hash join and selects both join keys (t1.c1 and
+        // t2.c1), so its output partitioning is a PartitioningCollection(Hash(c1), Hash(c1#..))
+        // that a downstream ProjectExec cannot narrow to a single member. The second branch is a
+        // left join, whose output partitioning is a single HashPartitioning(c1). Referencing `k`
+        // in the group-by keeps the first branch's collection alive (otherwise ColumnPruning drops
+        // it and the join narrows to a single Hash(c1)); aliasing t4.c1 (the build side of the
+        // left join) keeps the second branch a single Hash(c1). The union intersects the two to a
+        // single HashPartitioning(c1) and lets the group-by skip a shuffle.
+        def unionDF: DataFrame = sql(
+          """SELECT c1, c2, c3, k, count(*) FROM (
+            |  SELECT /*+ SHUFFLE_HASH(t2) */ t1.c1, t1.c2, t1.c3, t2.c1 AS k
+            |  FROM t1 JOIN t2 ON t1.c1 = t2.c1
+            |  UNION ALL
+            |  SELECT /*+ SHUFFLE_HASH(t4) */ t3.c1, t3.c2, t3.c3, t4.c1 AS k
+            |  FROM t3 LEFT JOIN t4 ON t3.c1 = t4.c1
+            |) GROUP BY c1, c2, c3, k
+            |""".stripMargin)
+
+        val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") {
+          unionDF.collect()
+        }
+
+        val shuffleNums = Seq(true, false).map { enabled =>
+          withSQLConf(
+              SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+              SQLConf.UNION_OUTPUT_PARTITIONING.key -> enabled.toString) {
+            val union = unionDF
+            val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u }
+            assert(unionExec.size == 1)
+
+            val partitioning = unionExec.head.outputPartitioning
+            if (enabled) {
+              assert(partitioning.isInstanceOf[HashPartitioning],
+                s"expected a HashPartitioning pass-through but got $partitioning")
+            } else {
+              assert(partitioning.isInstanceOf[UnknownPartitioning])
+            }
+
+            checkAnswer(union, correctResult)
+            union.queryExecution.executedPlan.collect {
+              case s: ShuffleExchangeExec => s
+            }.size
+          }
+        }
+        // Enabling the pass-through removes the shuffle before the aggregate.
+        assert(shuffleNums.head + 1 == shuffleNums.last)
+      }
+    }
+  }
+
+  test("SPARK-58317: union partitioning - all PartitioningCollection children pass through") {
+    withSQLConf(
+        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+        SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+      withTempView("t1", "t2", "t3", "t4") {
+        Seq((1, 2), (2, 3)).toDF("c1", "c2").createOrReplaceTempView("t1")
+        Seq((1, 9), (2, 9)).toDF("c1", "x").createOrReplaceTempView("t2")
+        Seq((1, 2), (3, 4)).toDF("c1", "c2").createOrReplaceTempView("t3")
+        Seq((1, 9), (3, 9)).toDF("c1", "y").createOrReplaceTempView("t4")
+
+        // Both branches are inner shuffled-hash joins on a single key and select both sides' join
+        // key (t1.c1 and t2.c1 AS k), so a downstream ProjectExec cannot narrow either child's
+        // PartitioningCollection(Hash(c1), Hash(k)) to a single member. The union should intersect
+        // to a PartitioningCollection carrying both members.
+        def unionDF: DataFrame = sql(
+          """SELECT /*+ SHUFFLE_HASH(t2) */ t1.c1, t2.c1 AS k
+            |FROM t1 JOIN t2 ON t1.c1 = t2.c1
+            |UNION ALL
+            |SELECT /*+ SHUFFLE_HASH(t4) */ t3.c1, t4.c1 AS k
+            |FROM t3 JOIN t4 ON t3.c1 = t4.c1
+            |""".stripMargin)
+
+        val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") {
+          unionDF.collect()
+        }
+
+        Seq(true, false).foreach { enabled =>
+          withSQLConf(
+              SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+              SQLConf.UNION_OUTPUT_PARTITIONING.key -> enabled.toString) {
+            val union = unionDF
+            val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u }
+            assert(unionExec.size == 1)
+
+            val partitioning = unionExec.head.outputPartitioning
+            if (enabled) {
+              assert(partitioning.isInstanceOf[PartitioningCollection],
+                s"expected a PartitioningCollection pass-through but got $partitioning")
+              val members = partitioning.asInstanceOf[PartitioningCollection].partitionings
+              assert(members.forall(_.isInstanceOf[HashPartitioning]))
+              assert(members.size == 2)
+            } else {
+              assert(partitioning.isInstanceOf[UnknownPartitioning])
+            }
+
+            checkAnswer(union, correctResult)
+          }
+        }
+      }
+    }
+  }
+
+  test("SPARK-58317: union partitioning - empty intersection falls back") {
+    withSQLConf(
+        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+        SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+      withTempView("t1", "t2") {
+        Seq((1, 2, 4), (2, 3, 5)).toDF("c1", "c2", "c3").createOrReplaceTempView("t1")
+        Seq((1, 9), (2, 9)).toDF("c1", "x").createOrReplaceTempView("t2")
+
+        // First branch reports PartitioningCollection(Hash(c1), Hash(c1#..)); the second branch
+        // is repartitioned on a disjoint column, so the intersection is empty.
+        def unionDF: DataFrame = sql(
+          """SELECT /*+ SHUFFLE_HASH(t2) */ t1.c1, t1.c2, t1.c3, t2.c1 AS k
+            |FROM t1 JOIN t2 ON t1.c1 = t2.c1
+            |UNION ALL
+            |SELECT c1, c2, c3, c2 AS k FROM t1 DISTRIBUTE BY c2
+            |""".stripMargin)
+
+        val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") {
+          unionDF.collect()
+        }
+
+        Seq(true, false).foreach { enabled =>
+          withSQLConf(
+              SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+              SQLConf.UNION_OUTPUT_PARTITIONING.key -> enabled.toString) {
+            val union = unionDF
+            val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u }
+            assert(unionExec.size == 1)
+            assert(unionExec.head.outputPartitioning.isInstanceOf[UnknownPartitioning])
+            checkAnswer(union, correctResult)
+          }
+        }
+      }
+    }
+  }
+
+  test("SPARK-58317: union partitioning - PartitioningCollection pass-through under AQE") {
+    // AQE is enabled by default in production; the collection pass-through must produce correct
+    // results there too, where the union output flows through the coalesce-compatibility path.
+    withSQLConf(
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+        SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+      withTempView("t1", "t2", "t3", "t4") {
+        Seq((1, 2), (2, 3), (1, 4)).toDF("c1", "c2").createOrReplaceTempView("t1")
+        Seq((1, 9), (2, 9)).toDF("c1", "x").createOrReplaceTempView("t2")
+        Seq((1, 2), (3, 4)).toDF("c1", "c2").createOrReplaceTempView("t3")
+        Seq((1, 9), (3, 9)).toDF("c1", "y").createOrReplaceTempView("t4")
+
+        // Referencing `k` in the group-by keeps each branch's PartitioningCollection alive under
+        // AQE (otherwise ColumnPruning drops it and the join narrows to a single Hash(c1)).
+        def unionDF: DataFrame = sql(
+          """SELECT c1, k, count(*) FROM (
+            |  SELECT /*+ SHUFFLE_HASH(t2) */ t1.c1, t2.c1 AS k
+            |  FROM t1 JOIN t2 ON t1.c1 = t2.c1
+            |  UNION ALL
+            |  SELECT /*+ SHUFFLE_HASH(t4) */ t3.c1, t4.c1 AS k
+            |  FROM t3 JOIN t4 ON t3.c1 = t4.c1
+            |) GROUP BY c1, k
+            |""".stripMargin)
+
+        val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") {
+          unionDF.collect()
+        }
+
+        Seq(true, false).foreach { enabled =>
+          withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> enabled.toString) {
+            checkAnswer(unionDF, correctResult)
+          }
+        }
+      }
+    }
+  }
+
   test("SPARK-53550: union partitioning should compare canonicalized attributes") {
     withSQLConf(
       SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {