ulysses-you commented on code in PR #58044:
URL: https://github.com/apache/spark/pull/58044#discussion_r3809358402
##########
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:
Fixed. The oracle now re-plans inside the `withSQLConf` block (`val baseline
= spark.sql(sqlText)`), so `checkAnswer` compares against a fresh 3-shuffle
baseline rather than the memoized plan.
##########
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:
Fixed. Hoisted `output` into a `val unionOutput` outside the per-child
closure.
##########
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:
Fixed. Reworded the comment to "both sides' expressions".
--
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]
