You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "cloud-fan (via GitHub)" <gi...@apache.org> on 2026/09/01 13:32:50 UTC
Re: [PR] [SPARK-59050][SQL] SPJ one-side shuffle with out-of-set keys produces wrong results in multi-joins [spark]
cloud-fan commented on code in PR #58339:
URL: https://github.com/apache/spark/pull/58339#discussion_r3904614801
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -573,12 +577,29 @@ 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 to arbitrary partitions when a side is
+ * re-shuffled onto this partitioning (see
+ * `KeyedShuffleSpec.createPartitioning`), so only the declared
Review Comment:
**Nit (P3):** This is true for one marked side considered against an unrelated layout, but the implementation also deliberately accepts two marked sides with the same ordered declared layout: equal unknown keys use the same deterministic hash fallback and are co-located. Please document that exception, since `areKeysCompatible` relies on it to omit a shuffle.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5463,6 +5476,675 @@ 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 = collect(plan) { case s: ShuffleExchangeExec => s }
+ 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 are the concatenation of every leg's keys -- a superset of each
+ // leg's declared set. When a leg's partitioning may contain unknown partition keys, another
+ // leg can declare exactly the key the unknown-keyed leg holds out-of-set, so the merged
+ // claim would promise co-location the unknown-keyed leg cannot honor. The union must drop
+ // the keyed partitioning entirely (with multiple legs the merged set is always larger than
+ // any single leg's, so there is no sound way to keep it).
+ 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 inner join's PartitioningCollection holds the keyed side's own (unmarked) KP next to
+ // the re-shuffled side's unknown-keyed one. The marker is spurious -- the inner join
+ // dropped every out-of-set row -- so a following reduced storage-partitioned join (bucket(4)
+ // reduced onto bucket(2)) must still work. Reading the marker off the collection 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 join's PartitioningCollection carries the keyed side's own (unmarked) KP
+ // next to the re-shuffled side's unknown-keyed one, and the join has already dropped every
+ // out-of-set row, so the marker is spurious and 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") {
+ // The inner join's output collection pairs the re-shuffled side's unknown-keyed KP with
+ // the keyed side's own unmarked KP. The join already dropped every out-of-set row, so the
+ // marker is spurious and must be cleared when the collection is built -- the downstream join
Review Comment:
**Nit (P3):** The inner join's `PartitioningCollection` still contains the marker values `[true, false]`; construction does not clear either marker. The downstream projection/consumer interprets the mixed collection and derives an unmarked partitioning. Please attribute the transition to that consumer so this comment does not send future fixes to `ShuffledJoin`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5463,6 +5476,675 @@ 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 = collect(plan) { case s: ShuffleExchangeExec => s }
+ 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 are the concatenation of every leg's keys -- a superset of each
+ // leg's declared set. When a leg's partitioning may contain unknown partition keys, another
+ // leg can declare exactly the key the unknown-keyed leg holds out-of-set, so the merged
+ // claim would promise co-location the unknown-keyed leg cannot honor. The union must drop
+ // the keyed partitioning entirely (with multiple legs the merged set is always larger than
+ // any single leg's, so there is no sound way to keep it).
+ 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 inner join's PartitioningCollection holds the keyed side's own (unmarked) KP next to
+ // the re-shuffled side's unknown-keyed one. The marker is spurious -- the inner join
+ // dropped every out-of-set row -- so a following reduced storage-partitioned join (bucket(4)
+ // reduced onto bucket(2)) must still work. Reading the marker off the collection 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 join's PartitioningCollection carries the keyed side's own (unmarked) KP
+ // next to the re-shuffled side's unknown-keyed one, and the join has already dropped every
+ // out-of-set row, so the marker is spurious and 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") {
+ // The inner join's output collection pairs the re-shuffled side's unknown-keyed KP with
+ // the keyed side's own unmarked KP. The join already dropped every out-of-set row, so the
+ // marker is spurious and must be cleared when the collection is built -- the downstream join
+ // matches the join key against the unknown-keyed member's expressions (the re-shuffled side
+ // is the join key's attribute origin), so a sibling-order preference in
+ // `createKeyedShuffleSpec` could not rescue the SPJ here. Both join orders must plan the
+ // single first-join 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, 'x', 't1'), (2, 'x', 't2'), (3, 'x', 't3'), (4, 'x', 't4')")
+ val expected = Seq(Row(1, "x", "u1"), Row(2, "x", "u2"), Row(3, "x", "u3"), Row(4, "x", "u4"))
+
+ // `t` first: the collection lists the unknown-keyed member before its unmarked sibling;
+ // `a` first: the mirror order (the negative control pinning order-insensitivity).
+ for (side <- Seq("t", "a")) {
+ val query = if (side == "t") {
+ """
+ |SELECT t.id, a.k, u.data
+ |FROM t JOIN testcat.ns.a a ON a.id = t.id AND a.k = t.k
+ |JOIN testcat.ns.u u ON t.id = u.id
+ |""".stripMargin
+ } else {
+ """
+ |SELECT a.id, a.k, u.data
+ |FROM testcat.ns.a a JOIN t ON a.id = t.id AND a.k = t.k
+ |JOIN testcat.ns.u u ON a.id = u.id
+ |""".stripMargin
+ }
+ 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)
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan, Seq(true))
+ }
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: inner join drops out-of-set rows before the cleared marker is trusted") {
+ // Data-level companion to the spurious-marker tests: `t` here genuinely holds an out-of-set
+ // row (5, 'z') that a's declared keys cannot match. The inner join drops it, so every row
+ // reaching the cleared-marker layout carries a declared key -- the downstream storage-
+ // partitioned join must then return exactly the matched rows: ids 1..4, and NOT id=5 even
+ // though u holds it. If the marker were cleared for a shape that keeps out-of-set rows,
+ // this query would silently lose or misplace 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'), (5, NULL, 'u5')")
+
+ 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'), (5, 'z', 't5')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM t JOIN testcat.ns.a a ON a.id = t.id AND a.k = t.k) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ // id=5 dropped by the inner join (a has no (5, *)); u's id=5 row has no partner.
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"), Row(4, "u4"))
+
+ // Baseline without SPJ.
+ 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 marker cleared, so the second join storage-partitions: only the first join's
+ // one-side shuffle remains.
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan, Seq(true))
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: genuine unknown keys survive an inner join on a key subset") {
+ // Adversarial shape against the spurious-marker clearing: the RIGHT OUTER output `r` is a
+ // single unknown-keyed KP carrying a GENUINE unknown row (1, 'zz'); it inner-joins `x`
+ // (a keyed table declaring the same key set) on a strict subset of the key columns; `u`
+ // above matches on the full (id, k). The (1,'zz','xa','uz') match must survive: clearing
Review Comment:
**Nit (P3):** `u` is created below this comment and joined later in the test, so "above" points at the wrong source location. Please use "below", or simply say that `u` matches on the full `(id, k)` key.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1678,7 +1751,18 @@ case class KeyedShuffleSpec(
// The shuffled side is laid out on this side's partition keys, so it inherits the flag. That
// is conservative rather than strictly true, and it can only ever add a shuffle: a later
// grouping of the shared key set carries the collapsed side's risk.
- partitioning.copy(expressions = newExpressions)
+ //
+ // The child re-shuffled onto this layout may hold keys outside the declared set, so every
+ // partitioning produced here carries the unknown-keys marker too (see
+ // `KeyedPartitioning.mayContainUnknownPartitionKeys`). `createPartitioning` is only reached
+ // from `EnsureRequirements`' shuffle loop to re-shuffle a join child onto the best spec's
+ // layout, and the re-shuffled child's keys are never provably a subset of the declared keys:
Review Comment:
**Nit (P3):** The keys are not inherently unknowable: for example, identity keys `[1]` are visibly a subset of identity keys `[1, 2]`. The current selection path simply does not carry or attempt a same-domain subset proof. Please describe this as conservative implementation behavior so the comment does not rule out a valid future precision improvement.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -754,15 +775,23 @@ case class KeyedPartitioning(
override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = {
val result = KeyedShuffleSpec(this, distribution)
if (SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys) {
+ val joinKeyPositions = result.keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2)
+ // The projection coarsens the declared key set, which an unknown-keyed claim cannot
+ // survive (see `mayContainUnknownPartitionKeys`). Return the unprojected spec instead: it
+ // disagrees on arity with every other spec, so co-location is refused and the child falls
Review Comment:
**Nit (P3):** Two marked `(id, k)` specs narrowed to `id` retain the same two-expression arity, so arity disagreement is not the general safeguard here. They are rejected because the omitted `k` position has an empty key-position intersection, and `canCreatePartitioning` likewise requires every position to map. Please cite those checks instead.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -573,12 +577,29 @@ 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 to arbitrary partitions when a side is
+ * re-shuffled onto this partitioning (see
+ * `KeyedShuffleSpec.createPartitioning`), so only the declared
+ * keys are guaranteed to be co-located;
+ * `KeyedShuffleSpec.areKeysCompatible` reflects this. Two carry
+ * rules: the claim survives only a node that keeps the declared
+ * key set verbatim -- a node that coarsens it (key-dropping
Review Comment:
**Nit (P3):** Union does not coarsen these keys into a smaller set: `KeyedPartitioning.concat` combines every child's keys, so `{1, 2}` union `{3}` declares `{1, 2, 3}`. The marker must still be rejected because a key unknown to one leg may become declared by another while remaining in the wrong partition. Please separate that key-set expansion rationale from the projection/reducer coarsening cases.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1678,7 +1751,18 @@ case class KeyedShuffleSpec(
// The shuffled side is laid out on this side's partition keys, so it inherits the flag. That
// is conservative rather than strictly true, and it can only ever add a shuffle: a later
// grouping of the shared key set carries the collapsed side's risk.
- partitioning.copy(expressions = newExpressions)
+ //
+ // The child re-shuffled onto this layout may hold keys outside the declared set, so every
+ // partitioning produced here carries the unknown-keys marker too (see
+ // `KeyedPartitioning.mayContainUnknownPartitionKeys`). `createPartitioning` is only reached
+ // from `EnsureRequirements`' shuffle loop to re-shuffle a join child onto the best spec's
+ // layout, and the re-shuffled child's keys are never provably a subset of the declared keys:
+ // a non-keyed (v1) child's keys are unknown to the planner, and a keyed child with an
+ // incompatible partitioning is re-evaluated in this spec's key space, which the planner
+ // cannot bound. Marking every such partitioning is therefore
+ // sound (conservative only when the two sides happen to share a transform and the re-shuffled
+ // keys are a known subset).
+ partitioning.copy(expressions = newExpressions, mayContainUnknownPartitionKeys = true)
Review Comment:
**Blocking (P1):** This marker also needs to participate in `KeyedPartitioning.keysSatisfy`. With subset keys enabled, a marked `(id, k)` layout still satisfies `ClusteredDistribution(id)`, although a declared `(1, 0)` row and a hash-fallback `(1, 1)` row can be in different physical partitions. A following `count(*) OVER (PARTITION BY id)` then skips its exchange and can produce two counts of 1 instead of 2. Please reject marked strict-subset clustering (while retaining full-key clustering) and add a unary window or aggregate regression test.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala:
##########
@@ -192,6 +192,35 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
assert(gpe.outputOrdering === Nil)
}
+ test("SPARK-59050: unknown-keyed child with reducers gives up the keyed claim at the real " +
+ "count") {
+ // No planner path reaches the give-up branch with the built-in transforms --
+ // `createShuffleSpec` refuses a narrowing projection of an unknown-keyed partitioning one
+ // hop earlier, and `areKeysCompatible` pairs it only with same-function partners that have
+ // no reducer (see `GroupPartitionsExec.outputPartitioning`) -- so pin the contract
+ // directly: the node must report `UnknownPartitioning` with its physical grouped count.
+ // Reporting zero partitions is what threw once a parent join built a
+ // `PartitioningCollection` over both sides.
+ val partitionKeys = Seq(row(1), row(2), row(1))
+ val child = DummySparkPlan(
+ outputPartitioning = KeyedPartitioning(Seq(exprA), partitionKeys)
+ .copy(mayContainUnknownPartitionKeys = true))
+ val gpe = GroupPartitionsExec(child, reducers = Some(Seq(None)))
Review Comment:
**Non-blocking (P2):** `Some(Seq(None))` makes `reducers.isDefined` true, but it never applies a `KeyReducer`; `reduceKeys` leaves `[1, 2, 1]` and the key type unchanged. Please provide a small real reducer via `Some(Seq(Some(...)))` that maps to a coarser domain, then assert the reduced grouping count and the parent collection's fallback partition count.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1545,6 +1574,45 @@ case class KeyedShuffleSpec(
}
} && expressions.zip(otherExpressions).forall {
case (l, r) => isExpressionCompatible(l, r)
+ } && {
+ // An unknown-keyed side (see `KeyedPartitioning.mayContainUnknownPartitionKeys`) guarantees
+ // co-location only for its declared keys, and `KeyGroupedPartitioner`'s out-of-set routing
+ // is a deterministic hash -- so it can pair only with a side whose keys are a subset of the
+ // declared keys.
+ //
+ // The key comparison below must also happen in a single domain: `isExpressionCompatible`
+ // admits an `AttributeReference` against a `TransformExpression` (and two different-but-
+ // compatible transforms) when `v2BucketingAllowCompatibleTransforms` is on, and in those
+ // cases the two sides' `partitionKeys` hold raw values on one side and transform outputs on
+ // the other, so the subset test would compare unrelated values. Require the partition
+ // expressions to be the same function per position before comparing keys.
+ //
+ // Two unknown-keyed sides are compatible only when they agree on the declared keys *and*
+ // their order: the out-of-set keys hash to the same-index partition on both sides, and a
+ // `GroupPartitionsExec` regrouping re-labels each partition by that side's declared key, so
+ // a differing declared order would push the out-of-set keys into different output
+ // partitions and lose their matches.
+ if (partitioning.mayContainUnknownPartitionKeys ||
+ other.partitioning.mayContainUnknownPartitionKeys) {
+ expressions.zip(otherExpressions).forall {
+ case (_: AttributeReference, _: AttributeReference) => true
+ case (l: TransformExpression, r: TransformExpression) => l.isSameFunction(r)
Review Comment:
**Non-blocking (P2):** Please add a focused negative test for two marked specs whose transforms are compatible but not the same function, such as `bucket(4)` versus `bucket(8)`, in both argument orders. The current negative test is attribute-versus-transform and every bucket case uses `bucket(4)`, so replacing this check with the broader compatible-transform relation would leave all added tests green.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1678,7 +1751,18 @@ case class KeyedShuffleSpec(
// The shuffled side is laid out on this side's partition keys, so it inherits the flag. That
// is conservative rather than strictly true, and it can only ever add a shuffle: a later
// grouping of the shared key set carries the collapsed side's risk.
- partitioning.copy(expressions = newExpressions)
+ //
+ // The child re-shuffled onto this layout may hold keys outside the declared set, so every
+ // partitioning produced here carries the unknown-keys marker too (see
+ // `KeyedPartitioning.mayContainUnknownPartitionKeys`). `createPartitioning` is only reached
+ // from `EnsureRequirements`' shuffle loop to re-shuffle a join child onto the best spec's
+ // layout, and the re-shuffled child's keys are never provably a subset of the declared keys:
+ // a non-keyed (v1) child's keys are unknown to the planner, and a keyed child with an
+ // incompatible partitioning is re-evaluated in this spec's key space, which the planner
+ // cannot bound. Marking every such partitioning is therefore
+ // sound (conservative only when the two sides happen to share a transform and the re-shuffled
+ // keys are a known subset).
+ partitioning.copy(expressions = newExpressions, mayContainUnknownPartitionKeys = true)
Review Comment:
**Blocking (P1):** A marked multi-partition layout must not satisfy `OrderedDistribution`. For declared keys `[1, 2]`, an out-of-set integer key `4` hashes to partition 0, so the physical partitions can contain `[1, 4]` and `[2]`. `keysSatisfy` currently accepts the requested ordering and `EnsureRequirements` skips range partitioning; the global sort is then only local and emits `[1, 4, 2]`. Please make the ordered-distribution branch reject marked layouts with more than one partition and add an end-to-end `ORDER BY` regression test.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org