This is an automated email from the ASF dual-hosted git repository.
ulysses-you 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 2d4b57ac7802 [SPARK-58819][SQL] UnionExec outputPartitioning should
compare children in the union's attribute space
2d4b57ac7802 is described below
commit 2d4b57ac7802e9f4e84babcc5e4bb0a36e814ac4
Author: Xiduo You <[email protected]>
AuthorDate: Fri Aug 21 09:26:56 2026 +0800
[SPARK-58819][SQL] UnionExec outputPartitioning should compare children in
the union's attribute space
### What changes were proposed in this pull request?
Change `UnionExec.prepareOutputPartitioning` to map every child's
partitioning (including the first child's) to this union's output attributes,
instead of only remapping the non-first children to the first child's
attributes. The now-redundant `toUnionOutput` step is removed, and
`comparePartitioning` keeps using structural equality (`l == r`) for hash
partitionings.
### Why are the changes needed?
Previously `prepareOutputPartitioning` left the first child's
`outputPartitioning` untouched while remapping the other children to the first
child's output attributes. A child's `outputPartitioning` is not necessarily
expressed in its own output attributes: `AttributeReference.equals` compares
name, nullability, metadata, and qualifier, and a child's partitioning can
differ from its output in any of them. For example,
`FilterExec.outputPartitioning` passes through its child's partitio [...]
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added two tests in `DataFrameSetOperationsSuite`:
- `SPARK-58819: union outputPartitioning compares children in the union's
attribute space` (SQL), where the first child's partition key carries a
subquery qualifier.
- `SPARK-58819: union outputPartitioning ignores partition key nullability`
(DataFrame), where a `Filter` narrows the partition key's nullability.
Both verify that the union still propagates `HashPartitioning(id)`, that a
downstream `groupBy(id).count()` reuses the partitioning instead of adding a
shuffle, and that the propagated partitioning is expressed in the union's own
output attributes.
Also ran `DataFrameSetOperationsSuite`, `UnionCodegenSuite`,
`CoalesceShufflePartitionsSuite`, `PlannerSuite`, `AdaptiveQueryExecSuite`, and
the TPC-DS plan stability suites (`TPCDSV1_4_PlanStabilitySuite`,
`TPCDSV2_7_PlanStabilitySuite`).
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude Code
Closes #58044 from ulysses-you/union-partitioning-miss.
Authored-by: Xiduo You <[email protected]>
Signed-off-by: Xiduo You <[email protected]>
---
.../sql/execution/basicPhysicalOperators.scala | 52 +++----
.../spark/sql/DataFrameSetOperationsSuite.scala | 154 +++++++++++++++++++++
2 files changed, 175 insertions(+), 31 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 9f9eae3d86f2..bea86501e6f3 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
@@ -915,22 +915,24 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
}
/**
- * Returns the output partitionings of the children, with the attributes
converted to
- * the first child's attributes at the same position.
+ * Returns the output partitionings of the children, with the attributes
converted to this
+ * union's output attributes at the same position.
*/
private def prepareOutputPartitioning(): Seq[Partitioning] = {
- // Create a map of attributes from the other children to the first child.
- val firstAttrs = children.head.output
- val attributesMap = children.tail.map(_.output).map { otherAttrs =>
- AttributeMap(otherAttrs.zip(firstAttrs))
+ // Map every child's partitioning attributes to this union's output
attributes, so all
+ // partitionings are expressed in the same attribute space before
comparison. A child's
+ // `outputPartitioning` may reference attributes that differ from its own
`output` in any
+ // field `AttributeReference.equals` compares other than `dataType`, which
two attributes
+ // sharing an `ExprId` agree on (so name, nullability, metadata,
qualifier): a Filter
+ // narrows nullability via `IsNotNull` while passing its child's
partitioning through, and
+ // a partitioning built inside a view or subquery carries that relation's
qualifier.
+ // `AttributeMap` is keyed by `ExprId`, so remapping every child
(including the first)
+ // normalizes all of those.
+ val unionOutput = output
+ val attributesMap = children.map(_.output).map { childAttrs =>
+ AttributeMap(childAttrs.zip(unionOutput))
}
-
- val partitionings = children.map(_.outputPartitioning)
- val firstPartitioning = partitionings.head
- val otherPartitionings = partitionings.tail
-
- val convertedOtherPartitionings = otherPartitionings.zipWithIndex.map {
case (p, idx) =>
- val attributeMap = attributesMap(idx)
+ children.map(_.outputPartitioning).zip(attributesMap).map { case (p,
attributeMap) =>
p match {
case e: Expression =>
e.transform {
@@ -940,7 +942,6 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
case _ => p
}
}
- Seq(firstPartitioning) ++ convertedOtherPartitionings
}
// Compares two leaf partitionings for union pass-through equivalence.
Callers pass leaf
@@ -950,8 +951,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
(left, right) match {
case (SinglePartition, SinglePartition) => true
case (l: HashPartitioningLike, r: HashPartitioningLike) => l == r
- // For `KeyedPartitioning`, only the partition expressions must match
(the other child's
- // expressions have already been remapped to the first child's
attributes by
+ // For `KeyedPartitioning`, only the partition expressions must match
(both sides'
+ // expressions have already been remapped to this union's output
attributes by
// `prepareOutputPartitioning`). The partition keys are intentionally
not compared here:
// children typically carry different key sets, and `outputPartitioning`
merges them.
case (l: KeyedPartitioning, r: KeyedPartitioning) =>
@@ -968,17 +969,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
return super.outputPartitioning
}
- // Children's partitionings with attributes remapped to the first child's
attributes.
+ // Children's partitionings with attributes remapped to this union's
output 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
@@ -995,9 +987,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
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 mergedExpressions = headKp.expressions
val isGrouped = mergedKeys.distinct.size == mergedKeys.size
val isNarrowed = kps.exists(_.isNarrowed)
return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped,
isNarrowed)
@@ -1028,8 +1018,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
}
intersection match {
case Seq() => super.outputPartitioning
- case Seq(p) => toUnionOutput(p)
- case ps =>
PartitioningCollection.fromPartitionings(ps.map(toUnionOutput))
+ case Seq(p) => p
+ case ps => PartitioningCollection.fromPartitionings(ps)
}
}
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 faefbab6c8ca..32b6da3a30b1 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
@@ -20,6 +20,7 @@ package org.apache.spark.sql
import java.sql.{Date, Timestamp}
import java.util.Locale
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
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.{HashPartitioning,
KeyedPartitioning, PartitioningCollection, UnknownPartitioning}
@@ -1589,6 +1590,159 @@ class DataFrameSetOperationsSuite extends
SharedSparkSession with AdaptiveSparkP
}
}
+ test("SPARK-58819: union outputPartitioning compares children in the union's
attribute space") {
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ withTempView("t1", "t2") {
+ // `DISTRIBUTE BY id` resolves its key with the `t1` qualifier inside
the subquery, but
+ // the outer `Project` exposes `id` with the subquery qualifier, so
child 0's partitioning
+ // differs from its output in qualifier, not nullability. The branches
overlap on `id = 1`,
+ // so a shared key's grouping depends on the union's co-location claim
being honored
+ // (guarded by `checkAnswer(grouped, ...)` below).
+ Seq((Option(1), 10), (Option(2), 20)).toDF("id",
"v").createOrReplaceTempView("t1")
+ Seq((Option(1), 30), (Option(3), 40)).toDF("id",
"v").createOrReplaceTempView("t2")
+
+ val sqlText =
+ """
+ |SELECT id, v FROM (SELECT id, v FROM t1 DISTRIBUTE BY id) WHERE
id IS NOT NULL
+ |UNION ALL
+ |SELECT id, sum(v) AS v FROM t2 GROUP BY id
+ |""".stripMargin
+ val union = spark.sql(sqlText)
+ val unionExec = union.queryExecution.executedPlan.collect { case u:
UnionExec => u }
+ assert(unionExec.size == 1)
+
+ // Pin the discriminator to qualifier: child 0's partitioning
references the same column
+ // (`id`) as its output, differing only in qualifier (`[t1]` vs the
subquery qualifier),
+ // never nullability. The explicit nullability check rules out the
nullability cause the
+ // DataFrame test below exercises; a future refactor leaking a
nullability difference
+ // fails loudly instead of silently changing what the test covers.
+ val child0 = unionExec.head.children.head
+ val outputId = child0.output.head.asInstanceOf[AttributeReference]
+ val partitioningId =
+
child0.outputPartitioning.asInstanceOf[HashPartitioning].expressions.head
+ .asInstanceOf[AttributeReference]
+ assert(outputId.exprId == partitioningId.exprId,
+ s"expected the same column: $outputId vs $partitioningId")
+ assert(outputId.nullable == partitioningId.nullable,
+ s"nullability must match, qualifier is the sole difference: " +
+ s"$outputId vs $partitioningId")
+ assert(outputId.qualifier != partitioningId.qualifier,
+ s"expected a qualifier difference: $outputId vs $partitioningId")
+
+ // Child 0's `HashPartitioning` references `id` with the `t1`
qualifier while its output
+ // `id` carries the subquery qualifier; remapping both to the union's
output attributes
+ // still propagates the hash partitioning despite the qualifier
difference. The propagated
+ // partitioning must be expressed in the union's own output attribute
(the subquery
+ // qualifier), not child 0's `[t1]` attribute, since `toUnionOutput`
was removed.
+
assert(unionExec.head.outputPartitioning.isInstanceOf[HashPartitioning],
+ s"expected a HashPartitioning pass-through but got
${unionExec.head.outputPartitioning}")
+ val hashPartitioning =
+ unionExec.head.outputPartitioning.asInstanceOf[HashPartitioning]
+ assert(hashPartitioning.expressions == Seq(unionExec.head.output.head))
+
+ // The two branches contribute one shuffle each (DISTRIBUTE BY and
GROUP BY). The propagated
+ // HashPartitioning lets the downstream group-by reuse them instead of
adding a third.
+ val unionShuffles = union.queryExecution.executedPlan.collect {
+ case s: ShuffleExchangeExec => s
+ }.size
+ val grouped = union.groupBy($"id").count()
+ val groupedShuffles = grouped.queryExecution.executedPlan.collect {
+ case s: ShuffleExchangeExec => s
+ }.size
+ assert(unionShuffles == 2, s"union should have 2 shuffles but got
$unionShuffles")
+ assert(groupedShuffles == 2,
+ s"group-by should reuse the union's partitioning (expect 2 shuffles)
but got " +
+ s"$groupedShuffles\n${grouped.queryExecution.executedPlan}")
+
+ // `UNION_OUTPUT_PARTITIONING=false` drops the pass-through so the
group-by adds its own
+ // shuffle; that freshly-planned path is the oracle for both the raw
union rows and the
+ // grouped result.
+ val (correctResult, correctGrouped) =
+ withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") {
+ val baseline = spark.sql(sqlText)
+ (baseline.collect(), baseline.groupBy($"id").count().collect())
+ }
+ checkAnswer(union, correctResult)
+ checkAnswer(grouped, correctGrouped)
+ }
+ }
+ }
+
+ test("SPARK-58819: union outputPartitioning ignores partition key
nullability") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ // `.sample(1.0)` keeps `FilterExec` above the shuffle without disabling
any rule:
+ // `PushPredicateThroughNonJoin.canPushThrough` has no `Sample` case, so
`IsNotNull` stays
+ // above the `Sample`, which passes its child's output and partitioning
through verbatim.
+ // The filter's output nullability is thus narrowed while its
passed-through partitioning
+ // keeps the nullable key. `fraction = 1.0` passes every row, keeping
counts deterministic.
+ val df1 = Seq((Option(1), 10L), (Option(2), 20L)).toDF("id", "v")
+ .repartition($"id").sample(withReplacement = false, fraction = 1.0,
seed = 42)
+ .filter($"id".isNotNull)
+ val df2 = Seq((Option(1), 30L), (Option(3), 40L)).toDF("id", "v")
+ .groupBy($"id").agg(sum($"v").as("v"))
+
+ val union = df1.union(df2)
+ val unionExec = union.queryExecution.executedPlan.collect { case u:
UnionExec => u }
+ assert(unionExec.size == 1)
+
+ // Pin the discriminator to nullability: child 0 is a `FilterExec` whose
partitioning
+ // (passed through verbatim) references the same column as its `output`,
differing only in
+ // nullability. The explicit qualifier check rules out the qualifier
cause the SQL-variant
+ // test above exercises; a future refactor leaking a qualifier
difference fails loudly
+ // instead of silently changing what the test covers.
+ val child0 = unionExec.head.children.head
+ val outputId = child0.output.head.asInstanceOf[AttributeReference]
+ val partitioningId =
+
child0.outputPartitioning.asInstanceOf[HashPartitioning].expressions.head
+ .asInstanceOf[AttributeReference]
+ assert(outputId.exprId == partitioningId.exprId,
+ s"expected the same column: $outputId vs $partitioningId")
+ assert(outputId.qualifier == partitioningId.qualifier,
+ s"qualifier must match (nullability is the sole difference): $outputId
vs $partitioningId")
+ assert(outputId.nullable != partitioningId.nullable,
+ s"expected a nullability difference: $outputId vs $partitioningId")
+ assert(outputId.withNullability(partitioningId.nullable) ==
partitioningId,
+ s"only nullability should differ: $outputId vs $partitioningId")
+
+ // child1 (groupBy) is the well-behaved branch: its output and
partitioning reference the
+ // same nullable group key, so the nullability mismatch is confined to
child0's Filter.
+ val child1 = unionExec.head.children(1)
+ val child1OutputId = child1.output.head.asInstanceOf[AttributeReference]
+ val child1PartitioningId =
+
child1.outputPartitioning.asInstanceOf[HashPartitioning].expressions.head
+ .asInstanceOf[AttributeReference]
+ assert(child1OutputId.nullable, s"group key should be nullable:
$child1OutputId")
+ assert(child1OutputId.exprId == child1PartitioningId.exprId,
+ s"child1 output and partitioning should reference the same column")
+ assert(child1OutputId.nullable == child1PartitioningId.nullable,
+ s"child1 has no nullability mismatch: $child1OutputId vs
$child1PartitioningId")
+
+ // The union propagates the co-located HashPartitioning in its own
output attribute space
+ // with the merged nullability (true, from child1's nullable group key).
This also pins that
+ // the partitioning attribute is the union's output attribute, not a
leaked child attribute.
+ assert(unionExec.head.outputPartitioning.isInstanceOf[HashPartitioning],
+ s"expected HashPartitioning pass-through but got
${unionExec.head.outputPartitioning}")
+ val unionPartitioning =
unionExec.head.outputPartitioning.asInstanceOf[HashPartitioning]
+ assert(unionPartitioning.expressions == Seq(unionExec.head.output.head))
+
+ // Without the fix the union reports UnknownPartitioning and the
downstream group-by adds a
+ // third shuffle; with it, the pass-through lets the group-by reuse the
two existing ones.
+ val grouped = union.groupBy($"id").count()
+ val groupedShuffles = grouped.queryExecution.executedPlan.collect {
+ case s: ShuffleExchangeExec => s
+ }.size
+ assert(groupedShuffles == 2,
+ s"group-by should reuse the union's partitioning (expect 2 shuffles)
but got " +
+ s"$groupedShuffles\n${grouped.queryExecution.executedPlan}")
+
+ // The two branches overlap on `id = 1`, so the correct count proves the
co-location claim
+ // is actually honored (a false claim would split id = 1 into duplicate
groups).
+ checkAnswer(grouped, Row(1, 2L) :: Row(2, 1L) :: Row(3, 1L) :: Nil)
+ }
+ }
+
test("SPARK-52921: union partitioning - range partitioning") {
val df1 = Seq((1, 2, 4), (1, 3, 5), (2, 2, 3), (2, 4, 5)).toDF("a", "b",
"c")
val df2 = Seq((4, 1, 5), (2, 4, 6), (1, 4, 2), (3, 5, 1)).toDF("d", "e",
"f")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
