dongjoon-hyun commented on code in PR #57491:
URL: https://github.com/apache/spark/pull/57491#discussion_r3646773557
##########
sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala:
##########
@@ -1659,6 +1659,187 @@ class DataFrameSetOperationsSuite extends
SharedSparkSession with AdaptiveSparkP
}
}
+ test("SPARK-58317: union partitioning - PartitioningCollection child
intersects to single") {
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.PREFER_SORTMERGEJOIN.key -> "false") {
+ withTempView("t1", "t2", "t3", "t4") {
+ Seq((1, 2, 4), (1, 3, 5), (2, 2, 3)).toDF("c1", "c2",
"c3").createOrReplaceTempView("t1")
+ Seq((1, 9), (2, 9)).toDF("c1", "x").createOrReplaceTempView("t2")
+ Seq((1, 2, 4), (2, 4, 5), (3, 6, 7)).toDF("c1", "c2",
"c3").createOrReplaceTempView("t3")
+ Seq((1, 9), (3, 9)).toDF("c1", "y").createOrReplaceTempView("t4")
+
+ // The first branch is an inner shuffled-hash join and selects both
join keys (t1.c1 and
+ // t2.c1), so its output partitioning is a
PartitioningCollection(Hash(c1), Hash(c1#..))
+ // that a downstream ProjectExec cannot narrow to a single member. The
second branch is a
+ // left join, whose output partitioning is a single
HashPartitioning(c1). The union should
+ // intersect the two to a single HashPartitioning(c1) and let the
group-by skip a shuffle.
+ def unionDF: DataFrame = sql(
+ """SELECT c1, c2, c3, count(*) FROM (
+ | SELECT /*+ SHUFFLE_HASH(t2) */ t1.c1, t1.c2, t1.c3, t2.c1 AS k
+ | FROM t1 JOIN t2 ON t1.c1 = t2.c1
+ | UNION ALL
+ | SELECT /*+ SHUFFLE_HASH(t4) */ t3.c1, t3.c2, t3.c3, t3.c1 AS k
+ | FROM t3 LEFT JOIN t4 ON t3.c1 = t4.c1
+ |) GROUP BY c1, c2, c3
Review Comment:
Since `k` is unused above the union, `ColumnPruning` drops it from both
branches and each `ProjectExec` narrows the join's collection to a single
`HashPartitioning(c1)`, so no `PartitioningCollection` reaches the union — this
test passes on master without this PR. Referencing `k` in the group-by keeps
the first branch's collection alive; aliasing `t4.c1` instead of `t3.c1` keeps
the second branch a single `Hash(c1)`, so the intersection is single as
intended.
```suggestion
| SELECT /*+ SHUFFLE_HASH(t4) */ t3.c1, t3.c2, t3.c3, t4.c1 AS k
| FROM t3 LEFT JOIN t4 ON t3.c1 = t4.c1
|) GROUP BY c1, c2, c3, k
```
--
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]
