ulysses-you commented on code in PR #59535:
URL: https://github.com/apache/spark/pull/59535#discussion_r4332338844
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -727,7 +758,7 @@ case class KeyedPartitioning(
*/
override protected def stringArgs: Iterator[Any] =
Iterator(expressions, partitionKeys, isGrouped, isCollapsed,
- mayContainUnknownPartitionKeys) ++ ungroupingOrigin
+ unknownKeysOrigin) ++ ungroupingOrigin
Review Comment:
Flattened `unknownKeysOrigin` like `ungroupingOrigin` in e2c076e9748. Added
rendering tests for both bare origin names and the unmarked case. ExplainSuite
passes.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -937,9 +937,10 @@ case class EnsureRequirements(
// The pairing is only worth committing to if both children still declare
the same aligned key
// sequence once the grouping has been pushed into them. A
`GroupPartitionsExec` gives up its
- // keyed claim when it turns out to regroup a layout that pins undeclared
rows to
- // `hash(key) % numPartitions` (see
`KeyLayout.mayContainUnknownPartitionKeys`), and only the
- // node knows the permutation it performs, so that answer arrives after
the pairing was chosen.
+ // keyed claim when the grouping re-labels the keys or replicates them
into expected slots,
Review Comment:
Updated the comments in e2c076e9748 to distinguish plain regrouping, which
retains clustering and loses routing, from key rewriting and non-identity
expected-key alignment, which drop the claim. Also replaced the stale
`keyedMarkerOf` reference with `representativeOf`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -447,6 +447,21 @@ case object REPLICATED_FOR_JOIN extends UngroupingOrigin
// A spread an operator's ordering requirement is read through, which pairs
with nothing.
case object SPLIT_FOR_ORDERING extends UngroupingOrigin
+/**
+ * Why a layout may hold rows whose key it does not declare, and what became
of the hash routing
+ * that placed them. `KeyGroupedPartitioner` routes such a row by
+ * `nonNegativeMod(key.hashCode, numPartitions)`, so a marked layout
co-locates whole keys, never a
+ * subset of their columns. Anything that regrouped, sorted or projected the
keys moves the rows off
+ * that routing, and only the co-location survives.
+ */
+sealed trait UnknownKeysOrigin
+
+// An undeclared row sits at `nonNegativeMod(key.hashCode, numPartitions)`;
the relation holds.
+case object HASH_ROUTING extends UnknownKeysOrigin
+
+// A regrouping, a sort or a projection moved the rows off that relation; only
co-location remains.
Review Comment:
Corrected the routing-loss producer list and the union carry rule in
e2c076e9748. The projection comment now explains the undeclared-row versus
partition-label mismatch, rather than declared-key collapse.
##########
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:
Added both SQL regressions in e2c076e9748 with AQE on/off and AQE-off
shuffle counts of 3 and 4. Verified discrimination: replacing the ordered-list
check with set equality produces 76 groups instead of 40; removing the
marked/marked routing guard produces 4 groups instead of 40. Updated the
hand-built-test comment. Also removed the session-log paragraph from the PR
testing section and updated the coverage and test results.
--
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]
