ulysses-you commented on code in PR #58339:
URL: https://github.com/apache/spark/pull/58339#discussion_r3930476105
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -219,6 +236,21 @@ case class GroupPartitionsExec(
PartitionGrouping(partitions, isGrouped, isCollapsed)
}
+ /**
+ * Whether this node's grouping leaves every partition where it was and
keeps all of them, i.e.
+ * output partition i holds exactly input partition i and there is one
output per input. That is
+ * the only grouping that keeps a marked layout's undeclared rows at
hash(key) % numPartitions:
+ * a grouping that drops trailing declared keys still reads identity for
every group it keeps,
+ * but the partition count shrinks and the hash modulus with it. The
`forall` stops at the first
+ * moved partition, so a reorder or coalesce is rejected without a full scan.
+ */
+ @transient private lazy val identityGrouping: Boolean =
Review Comment:
@/tmp/r8-replies/f26.md
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -78,6 +78,18 @@ case class GroupPartitionsExec(
// data types match the reduced partition keys for the
identity-vs-transform and
// single-side-transform reducers; for the both-sides-reduce shape no
single transform
// describes the keys (see `KeyedShuffleSpec.reducersBothWays`).
+ //
+ // A marked claim pins undeclared rows to hash(key) % numPartitions
(see
+ // `KeyedPartitioning.mayContainUnknownPartitionKeys`). Only an
identity grouping keeps
+ // that relationship: any other grouping -- a reorder, a coalesce, a
resize, or the
+ // collapse a reduction applies -- moves those rows. Clearing only the
marker would
+ // misreport the undeclared rows that remain, so give up the keyed
partitioning at the
+ // physical output count (one per group, padding included) that a
parent's
+ // `PartitioningCollection` requires for uniformity.
`identityGrouping` is a lazy val, so
+ // repeated `outputPartitioning` calls scan it at most once.
+ if (PartitioningCollection.keyedMarkerOf(p).contains(true) &&
!identityGrouping) {
Review Comment:
@/tmp/r8-replies/recheck.md
##########
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 marked and unmarked keyed inputs meet; within a
collection the
+ * constructor makes that unrepresentable. Only `ShuffledJoin`'s `InnerLike`
arm can mix
+ * inputs this way; see the call site for the argument.
+ */
+ private def clearUnknownPartitionKeys(
+ partitionings: Seq[Partitioning]): Seq[Partitioning] = {
+ // One cached keyed member answers per input instead of re-flattening a
left-deep join
+ // chain, and keyless inputs drop out of the `flatMap` rather than reading
as unmarked. The
+ // all-unmarked path must reach no `copy`: `transform`'s `fastEquals`
would compare every
+ // partition key.
+ val markers = partitionings.flatMap(PartitioningCollection.keyedMarkerOf)
Review Comment:
@/tmp/r8-replies/keyless.md
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1142,29 +1205,33 @@ object PartitioningCollection {
* Note: this can't be implemented with `TreeNode.transform`.
*/
def fromPartitionings(partitionings: Seq[Partitioning]):
PartitioningCollection = {
- // See the class doc for why the flag is normalized by OR rather than
required to agree. One
- // representative per member is enough, because every collection agrees on
the flag internally
- // by this same construction, and only a member that disagrees is rebuilt.
+ // See the class doc for why the flags are normalized by OR rather than
required to agree. One
+ // representative per member is enough, because every collection agrees on
the flags
+ // internally by this same construction, and only a member that disagrees
is rebuilt.
val anyCollapsed =
partitionings.exists(representativeOf(_).exists(_.isCollapsed))
+ val anyUnknownKeys =
+
partitionings.exists(representativeOf(_).exists(_.mayContainUnknownPartitionKeys))
var canonicalKeys: Seq[InternalRowComparableWrapper] = null
// A partitioning with no `KeyedPartitioning` in it has nothing to
normalize, and one that
- // already agrees on both the keys and the flag is returned as it is. That
is what keeps
+ // already agrees on the keys and both flags is returned as it is. That is
what keeps
// repeated `outputPartitioning` computations over deeply nested
collections (e.g. chains of
// same-key joins) O(1) per level.
def intern(p: Partitioning): Partitioning = representativeOf(p) match {
case None => p
case Some(representative) =>
if (canonicalKeys == null) canonicalKeys = representative.partitionKeys
if ((representative.partitionKeys eq canonicalKeys) &&
- representative.isCollapsed == anyCollapsed) {
+ representative.isCollapsed == anyCollapsed &&
+ representative.mayContainUnknownPartitionKeys == anyUnknownKeys) {
Review Comment:
@/tmp/r8-replies/nested.md
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5933,6 +5946,945 @@ 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. SQL cannot
produce that shape:
Review Comment:
@/tmp/r8-replies/unreachable.md
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -578,12 +582,41 @@ 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
+ * same keys in the same order and using the
same partition
+ * function per position still pair (equal
undeclared keys hash
+ * to the same partition), but a row of an
undeclared key sits in
+ * the partition of some other declared key,
away from rows
Review Comment:
@/tmp/r8-replies/colocation.md
--
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]
