peter-toth commented on code in PR #58339:
URL: https://github.com/apache/spark/pull/58339#discussion_r3873427831
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala:
##########
@@ -152,7 +152,8 @@ trait PartitioningPreservingUnaryExecNode extends
UnaryExecNode
MultiTransform.generateCartesianProduct(
projectablePositions.map(i => () => alternativesPerPosition(i)))
.map(projectedExprs =>
- new KeyedPartitioning(projectedExprs, sharedKeys, isGrouped,
isNarrowed))
+ new KeyedPartitioning(projectedExprs, sharedKeys, isGrouped,
isNarrowed,
+ kps.exists(_.mayContainUnknownPartitionKeys)))
Review Comment:
**Finding 1.** The marker travels through the projection but the declared
key set does not. When `projectablePositions.length < numPositions`,
`sharedKeys` is `keySource.projectKeys(...)`, so two different full keys
collapse onto one projected key -- and a key that was *outside* the declared
set before the projection can be *inside* it after. `areKeysCompatible`'s
subset test then reads as a guarantee that no longer holds.
Measured on `83af168`, only
`spark.sql.sources.v2.bucketing.shuffle.enabled=true` and AQE off (as in the
new tests):
```sql
-- a: keyed on (id, k), keys {(1,x),(2,x),(3,x),(4,x)}; u: keyed on id,
keys {1,2,3,4}
-- t: v1 parquet with (1,z),(2,z),(3,z),(4,z) -- every (id,k) is out-of-set,
no id is
SELECT r.id, u.k
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
```
Expected 4 rows, got 2 (`[1,u1], [3,u3]`); with
`spark.sql.sources.v2.bucketing.enabled=false` all 4 come back. The `Project
[id]` above the RIGHT OUTER join drops key position 1, so the flagged
partitioning declares `{1,2,3,4}`, `u`'s keys are a subset, the second
`SortMergeJoin` storage-partitions with no exchange, and the `t` rows whose
`(id,k)` hashed to the wrong partition never meet their `u` match.
The marker is only sound while the key set survives verbatim, so drop the
keyed claim when it doesn't -- right after `projectablePositions` is computed:
```scala
if (projectablePositions.length < numPositions &&
kps.exists(_.mayContainUnknownPartitionKeys)) {
return LazyList.empty
}
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -990,7 +990,9 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
val mergedExpressions = headKp.expressions
val isGrouped = mergedKeys.distinct.size == mergedKeys.size
val isNarrowed = kps.exists(_.isNarrowed)
- return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped,
isNarrowed)
+ val mayContainUnknownPartitionKeys =
kps.exists(_.mayContainUnknownPartitionKeys)
Review Comment:
**Finding 2.** `mergedKeys` is the concatenation of every leg's keys, so the
merged declared set is a *superset* of the flagged leg's. A key that the
flagged leg holds out-of-set can be contributed as a declared key by another
leg, and `areKeysCompatible` then accepts a partner that holds it.
Measured on `83af168`, same single config. `s` is keyed on `id` with key
`{3}` -- exactly the key `t` has and `a` does not:
```sql
-- a keyed {1,2}; t v1 parquet {1,2,3}; s keyed {3}; u keyed {1,2,3}
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
```
Expected 4 rows (`id=3` matches `u` twice, once via `t` and once via `s`),
got 3 -- `t`'s `id=3` is lost. The union declares `[1,2,3]`, which equals `u`'s
keys, so the second join storage-partitions with no `GroupPartitionsExec` and
no exchange.
This is the shape "SPJ: union preserves the unknown-partition-keys marker"
covers, but with a deliberately *disjoint* leg (`s` = `{4,5}`), which is why it
passes.
The merged claim can't be trusted once a leg is flagged, so fall back before
building it:
```scala
if (compatible) {
if (kps.exists(_.mayContainUnknownPartitionKeys)) {
return super.outputPartitioning
}
```
A laxer rule would keep the merged partitioning when every flagged leg's key
set already equals `mergedKeys`, but with more than one leg that is only the
all-legs-identical case. The existing union test would then assert the opposite
-- that the keyed partitioning is dropped.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1336,6 +1349,31 @@ case class KeyedShuffleSpec(
}
} && expressions.zip(otherExpressions).forall {
case (l, r) => isExpressionCompatible(l, r)
+ } && {
+ // A partitioning that may contain unknown partition keys (see
+ // `mayContainUnknownPartitionKeys`) only guarantees co-location for its
declared keys: keys
+ // outside the declared set were routed to arbitrary partitions by
`KeyGroupedPartitioner`.
+ // That routing is a deterministic hash of the key, so a flagged side
can only be
+ // co-partitioned with a side whose keys are a subset of the declared
keys -- the other
+ // side's keys must all land in the partitions this side declares.
+ //
+ // Two flagged 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) {
+ partitioning.partitionKeys == other.partitioning.partitionKeys
+ } else if (partitioning.mayContainUnknownPartitionKeys) {
Review Comment:
**Finding 5.** This compares the two sides' `partitionKeys` directly, but
`isExpressionCompatible` just above admits an `AttributeReference` against a
`TransformExpression` (and two different-but-compatible transforms) when
`allowCompatibleTransforms` is on. In those cases the two key sequences live in
different domains -- raw `id` values on one side, bucket ids on the other -- so
`declared.contains` is comparing unrelated numbers.
It can reject a sound pairing, and it can accept an unsound one: flagged
identity keys `{0,1,2,3}` "contain" bucket keys `{0,1}`, so the pairing passes
while the question that matters (are the flagged side's out-of-set ids in the
partition their bucket label points at?) is never asked. Also from reading, not
measured.
`EnsureRequirements` computes `leftReducedKeys` / `rightReducedKeys` a few
lines below the `areKeysCompatible` call, so comparing in the reduced key space
is available; refusing the marker path outright unless both sides are the same
function would also close it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -74,7 +74,8 @@ case class GroupPartitionsExec(
p.transform {
case k: KeyedPartitioning =>
val projectedExpressions =
joinKeyPositions.fold(k.expressions)(_.map(k.expressions))
- KeyedPartitioning(projectedExpressions, partitionKeys, isGrouped =
isGrouped)
+ KeyedPartitioning(projectedExpressions, partitionKeys, isGrouped =
isGrouped,
+ mayContainUnknownPartitionKeys =
k.mayContainUnknownPartitionKeys)
Review Comment:
**Finding 4.** Same shape as findings 1 and 2, one hop earlier. With
`reducers` defined, `groupedPartitions` is keyed by the *reduced* keys, so this
partitioning declares a coarsened set -- and an out-of-set key can reduce into
it. A flagged side with identity keys `{0,1,2,3}` that holds an out-of-set
`id=4`, reduced by `bucket(4, .)`, declares `{0,1,2,3}` again, and
`bucket(4,4)=0` is one of them, while the `id=4` rows sit in
`nonNegativeMod(hash(4), 4)`.
I did not build this one -- it needs
`spark.sql.sources.v2.bucketing.allowCompatibleTransforms=true` on top of the
shuffle config, and the reducer has to be picked for the flagged side. The
experiment that settles it: `a` keyed `identity(id)` with ids 0..3, `t` v1
holding id 4, `u` keyed `bucket(4, id)` holding id 4, then `a RIGHT OUTER JOIN
t` joined to `u` on `id`. Same fix shape as findings 1 and 2: don't propagate
the keyed claim when `reducers.isDefined`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -4727,4 +4751,266 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
assert(shuffles.isEmpty, "should not contain any shuffle")
checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(2, "bb", 10.0, 19.5)))
}
+
+ test("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("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"))
+
+ Seq("RIGHT OUTER").foreach { joinType =>
Review Comment:
**Finding 7.** This iterates a single element -- and LEFT OUTER and FULL
OUTER each get their own block below rather than joining the loop. Either
inline the body and drop the loop, or fold the two later blocks in (they differ
in the assertions, so inlining is probably simpler).
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -4727,4 +4751,266 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
assert(shuffles.isEmpty, "should not contain any shuffle")
checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(2, "bb", 10.0, 19.5)))
}
+
+ test("SPJ: one-side shuffle with out-of-set keys loses matches in a
following SPJ join") {
Review Comment:
**Finding 8.** 89 of this suite's 110 tests are prefixed with their ticket
id; these six aren't. Worth `SPARK-59050: SPJ: ...` on each so they are
greppable from the ticket.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -4727,4 +4751,266 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
assert(shuffles.isEmpty, "should not contain any shuffle")
checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(2, "bb", 10.0, 19.5)))
}
+
+ test("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("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"))
+
+ Seq("RIGHT OUTER").foreach { joinType =>
+ val query =
+ s"""
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a $joinType 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, 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,
Review Comment:
**Finding 9.** The other three blocks in this test also assert
`collectGroupPartitions(...).isEmpty`. The comment here says FULL OUTER is safe
*because* it exposes `UnknownPartitioning`, and that is exactly what the
missing assertion would pin -- as written, nothing fails if a future change
lets the downstream join storage-partition on this layout.
--
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]
