peter-toth commented on code in PR #58044:
URL: https://github.com/apache/spark/pull/58044#discussion_r3819934258
##########
sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala:
##########
@@ -1589,6 +1590,141 @@ 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`,
Review Comment:
**Finding 8.** The nullability test pins its discriminator four ways — same
`exprId`, same `qualifier`, `nullable` differs, `withNullability` round-trips —
so a change that shifts the cause fails loudly there. This test states its
discriminator only in prose, and it is the more fragile of the two: the drift
comes from the auto-generated derived-table alias re-qualifying `id`, i.e. from
name resolution, not from an operator's own contract. If that ever stops
happening, child 0's output attribute and its partitioning attribute become
`==`, base propagates `hashpartitioning(id#9, 5)` on its own, and this test
goes green on unfixed code with nothing to flag it.
Same shape as the sibling test:
```scala
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): $outputId vs
$partitioningId")
assert(outputId.qualifier != partitioningId.qualifier,
s"expected a qualifier difference: $outputId vs $partitioningId")
```
Measured on this head: `out qual=List(__auto_generated_subquery_name)` vs
`part qual=List(t1)`, `nullable=true` on both sides.
##########
sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala:
##########
@@ -1589,6 +1590,141 @@ 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)
+
+ // 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",
+ // `PushDownPredicates` would otherwise push `IsNotNull` below the
shuffle, turning child
+ // 0 into a `ShuffleExchangeExec` whose partitioning and output carry
the same nullability.
+ // Disabling it keeps `FilterExec` directly above the shuffle, so its
output (nullability
+ // narrowed) differs from its passed-through partitioning only in
nullability.
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.catalyst.optimizer.PushDownPredicates") {
Review Comment:
**Finding 7.** The comment is right that `PushPredicateThroughNonJoin`
relocates `IsNotNull` below the shuffle *for the shape written here*, but the
drift itself does not need the rule disabled — `canPushThrough`
(`Optimizer.scala:2393`) has no `Sample` case, so a filter above a `Sample`
stays put with the full optimizer running.
Measured, dropping the excluded-rules entry and inserting one `.sample(...)`:
```scala
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)
```
- head: `child0 out(null=false, qual=List()) part(null=true, qual=List())
sameExprId=true onlyNullabilityDiffers=true`,
`union.outputPartitioning=hashpartitioning(id#9, 2)`, 2 shuffles, rows `[1,2]
[2,1] [3,1]`.
- base (`5460c109699`): same drift, `UnknownPartitioning(0)`, 3 shuffles.
`fraction = 1.0` passes every row, so the counts and `checkAnswer` stay
deterministic. That buys two things: the test stops asserting on a plan shape
the optimizer would never hand you, and the "not reachable through the normal
SQL path" note in the PR thread can go — the shape is reachable, it just needs
a node `canPushThrough` doesn't cover instead of a config.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -892,22 +892,23 @@ 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 (name, nullability,
metadata, qualifier): a
Review Comment:
**Finding 9.** `AttributeReference.equals` (`namedExpressions.scala:301`)
compares `name`, `dataType`, `nullable`, `metadata`, `exprId` and `qualifier`.
The parenthetical drops `dataType`, which is the one entry that isn't cosmetic:
name, nullability, metadata and qualifier can all differ without changing what
the partitioning means, because `HashPartitioning.partitionIdExpression` hashes
values (`Pmod(CollationAwareMurmur3Hash(expressions), n)`,
`partitioning.scala:328`), while remapping an `int` attribute onto a `bigint`
one would turn a true co-location claim into a false one. The remap is safe
because two attributes sharing an `exprId` agree on `dataType` in a resolved
plan — that is the assumption worth writing down rather than leaving off the
list.
```scala
// `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.
```
--
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]
