ulysses-you commented on code in PR #58044:
URL: https://github.com/apache/spark/pull/58044#discussion_r3820514522
##########
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:
Addressed: dropped `OPTIMIZER_EXCLUDED_RULES` and inserted
`.sample(withReplacement = false, fraction = 1.0, seed = 42)` between the
shuffle and the filter, so `FilterExec` stays above the shuffle with every rule
on (`canPushThrough` has no `Sample` case). Both SPARK-58819 tests still pass.
##########
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:
Addressed: pinned the qualifier discriminator with the same assert-based
checks as the nullability test -- same `exprId`, same `nullable`, differing
`qualifier`.
##########
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:
Addressed: named `dataType` as the one field a remap could not safely
normalize, and the `ExprId`-shared-`dataType` assumption behind it.
--
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]
