peter-toth commented on code in PR #59535:
URL: https://github.com/apache/spark/pull/59535#discussion_r4327547405
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -9351,6 +9358,313 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-60277: SPJ: a union merges a marked leg that declares what the
other legs " +
+ "declare") {
+ // The marked leg is `a RIGHT OUTER JOIN t`: t reports no partitioning, so
its side is
+ // shuffled onto a's declared keys {1, 2, 3}, and the right-outer arm
exposes that marked
+ // layout, holding t's id=4 rows out-of-set. The second leg `s` declares
{1, 2}, a subset, so
+ // the keys merge and the marker travels onto them; the out-of-set rows
ride one partition of
+ // the marked leg, so they stay together in its group, which is what the
counts assert.
+ createTable("a", columns, Array(identity("id")))
+ createTable("s", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL), (3,
'a3', NULL)")
+ sql("INSERT INTO testcat.ns.s VALUES (1, 's1', NULL), (2, 's2', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (4, 't4'), (4, 't4b')")
+
+ val query =
+ """
+ |SELECT id, count(*) AS cnt
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON
a.id = t.id
+ | UNION ALL
+ | SELECT id FROM testcat.ns.s)
+ |GROUP BY id
+ |""".stripMargin
+ val expected = Seq(Row(1, 2), Row(2, 2), Row(4, 2))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ val plan = df.queryExecution.executedPlan
+ val union = collect(plan) { case u: UnionExec => u }.head
+ val merged = union.outputPartitioning match {
+ case kp: physical.KeyedPartitioning => kp
+ case other => fail(s"the union should merge the legs' keys, got
$other")
+ }
+ assert(merged.mayContainUnknownPartitionKeys,
+ s"the merged layout must carry the marked leg's marker, got $merged")
+ assert(!merged.unknownKeysRouted,
+ s"an ungrouped concatenation gives the routing up, got $merged")
+ // The regrouping leaves every key's rows in one group, so the plan
satisfies the
+ // aggregate's distribution; that is what keeps AQE's rules from
discarding it.
+ assert(ValidateRequirements.validate(plan),
+ s"the plan must satisfy its requirements:\n$plan")
+ assert(collectAllShuffles(plan).size == 1,
+ s"the join's one-side shuffle should be the only exchange:\n$plan")
+ }
+
+ // Negative control: `s` declares the very key the marked leg holds
out-of-set, so the
+ // merged claim would promise co-location for id=4 that the marked leg
cannot honor. The
+ // keys must not merge, and the aggregate pays its own exchange.
+ sql("DROP TABLE IF EXISTS testcat.ns.s")
+ createTable("s", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.s VALUES (4, 's4', NULL)")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 1), Row(2, 1), Row(4, 3)))
+ val plan = df.queryExecution.executedPlan
+ val union = collect(plan) { case u: UnionExec => u }.head
+
assert(union.outputPartitioning.isInstanceOf[physical.UnknownPartitioning],
+ s"a leg declaring an out-of-set key must stop the merge, got " +
+ union.outputPartitioning)
+ assert(collectAllShuffles(plan).size == 2,
+ s"the one-side shuffle plus the aggregate's own exchange:\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-60277: SPJ: a join above a merged marked union pairs when the
keys line up") {
+ // The merged layout is marked and ungrouped, which `resolveChild` admits
and wraps in a
+ // `GroupPartitionsExec`. Coalescing the partitions that share a key keeps
every key's rows in
+ // the group named for it, so the node reports a keyed layout whose
routing is lost, and the
+ // join pairs it with the other side's scan instead of re-shuffling it.
The marked leg's
+ // out-of-set id=4 rows sit in a declared key's group, and the scan holds
no id=4 row, so they
+ // match nothing.
+ createTable("ja", columns, Array(identity("id")))
+ createTable("js", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.ja VALUES (1, 'a1', NULL), (2, 'a2', NULL),
(3, 'a3', NULL)")
+ sql("INSERT INTO testcat.ns.js VALUES (1, 's1', NULL), (2, 's2', NULL)")
+
+ withTable("jt") {
+ sql("CREATE TABLE jt (id INT, data STRING) USING parquet")
+ sql("INSERT INTO jt VALUES (1, 't1'), (2, 't2'), (4, 't4'), (4, 't4b')")
+
+ val query =
+ """
+ |SELECT u.id, count(*) AS cnt FROM (
+ | SELECT jt.id AS id FROM testcat.ns.ja a RIGHT OUTER JOIN jt ON
a.id = jt.id
+ | UNION ALL
+ | SELECT id FROM testcat.ns.js) u
+ |JOIN testcat.ns.ja a2 ON u.id = a2.id
+ |GROUP BY u.id
+ |""".stripMargin
+
+ Seq(false, true).foreach { aqe =>
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 2), Row(2, 2)))
+ if (!aqe) {
+ // The merge is what puts the union's layout in front of the
join's grouping; under
+ // AQE the plan is rewritten around the stages, so the counts
above carry that arm.
+ val plan = df.queryExecution.executedPlan
+ val union = collect(plan) { case u: UnionExec => u }.head
+ union.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+ assert(kp.mayContainUnknownPartitionKeys,
+ s"the merged layout must carry the marked leg's state, got
$kp")
+ case other =>
+ fail(s"the union should merge the legs' keys, got $other")
+ }
+ val grouped = collectAllGroupPartitions(plan)
+ assert(grouped.size == 1, s"the join pairs through the
grouping:\n$plan")
+ grouped.head.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+ assert(kp.isGrouped && kp.mayContainUnknownPartitionKeys &&
+ !kp.unknownKeysRouted,
+ s"the grouping keeps the claim and loses the routing, got
$kp")
+ case other =>
+ fail(s"the grouping should keep the keyed claim, got $other")
+ }
+ assert(collectAllShuffles(plan).size == 1,
+ s"the join pairs instead of re-shuffling the union side:\n$plan")
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-60277: a window above a merged marked union keeps its sort") {
+ // The grouping coalesces the partitions of one declared key, and one of
them may also hold an
+ // undeclared key's rows, so the key ordering is not constant across the
merged partition. The
+ // window needs the sort its required ordering asks for.
+ createTable("wa", columns, Array(identity("id")))
+ createTable("ws", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.wa VALUES (1, 'a1', NULL), (2, 'a2', NULL),
(3, 'a3', NULL)")
+ sql("INSERT INTO testcat.ns.ws VALUES (1, 's1', NULL), (2, 's2', NULL),
(3, 's3', NULL)")
+
+ withTable("wt") {
+ sql("CREATE TABLE wt (id INT, data STRING) USING parquet")
+ sql("INSERT INTO wt VALUES (1, 't1'), (2, 't2'), (3, 't3'), (4, 't4')")
+
+ val union =
+ """
+ |SELECT wt.id AS id FROM testcat.ns.wa a RIGHT OUTER JOIN wt ON a.id
= wt.id
+ |UNION ALL
+ |SELECT id FROM testcat.ns.ws
+ |""".stripMargin
+ val expected = Seq(
+ Row(1, 2), Row(1, 2), Row(2, 2), Row(2, 2), Row(3, 2), Row(3, 2),
Row(4, 1))
+
+ Seq(false, true).foreach { preserveKeyOrdering =>
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_PRESERVE_KEY_ORDERING_ON_COALESCE_ENABLED.key
->
+ preserveKeyOrdering.toString) {
+ val df = sql(union).sortWithinPartitions("id")
+ .selectExpr("id", "count(1) OVER (PARTITION BY id) AS cnt")
+ val mismatch = QueryTest.sameRows(expected, df.collect().toSeq)
+ assert(mismatch.isEmpty, mismatch.getOrElse(""))
+ val plan = df.queryExecution.executedPlan
+ val grouping = collectAllGroupPartitions(plan)
+ assert(grouping.size == 1, s"the union should be grouped
once:\n$plan")
+ assert(grouping.head.groupedPartitions.exists(_._2.size > 1))
+
assert(grouping.head.outputPartitioning.asInstanceOf[physical.KeyedPartitioning]
+ .mayContainUnknownPartitionKeys)
+ assert(grouping.head.outputOrdering.isEmpty)
+ assert(collect(plan) {
+ case s: SortExec if unwrapWrappers(s.child) == grouping.head => s
+ }.size == 1, s"the window needs a sort above the grouping:\n$plan")
+ assert(ValidateRequirements.validate(plan))
+ }
+ }
+ }
+ }
+
+ test("SPARK-60277: SPJ: two marked legs route the keys they do not declare
alike") {
+ // Both legs are marked and declare {1, 2, 3}, so each is shuffled onto
those keys and each
+ // right-outer arm exposes the marked layout. id=4 is out-of-set on both,
and both route it by
+ // the same hash over a count the shared declaration fixes, so its rows
are counted once; a
+ // different routing would split them, which the count asserts.
+ createTable("a", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL), (3,
'a3', NULL)")
+
+ withTable("t1", "t2") {
+ sql("CREATE TABLE t1 (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t1 VALUES (1, 'x1'), (4, 'x4'), (4, 'x4b')")
+ sql("CREATE TABLE t2 (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t2 VALUES (2, 'y2'), (4, 'y4'), (5, 'y5')")
+
+ val query =
+ """
+ |SELECT id, count(*) AS cnt
+ |FROM (SELECT t1.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t1 ON
a.id = t1.id
+ | UNION ALL
+ | SELECT t2.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t2 ON
a.id = t2.id)
+ |GROUP BY id
+ |""".stripMargin
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 1), Row(2, 1), Row(4, 3), Row(5, 1)))
+ val plan = df.queryExecution.executedPlan
+ val union = collect(plan) { case u: UnionExec => u }.head
+ val merged = union.outputPartitioning match {
+ case kp: physical.KeyedPartitioning => kp
+ case other => fail(s"the union should merge the legs' keys, got
$other")
+ }
+ assert(merged.mayContainUnknownPartitionKeys,
+ s"the merged layout must carry the legs' marker, got $merged")
+ assert(collectAllShuffles(plan).size == 2,
+ s"one one-side shuffle per leg should be the whole exchange
count:\n$plan")
+ assert(collectAllGroupPartitions(plan).nonEmpty,
+ s"the aggregate should group the union's partitions:\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-60277: marked legs must declare their keys in one order to be
merged") {
+ // Both legs send a key K no leg declares to `nonNegativeMod(K.hashCode,
3)`, so legs
+ // declaring [1, 2, 3] and [2, 3, 1] send it to a different named group
and K's rows split.
+ // The gate compares the key lists, not the key sets; a differing
multiplicity is refused for
+ // the same reason. Hand-built, since a planner lays both legs out on one
template.
Review Comment:
**Finding 1.** "Hand-built, since a planner lays both legs out on one
template" doesn't hold. A grouped union used as the one-side-shuffle template
gives a marked leg an unsorted key order. These two queries reach the guards
from SQL:
```scala
// testcat tables, identity(id): a1 {3, 4}, a2 {1, 2}, a3 {1, 2, 3, 4}, s1
{1, 2}, s2 {1, 2}
// parquet t1 and t2: ids 1 to 40
// V2_BUCKETING_SHUFFLE_ENABLED and UNION_OUTPUT_PARTITIONING on
val leg1 = "SELECT t1.id AS id FROM (SELECT id FROM testcat.ns.a1 UNION ALL
" +
"SELECT id FROM testcat.ns.a2) tmpl RIGHT OUTER JOIN t1 ON tmpl.id = t1.id"
val leg2 = "SELECT t2.id AS id FROM testcat.ns.a3 a RIGHT OUTER JOIN t2 ON
a.id = t2.id"
// (1) Marked legs over [3, 4, 1, 2] and [1, 2, 3, 4]: the key-order gate in
the union.
sql(s"SELECT id, count(*) FROM ($leg1 UNION ALL $leg2) GROUP BY id")
// (2) Two merged unions, both regrouped to [1, 2, 3, 4] with the routing
lost: the routing
// check in `areKeysCompatible`.
val u1 = s"$leg1 UNION ALL SELECT id FROM testcat.ns.s1"
val u2 = s"$leg2 UNION ALL SELECT id FROM testcat.ns.s2"
sql(s"SELECT x.id, count(*) FROM ($u1) x JOIN ($u2) y ON x.id = y.id GROUP
BY x.id")
```
Both return the 40 expected rows at the head, with AQE off and on. With a
guard weakened they don't:
- Comparing the marked legs' key sets instead of their lists
(`basicPhysicalOperators.scala:1062`): (1) returns 76 rows, because every
undeclared id splits into two groups.
- Dropping `unknownKeysRouted` from the marked/marked arm of
`areKeysCompatible` (`partitioning.scala:2324`): (2) returns 4 rows, because
every undeclared-key match is lost.
Today these guards are pinned by this test and the
`GroupPartitionsExecSuite` unit case. Adding both queries as tests, with their
shuffle counts (3 and 4 at the head), would pin them through the planner too.
The comment above could then say what the hand-built cases add. This
complements the pairing-sweep point at
[r4324836938](https://github.com/apache/spark/pull/59535#discussion_r4324836938).
--
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]
