ulysses-you commented on code in PR #58339:
URL: https://github.com/apache/spark/pull/58339#discussion_r3912378575
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -578,12 +582,36 @@ 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 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 sharing a subset of its columns.
`satisfies` and
+ * `KeyedShuffleSpec.areKeysCompatible`
therefore accept a marked
+ * partitioning only for full-key clustering,
never for a subset
+ * of its partition columns and never for a
global ordering
+ * across several partitions. Two carry rules:
(1) a node that
+ * changes the declared key set must drop the
keyed partitioning,
+ * whether it coarsens it (key-dropping
projection, reducer,
+ * join-key projection) or expands it over a
marked leg (a union,
+ * where another leg may declare exactly the
key that leg holds
+ * out-of-set); (2) a marked member beside an
unmarked sibling is
+ * spurious, and `ShuffledJoin`'s `InnerLike`
arm, the only site
+ * that mixes them, clears it at construction,
so members are
+ * uniformly marked or unmarked everywhere
downstream.
*/
case class KeyedPartitioning(
expressions: Seq[Expression],
@transient partitionKeys: Seq[InternalRowComparableWrapper],
isGrouped: Boolean,
- isCollapsed: Boolean) extends Expression with Partitioning with
Unevaluable {
+ isCollapsed: Boolean,
+ mayContainUnknownPartitionKeys: Boolean = false)
Review Comment:
Thanks for working through the premise honestly, including the
self-correction from round 3, @peter-toth. Implemented your `isCollapsed`
mirror in full in 6021d83: `fromPartitionings` OR-normalizes
`mayContainUnknownPartitionKeys` alongside `isCollapsed` (one pass over direct
representatives per flag, kept as two `exists` calls so each short-circuits),
the constructor gained the matching `require`, and both the class doc and the
`@param`'s carry rule (2) now name the enforcement. Consumers are unchanged --
`kps.head` and the per-member `GroupPartitionsExec` reads are now both provably
right by the invariant rather than by the site enumeration. The two new guard
rails are themselves tested in `DistributionSuite` (`fromPartitionings
normalizes the unknown-keys marker by OR` and `requires members to agree on the
marker`; the latter shares the keys reference deliberately so it reaches the
marker `require` past the reference check). With the `require` in place, a
future site that mixes
markers fails loudly at construction instead of silently trusting one member.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala:
##########
@@ -80,6 +87,29 @@ trait ShuffledJoin extends JoinCodegenSupport {
s"ShuffledJoin should not take $x as the JoinType")
}
+ /**
+ * Clears the `mayContainUnknownPartitionKeys` marker of every
`KeyedPartitioning` in
+ * `partitionings` when at least one member is unmarked. Only
`ShuffledJoin`'s `InnerLike` arm
+ * can mix the two; see the call site for the argument. The interning
`fromPartitionings` does
+ * afterwards changes no marker.
+ */
+ private def clearUnknownPartitionKeys(
+ partitionings: Seq[Partitioning]): Seq[Partitioning] = {
+ val members = partitionings.flatMap(PartitioningCollection.flatten)
+ .collect { case k: KeyedPartitioning => k }
+ if (members.isEmpty || members.forall(_.mayContainUnknownPartitionKeys)) {
+ partitionings
+ } else {
+ partitionings.map {
+ case e: Expression =>
+ e.transform {
+ case k: KeyedPartitioning => k.copy(mayContainUnknownPartitionKeys
= false)
Review Comment:
Fixed in 6021d83, thanks -- both layers: the transform case now only matches
marked members, and the whole rebuild is skipped behind an `exists`/`forall`
pair over a lazy view, so the all-unmarked inner join allocates nothing and
compares no partition key on `outputPartitioning`. I did not take the
`representativeOf` follow-on: it is `private[physical]`, and with the
copy/equals cost gone the remaining tree walk is the short-circuiting scan
itself; the comment states the `fastEquals` fall-through so nobody re-widens
the case later.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5721,6 +5734,828 @@ 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 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
Review Comment:
Reworded in 6021d83, thanks -- the comment now says the merged set is a
superset, possibly equal, and that keeping the marker for same-set legs would
in fact be sound (its out-of-set row rides its own leg's partition into that
partition's group); the check does not try to tell that case from the rest, and
refuses.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5721,6 +5734,828 @@ class KeyGroupedPartitioningSuite
}
}
+
Review Comment:
Both dropped in 6021d83, thanks -- the blank before the assert helper and
the pair at end of file.
--
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]
