cloud-fan commented on code in PR #58339:
URL: https://github.com/apache/spark/pull/58339#discussion_r3913970103
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -578,12 +582,39 @@ case class CoalescedNullAwareHashPartitioning(
* partitioning this one was derived from onto the same
key, so one key here can
* stand for several of the original ones. Sticky. See "Key
Collapse" above for
* what it gates and how it travels.
+ * @param mayContainUnknownPartitionKeys Whether the data may contain rows
whose partition key is
+ * not among the declared `partitionKeys`.
`KeyGroupedPartitioner`
+ * routes such rows by a deterministic hash
when a side is
+ * re-shuffled onto this partitioning (see
+ * `KeyedShuffleSpec.createPartitioning`), so
co-location holds
+ * for whole keys only: two marked
partitionings declaring the
Review Comment:
**Nit (P3):** This compatibility statement is incomplete: equal declared
keys and order are not sufficient. The marker path in `areKeysCompatible` also
requires the corresponding partition expressions to use the same function; the
added bucket(4)-versus-bucket(8) test relies on that distinction. Please
qualify this sentence with the same-function requirement so the documented
contract matches the implementation.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala:
##########
@@ -80,6 +87,34 @@ trait ShuffledJoin extends JoinCodegenSupport {
s"ShuffledJoin should not take $x as the JoinType")
}
+ /**
+ * Clears the `mayContainUnknownPartitionKeys` marker of every
`KeyedPartitioning` in
+ * `partitionings` when at least one member is unmarked. Only
`ShuffledJoin`'s `InnerLike` arm
+ * can mix the two; see the call site for the argument. The interning
`fromPartitionings` does
+ * afterwards changes no marker.
Review Comment:
**Nit (P3):** This sentence is ungrammatical and makes the sequencing hard
to parse. For example: `The subsequent interning in fromPartitionings does not
change any markers.`
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala:
##########
@@ -80,6 +87,34 @@ trait ShuffledJoin extends JoinCodegenSupport {
s"ShuffledJoin should not take $x as the JoinType")
}
+ /**
+ * Clears the `mayContainUnknownPartitionKeys` marker of every
`KeyedPartitioning` in
+ * `partitionings` when at least one member is unmarked. Only
`ShuffledJoin`'s `InnerLike` arm
+ * can mix the two; see the call site for the argument. The interning
`fromPartitionings` does
+ * afterwards changes no marker.
+ */
+ private def clearUnknownPartitionKeys(
+ partitionings: Seq[Partitioning]): Seq[Partitioning] = {
+ // Only a marked member among an unmarked one needs the rebuild. On the
common all-unmarked
+ // path, even a no-op `copy` is not free: `transform`'s `fastEquals` would
fall through to
+ // comparing every partition key.
+ val marked = partitionings.view.flatMap(PartitioningCollection.flatten)
Review Comment:
**Non-blocking (P2):** This recursively materializes every keyed leaf merely
to read a marker that the `PartitioningCollection` invariant makes uniform. For
a left-deep chain of inner joins, each level re-flattens the nested child,
giving quadratic leaf visits in `outputPartitioning` even on the common
all-unmarked path. Please inspect one cached representative per direct input
and reserve the recursive transform for the genuinely mixed-marker case.
**Recommended change:** Expose an internal marker-summary or representative
helper backed by `PartitioningCollection.firstKeyedPartitioning` and use it for
the initial marked/all-marked decision.
**Why this works:** Read at most one cached representative from each direct
child partitioning; recursively traverse only when mixed markers require
rewriting marked descendants.
**Scope:** `PartitioningCollection`'s internal helper surface and
`ShuffledJoin.clearUnknownPartitionKeys`.
**Compatibility:** This changes only planning cost; marker normalization and
the mixed-marker rewrite remain behaviorally identical.
**Risks:** The summary must rely on the constructor invariant only after
every nested collection has validated uniform markers.
**Constraints:** Keep the all-marked and all-unmarked fast paths
allocation-free. Keep the recursive rewrite for a marked subtree beside an
unmarked sibling.
**Success:** Computing output partitioning for an N-deep all-unmarked join
chain visits O(N) direct members overall instead of repeatedly flattening every
accumulated leaf.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -983,13 +983,17 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
if (partitionings.forall(_.isInstanceOf[KeyedPartitioning])) {
val kps = partitionings.map(_.asInstanceOf[KeyedPartitioning])
val headKp = kps.head
- // The `KeyedPartitioning`s must agree on the partition expressions to
merge.
- val compatible = kps.forall(comparePartitioning(_, headKp))
+ // To merge, the `KeyedPartitioning`s must agree on the partition
expressions and no leg
+ // may carry the marker: the merged set declares the other legs' keys,
so an out-of-set row
+ // can sit in the wrong partition for the merged claim (rule (1) of the
+ // `KeyedPartitioning.mayContainUnknownPartitionKeys` doc). Unlike a
join, a union keeps
+ // every leg's rows, so an unmarked leg excuses a marked one nothing.
Review Comment:
**Nit (P3):** This clause is malformed. `An unmarked leg does not excuse a
marked one` states the intended rule directly.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -79,6 +79,23 @@ case class GroupPartitionsExec(
// single-side-transform reducers; for the both-sides-reduce shape no
single transform
// describes the keys (see `KeyedShuffleSpec.reducersBothWays`).
val partitionKeys = grouping.partitions.map(_._1)
+ // Members carry a uniform marker (mixed collections are cleared at
construction by
+ // `ShuffledJoin`), so any member answers for all of them.
+ val mayContainUnknownKeys = p.exists {
Review Comment:
**Nit (P3):** This recursive `exists` is only used by the fallback whose
other condition is `reducers.isDefined`, so the common `reducers = None` paths
scan the entire partitioning tree and immediately discard the result before the
following transform. Please put the cheap reducer check first and evaluate the
marker scan only on that branch.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5721,6 +5734,827 @@ class KeyGroupedPartitioningSuite
}
}
+ /**
+ * Asserts that the plan's shuffles (in tree order) are all
`KeyedPartitioning`s carrying the
+ * given `mayContainUnknownPartitionKeys` flags (a `KeyedPartitioning`
produced by
+ * `KeyedShuffleSpec.createPartitioning` always carries the marker).
+ */
+ private def assertShuffleMayContainUnknownPartitionKeys(
+ plan: SparkPlan,
+ expected: Seq[Boolean]): Unit = {
+ val shuffles = collectAllShuffles(plan)
+ assert(shuffles.size === expected.size,
+ s"expected ${expected.size} shuffles, got ${shuffles.size}:\n$plan")
+ shuffles.zip(expected).foreach { case (shuffle, hasUnknown) =>
+ shuffle.outputPartitioning match {
+ case k: KeyedPartitioning =>
+ assert(k.mayContainUnknownPartitionKeys === hasUnknown,
+ s"expected shuffle output
mayContainUnknownPartitionKeys=$hasUnknown, got " +
+ s"${k.mayContainUnknownPartitionKeys}:\n$plan")
+ case p =>
+ fail(s"expected a KeyedPartitioning shuffle, got $p:\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: one-side shuffle with out-of-set keys loses matches
in a following " +
+ "SPJ join") {
+ // a: keyed on id, keys {1, 2}. t: v1 parquet, keys {1, 2, 3}. u: keyed on
id, keys {1, 2, 3}.
+ // With shuffle.enabled, a RIGHT OUTER JOIN t shuffles t onto a's declared
keys {1, 2}; t's
+ // id=3 row is out-of-set, so the join output's partitioning has unknown
keys. A following
+ // storage-partitioned join against u must not trust it and falls back to
a shuffle.
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3,
'u3', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON
a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"))
+
+ // Baseline: no SPJ -> all three rows.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // Two one-side shuffles: t onto a's keys, then the first join's
output (unknown-keyed)
+ // onto u's keys. Both are keyed with unknown partition keys; neither
join GPEs.
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"second join must not storage-partition on an unknown-keyed layout,
got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: preserved non-keyed side of outer join falls back to
shuffle " +
+ "downstream") {
+ // Same hazard for every outer join type whose preserved side is the
non-keyed table: the
+ // one-side shuffle marks the preserved side's partitioning as having
unknown keys, so a
+ // downstream storage-partitioned join against a larger key set must fall
back to a shuffle.
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3,
'u3', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"))
+
+ // RIGHT OUTER preserves the non-keyed t on the right.
+ val rightQuery =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON
a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(rightQuery)
+ checkAnswer(df, expected)
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"downstream join must not storage-partition on an unknown-keyed
layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+
+ // FULL OUTER exposes UnknownPartitioning, so it is already safe
regardless of the shuffle
+ // direction; correctness is the guard.
+ val fullQuery =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a FULL OUTER JOIN t ON
a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(fullQuery)
+ checkAnswer(df, expected)
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ // The downstream join must not storage-partition on the first join's
unknown-keyed
+ // layout; FULL OUTER keeps it safe only because the join output
exposes
+ // UnknownPartitioning.
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"downstream join must not storage-partition on an unknown-keyed
layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+
+ // t LEFT OUTER JOIN a preserves the non-keyed t on the left.
+ val leftQuery =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM t LEFT OUTER JOIN testcat.ns.a a ON
t.id = a.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(leftQuery)
+ checkAnswer(df, expected)
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"downstream join must not storage-partition on an unknown-keyed
layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: keyed preserved side of outer join still uses the
one-side shuffle") {
+ // a (keyed) preserved on the left, t (non-keyed) nullable on the right: t
is shuffled onto
+ // a's keys (its partitioning is marked as having unknown keys), but the
LEFT OUTER join exposes
+ // only a's accurate partitioning, so the one-side shuffle stays sound and
the downstream SPJ
+ // still runs (no shuffle for the second join).
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3,
'u3', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT a.id AS id FROM testcat.ns.a a LEFT OUTER JOIN t ON
a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, "u1"), Row(2, "u2")))
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).nonEmpty,
+ s"downstream join should storage-partition on the accurate keyed
layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: one-side shuffle with out-of-set keys loses matches
in a following " +
+ "SPJ join (bucket)") {
+ // Same hazard as the identity variant, but the keyed sides are
partitioned by bucket(4, id):
+ // a covers buckets {0, 1, 2} (ids 0, 1, 2), while t holds id 3 (bucket
3), which a does
+ // not, so the one-side shuffle misplaces t's bucket-3 row while still
declaring a's layout.
+ // `id` is LONG because `BucketFunction` binds its value argument to
LongType.
+ val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ createTable("a", cols, Array(bucket(4, "id")))
+ createTable("u", cols, Array(bucket(4, "id")))
+ sql("INSERT INTO testcat.ns.a VALUES (0, 'a0'), (1, 'a1'), (2, 'a2')")
+ sql("INSERT INTO testcat.ns.u VALUES (0, 'u0'), (1, 'u1'), (2, 'u2'), (3,
'u3')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id BIGINT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (0, 't0'), (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON
a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(0L, "u0"), Row(1L, "u1"), Row(2L, "u2"), Row(3L,
"u3"))
+
+ // Baseline: no SPJ -> all four rows.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"second join must not storage-partition on an unknown-keyed layout,
got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: unknown-keyed partitioning still joins a
subset-keyed partner") {
+ // r (from a RIGHT OUTER JOIN t) has unknown partition keys {1, 2}, but
the downstream u is
+ // keyed on a subset {1}, so the storage-partitioned join stays compatible
and works: every
+ // key u can have is co-located on r's declared layout.
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON
a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, "u1")))
+ // Only the first join's one-side shuffle remains; the second join
storage-partitions.
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).nonEmpty,
+ s"subset-keyed partner should still storage-partition join, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: project dropping a key position drops the
unknown-keyed claim") {
+ // The first join's output is keyed on (id, k) and may contain unknown
keys (t's rows are all
+ // out-of-set: a holds k=x, t holds k=z). The Project below the second
join drops the k
+ // position, so the declared key set coarsens from {(1, x) ... (4, x)} to
{1, 2, 3, 4}, and
+ // an out-of-set (id, k) can then land inside the projected declared set.
The keyed claim must
+ // be dropped entirely, otherwise the second join trusts the coarsened
layout and silently
+ // loses the misplaced rows' matches.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", StringType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("id"), identity("k")))
+ createTable("u", cols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES " +
+ "(1, 'x', 'a1'), (2, 'x', 'a2'), (3, 'x', 'a3'), (4, 'x', 'a4')")
+ sql("INSERT INTO testcat.ns.u VALUES " +
+ "(1, NULL, 'u1'), (2, NULL, 'u2'), (3, NULL, 'u3'), (4, NULL, 'u4')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, k STRING, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 'z', 't1'), (2, 'z', 't2'), (3, 'z',
't3'), (4, 'z', 't4')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t
+ | ON a.id = t.id AND a.k = t.k) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"), Row(4,
"u4"))
+
+ // Baseline: no SPJ -> all four rows.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // The projection drops the unknown-keyed claim, so the second join
shuffles: two one-side
+ // shuffles, both keyed with unknown partition keys, and no
GroupPartitionsExec.
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"second join must not storage-partition on the coarsened layout,
got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: union of an unknown-keyed leg drops the merged keyed
partitioning") {
+ // The union's merged keys concatenate every leg's keys, a superset,
possibly equal, of each
+ // leg's declared set. When a leg may contain unknown partition keys,
another leg can declare
+ // exactly the key that leg holds out-of-set, so the merged claim would
promise co-location
+ // the marked leg cannot honor. Legs that declare the same key set would
in fact tolerate
+ // keeping the marker (its out-of-set row rides its own leg's partition
into that partition's
+ // group); this check does not try to tell that case from the rest, and
refuses.
+ createTable("a", columns, Array(identity("id")))
+ createTable("s", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.s VALUES (4, 's4', NULL), (5, 's5', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3,
'u3', NULL), " +
+ "(4, 'u4', NULL), (5, 'u5', NULL)")
+
+ // Disjoint-keyed second leg: the union's merged keys {1, 2, 4, 5} do not
cover t's out-of-set
+ // id=3, but the merged claim would still be a superset of the
unknown-keyed leg's
+ // declared keys.
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |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) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"), Row(4,
"u4"), Row(5, "u5"))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // The first join's one-side shuffle, then the union side re-shuffles
onto u's layout for
+ // the second join: the union exposes no keyed partitioning, so no
GroupPartitionsExec.
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"second join must not storage-partition on the merged layout, got:
" +
+ df.queryExecution.executedPlan)
+ }
+ }
+
+ // Overlapping second leg: s declares exactly the key {3} that the
unknown-keyed leg holds
+ // out-of-set, so the union's merged keys {1, 2, 3} equal u's keys and the
second join would
+ // storage-partition with no exchange, silently losing t's id=3 match.
+ sql("DROP TABLE IF EXISTS testcat.ns.s")
+ sql("DROP TABLE IF EXISTS testcat.ns.u")
+ createTable("s", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.s VALUES (3, 's3', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3,
'u3', NULL)")
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |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) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ // id=3 matches u once via t and once via s.
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"), Row(3,
"u3"))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"second join must not storage-partition on the merged layout, got:
" +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: join-key projection of an unknown-keyed layout drops
the claim") {
+ // Like the key-dropping-project repro, but the projection keeps both key
positions: the
+ // coarsening happens when the second join projects the declared keys down
to its join key
+ // (`id`) instead. A key that was out-of-set in the full key space lands
inside the projected
+ // declared set, so the unknown-keyed spec must be refused and the second
join must shuffle.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", StringType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("id"), identity("k")))
+ createTable("u", cols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES " +
+ "(1, 'x', 'a1'), (2, 'x', 'a2'), (3, 'x', 'a3'), (4, 'x', 'a4')")
+ sql("INSERT INTO testcat.ns.u VALUES " +
+ "(1, NULL, 'u1'), (2, NULL, 'u2'), (3, NULL, 'u3'), (4, NULL, 'u4')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, k STRING, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 'z', 't1'), (2, 'z', 't2'), (3, 'z',
't3'), (4, 'z', 't4')")
+
+ val query =
+ """
+ |SELECT r.id, r.k, u.data
+ |FROM (SELECT t.id AS id, t.k AS k FROM testcat.ns.a a RIGHT OUTER
JOIN t
+ | ON a.id = t.id AND a.k = t.k) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "z", "u1"), Row(2, "z", "u2"), Row(3, "z",
"u3"), Row(4, "z", "u4"))
+
+ // Baseline: no SPJ -> all four rows.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // The second join refuses the projected unknown-keyed spec and
shuffles instead: the
+ // first join's one-side shuffle plus the re-shuffle of the first
join's output.
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: spurious marker of an inner join keeps the reduced
SPJ") {
+ // The spurious marker on the inner join's collection is cleared at
construction (see
+ // `ShuffledJoin`), so a following reduced storage-partitioned join
(bucket(4) onto
+ // bucket(2)) must still work. Reading the marker with `exists` used to
make the
+ // GroupPartitionsExec give up with a zero-partition UnknownPartitioning,
which threw at
+ // planning when a parent asked for the partitioning.
+ val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ createTable("a", cols, Array(bucket(4, "id")))
+ createTable("u", cols, Array(bucket(2, "id")))
+ sql("INSERT INTO testcat.ns.a VALUES (0, 'a0'), (1, 'a1'), (2, 'a2'), (3,
'a3')")
+ sql("INSERT INTO testcat.ns.u VALUES (0, 'u0'), (1, 'u1'), (2, 'u2'), (3,
'u3')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id BIGINT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (0, 't0'), (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT a.id, u.data
+ |FROM testcat.ns.a a JOIN t ON a.id = t.id
+ |JOIN testcat.ns.u u ON a.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(0L, "u0"), Row(1L, "u1"), Row(2L, "u2"), Row(3L,
"u3"))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // The reduced SPJ still runs: the first join's one-side shuffle is
the only shuffle.
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
Seq(true))
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: inner join with in-set rows keeps the sound SPJ
downstream") {
+ // Same shape as the one-side-shuffle repro, but the inner join's second
side holds only
+ // in-set rows: the marker is spurious and cleared at construction (see
`ShuffledJoin`), so
+ // the following storage-partitioned join must not pay a shuffle for it.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", StringType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("id"), identity("k")))
+ createTable("u", cols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES " +
+ "(1, 'x', 'a1'), (2, 'x', 'a2'), (3, 'x', 'a3'), (4, 'x', 'a4')")
+ sql("INSERT INTO testcat.ns.u VALUES " +
+ "(1, NULL, 'u1'), (2, NULL, 'u2'), (3, NULL, 'u3'), (4, NULL, 'u4')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, k STRING, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 'x', 't1'), (2, 'x', 't2'), (3, 'x',
't3'), (4, 'x', 't4')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a JOIN t ON a.id = t.id
AND a.k = t.k) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"), Row(4,
"u4"))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // Only the first join's one-side shuffle: the second join
co-partitions on the
+ // spurious-marker-free layout without paying a shuffle.
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
Seq(true))
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: inner join clears the spurious marker on both member
orders") {
+ // `t`'s id=3 can match nothing on `a` (keys {1, 2}), so the inner join's
marker is spurious
+ // and cleared at construction (see `ShuffledJoin`). If it survived, the
subset gate in
+ // `areKeysCompatible` would refuse `u`'s wider key set and cost a shuffle
no sibling order
+ // can rescue. Both join orders must plan master's shape: the single
first-join shuffle
+ // plus a `GroupPartitionsExec` on each side of the storage-partitioned
second join.
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3,
'u3', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ // `t` first puts the marked member ahead of its unmarked sibling in the
collection;
+ // `a` first is the mirror order.
+ for (side <- Seq("t", "a")) {
+ val query = if (side == "t") {
+ """
+ |SELECT t.id, u.data
+ |FROM t JOIN testcat.ns.a a ON a.id = t.id
+ |JOIN testcat.ns.u u ON t.id = u.id
+ |""".stripMargin
+ } else {
+ """
+ |SELECT a.id, u.data
+ |FROM testcat.ns.a a JOIN t ON a.id = t.id
+ |JOIN testcat.ns.u u ON a.id = u.id
+ |""".stripMargin
+ }
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, "u1"), Row(2, "u2")))
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
Seq(true))
+
assert(collectGroupPartitions(df.queryExecution.executedPlan).nonEmpty,
+ s"side=$side: the second join should storage-partition, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: inner join marker clearing reaches nested
collections") {
+ // `ShuffledJoin`'s `InnerLike` arm passes each child's partitioning into
the joined
+ // collection as reported, so the next inner join can find marked members
nested inside a
+ // collection inherited from an all-marked inner join below. That nesting
cannot be planned
+ // through SQL (an inner join's own clearing already flattens what SQL
puts in it), so the
Review Comment:
**Nit (P3):** The clearing does not flatten the partitioning produced by
SQL. `clearUnknownPartitionKeys` uses `flatten` only to inspect the markers,
then returns the original hierarchy or transforms it in place; this test itself
later calls `PartitioningCollection.flatten` on `inner2.outputPartitioning`.
Please replace this rationale with the actual reason the all-marked nested
shape cannot be planned through SQL.
--
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]
