From 2252c4212d71377b146da0b75ac708ed536c67f9 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Fri, 24 Jul 2026 13:09:38 +0800 Subject: [PATCH 1/2] [SPARK-58317][SQL] Union output partitioning should support PartitioningCollection children ### What changes were proposed in this pull request? `UnionExec.outputPartitioning` (SPARK-52921) passes a child partitioning through the union when every child is compatibly partitioned, letting a downstream operator (e.g. an aggregate) skip a shuffle. This PR extends that pass-through to children whose `outputPartitioning` is a `PartitioningCollection`. - `outputPartitioning` now treats 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 passes through the intersection across all children: empty -> `UnknownPartitioning`; one member -> that partitioning; many -> a `PartitioningCollection`. Only index-co-locatable partitionings (`HashPartitioningLike` / `SinglePartition`) participate. - The existing `KeyedPartitioning` concatenation path is kept as a separate case that fires only when every child is a single `KeyedPartitioning`. It is a distinct physical strategy (`sparkContext.union`, numPartitions = sum) and cannot be folded into the co-located intersection (`SQLPartitioningAwareUnionRDD`, numPartitions = N), because a `PartitioningCollection` requires uniform numPartitions across its members. ### Why are the changes needed? `comparePartitioning` had no case for `PartitioningCollection`, so whenever a child reported one, the whole union collapsed to `UnknownPartitioning` and an extra shuffle was inserted. This is hit by a common query shape: a `UNION ALL` where one branch is an inner shuffled-hash join (whose `outputPartitioning` is `PartitioningCollection(Hash(k1), Hash(k2))`) and the other is a left join (single `HashPartitioning`), feeding a `GROUP BY`. The union failed to pass through and a redundant `ENSURE_REQUIREMENTS` shuffle appeared before the aggregate. ### Does this PR introduce _any_ user-facing change? No behavior change visible in query results. Under the default-on `spark.sql.unionOutputPartitioning`, affected plans drop a redundant shuffle. ### How was this patch tested? Added tests to `DataFrameSetOperationsSuite` covering: the mixed inner+left join union (intersection to a single `HashPartitioning`, shuffle eliminated); all children reporting collections (pass-through as a `PartitioningCollection`); empty intersection (fall back to `UnknownPartitioning`); and the collection pass-through under AQE. Existing `DataFrameSetOperationsSuite` and `KeyGroupedPartitioningSuite` (97 tests) regressions pass. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8) Co-Authored-By: Claude Opus 4.8 --- .../execution/basicPhysicalOperators.scala | 110 +++++++---- .../sql/DataFrameSetOperationsSuite.scala | 183 +++++++++++++++++- 2 files changed, 253 insertions(+), 40 deletions(-) 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 79cca1c4cbd83..3e1459693c481 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 @@ -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,76 @@ 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. + def flattenPartitioning(p: Partitioning): Seq[Partitioning] = p match { + case pc: PartitioningCollection => pc.partitionings.flatMap(flattenPartitioning) + case other => Seq(other) + } + val candidateSets = partitionings.map { p => + flattenPartitioning(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)) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala index 7e9a7ad745794..128896bfd15e2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/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,187 @@ 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). The union should + // intersect the two to a single HashPartitioning(c1) and let the group-by skip a shuffle. + def unionDF: DataFrame = sql( + """SELECT c1, c2, c3, 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, t3.c1 AS k + | FROM t3 LEFT JOIN t4 ON t3.c1 = t4.c1 + |) GROUP BY c1, c2, c3 + |""".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") + + def unionDF: DataFrame = sql( + """SELECT c1, 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 + |""".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") { From b79f1b265e452136bbbce9f3a35a9a316eb4651a Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Sat, 25 Jul 2026 05:26:27 +0800 Subject: [PATCH 2/2] [SPARK-58317][SQL] Address review: fix test coverage and dedup flatten helper Address review comments from dongjoon-hyun and viirya: - The "intersects to single" and AQE tests passed without this PR because `k` was unused above the union, so ColumnPruning dropped it and each ProjectExec narrowed the join's PartitioningCollection to a single HashPartitioning. Reference `k` in the group-by to keep the collection alive, and alias t4.c1 (the left join's build side) so the second branch stays a single Hash(c1). The tests now fail on pre-PR code and pass with the fix, genuinely exercising the PartitioningCollection intersection path. - Lift the duplicated `flattenPartitioning` local helper into a shared `PartitioningCollection.flatten` and reuse it from both UnionExec and PartitioningPreservingUnaryExecNode. Co-Authored-By: Claude Opus 4.8 --- .../plans/physical/partitioning.scala | 9 +++++++++ .../AliasAwareOutputExpression.scala | 12 ++---------- .../execution/basicPhysicalOperators.scala | 6 +----- .../sql/DataFrameSetOperationsSuite.scala | 19 ++++++++++++------- 4 files changed, 24 insertions(+), 22 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala index d2bb12d2053aa..76e667590a33c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala +++ b/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 + } } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala index b37e1b258e9bb..ef42595047489 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala +++ b/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 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 3e1459693c481..e94b774a37d19 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,12 +987,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // 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. - def flattenPartitioning(p: Partitioning): Seq[Partitioning] = p match { - case pc: PartitioningCollection => pc.partitionings.flatMap(flattenPartitioning) - case other => Seq(other) - } val candidateSets = partitionings.map { p => - flattenPartitioning(p).filter { + PartitioningCollection.flatten(p).filter { case _: HashPartitioningLike => true case SinglePartition => true case _ => false diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala index 128896bfd15e2..faefbab6c8ca9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -1672,16 +1672,19 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP // 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). The union should - // intersect the two to a single HashPartitioning(c1) and let the group-by skip a shuffle. + // 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, count(*) FROM ( + """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, t3.c1 AS k + | 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 + |) GROUP BY c1, c2, c3, k |""".stripMargin) val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { @@ -1817,14 +1820,16 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP 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, count(*) FROM ( + """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 + |) GROUP BY c1, k |""".stripMargin) val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") {