peter-toth commented on code in PR #58044:
URL: https://github.com/apache/spark/pull/58044#discussion_r3806064860
##########
sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala:
##########
@@ -1589,6 +1589,57 @@ class DataFrameSetOperationsSuite extends
SharedSparkSession with AdaptiveSparkP
}
}
+ test("SPARK-58819: union outputPartitioning ignores partition key
nullability") {
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ withTempView("t1", "t2") {
+ // `id` is nullable (Option[Int]) so that `IsNotNull` in the Filter
actually adjusts it.
+ // The two branches overlap on `id = 1`, so a shared key's grouping
result 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 union = spark.sql(
+ """
+ |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 unionExec = union.queryExecution.executedPlan.collect { case u:
UnionExec => u }
+ assert(unionExec.size == 1)
+
+ // `IsNotNull` in the Filter adjusts the nullability of the partition
key, but nullability
+ // does not affect the hash, so the union should still propagate the
hash partitioning.
+
assert(unionExec.head.outputPartitioning.isInstanceOf[HashPartitioning],
+ s"expected a HashPartitioning pass-through but got
${unionExec.head.outputPartitioning}")
+
+ // 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 path is the oracle for both the raw union rows and
the grouped result.
+ val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key
-> "false") {
Review Comment:
**Finding 1.** Neither of these blocks produces a baseline. `union` and
`grouped` were already planned above (`union.queryExecution.executedPlan` at
:1607, `grouped.queryExecution.executedPlan` at :1621) and `QueryExecution`
memoizes that — `lazyExecutedPlan` is a `LazyTry` held in a `val`
(`QueryExecution.scala:353`), and `Dataset.collect()` goes through
`withAction(...)` on the same `QueryExecution` instance
(`classic/Dataset.scala:1563`, `:2323`). So
`withSQLConf(UNION_OUTPUT_PARTITIONING -> "false")` never re-plans and both
`checkAnswer`s compare a plan against itself.
That makes the co-location guard a no-op: it stays green even if
`prepareOutputPartitioning`/`comparePartitioning` were to claim a co-location
the union RDD can't honor, which is exactly the regression the comment says it
catches.
I measured it with a probe inside the `correctGrouped` block:
```
PROBE stale-plan-shuffles=2 fresh-plan-shuffles=3
```
The "oracle" runs the 2-shuffle pass-through plan; only a freshly built
DataFrame gets the 3-shuffle baseline.
The sibling test at :1545 has the right shape — it builds the DataFrame
*inside* the conf block. Same fix here, hoisting the query text:
```scala
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 (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)
```
Or skip the config flip and assert the rows directly, which is a stronger
oracle and shorter (I ran the baseline to get these):
```scala
checkAnswer(union, Row(1, 10L) :: Row(2, 20L) :: Row(1, 30L) :: Row(3, 40L)
:: Nil)
checkAnswer(grouped, Row(1, 2L) :: Row(2, 1L) :: Row(3, 1L) :: Nil)
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -928,7 +924,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
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
+ // expressions have already been remapped to this union's output
attributes by
Review Comment:
**Finding 3.** The target of the remap got updated but not the subject: with
the first child remapped too, it isn't only "the other child's" expressions.
Suggest rewording 926-927 to:
```scala
// For `KeyedPartitioning`, only the partition expressions must match
(both sides'
// expressions have already been remapped to this union's output
attributes by
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -892,22 +892,19 @@ 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 its own input attributes (e.g. a
Filter passes through
+ // its child's partitioning but adjusts the output nullability), so even
the first child is
+ // remapped.
+ val attributesMap = children.map(_.output).map { childAttrs =>
+ AttributeMap(childAttrs.zip(output))
Review Comment:
**Finding 2.** `output` is inside the closure, so it runs once per child
instead of once. `UnionExec.output` is a plain `def` that rebuilds from
`children.map(_.output).transpose` (`:880`) and allocates a fresh
`AttributeReference` per column whenever a child widened the dataType, so this
turns an O(children x columns) step into O(children^2 x columns) — on a method
nothing memoizes and that `EnsureRequirements` plus every
`PartitioningPreservingUnaryExecNode` above the union calls repeatedly. Base
called `output` exactly once, in the `toUnionOutput` helper this PR removes, so
it's a regression rather than a pre-existing cost.
```scala
val unionOutput = output
val attributesMap = children.map(_.output).map { childAttrs =>
AttributeMap(childAttrs.zip(unionOutput))
}
```
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
