ulysses-you commented on code in PR #58339:
URL: https://github.com/apache/spark/pull/58339#discussion_r3877527559
##########
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:
Fixed in 355d38891f4: `projectKeyedPartitionings` now returns an empty list
when the projection drops a key position (`projectablePositions.length <
numPositions`) and any input KP is flagged, so the keyed claim is dropped
entirely instead of being carried across the coarsened declared set. Added a
wrong-results repro test (`SPARK-59050: SPJ: project dropping a key position
drops the unknown-keyed claim`) that fails on the previous commit with 4
expected / 2 returned, matching your measurement, and passes after the fix.
##########
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:
Fixed in 355d38891f4: `UnionExec` now returns `super.outputPartitioning`
whenever any leg is flagged, before building `mergedKeys`. Since the merged set
is the concatenation of every leg's keys, with multiple legs it is always
larger than any single leg's declared set, so there is no sound way to keep the
merged claim. Updated the union test accordingly: it now asserts the merged
keyed partitioning is dropped (no `GroupPartitionsExec`, the union side
re-shuffles), and covers both a disjoint second leg and your wider shape where
the second leg declares exactly the out-of-set key (fails 4 expected / 3
returned before the fix, passes after).
##########
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:
Guarded in 355d38891f4: `GroupPartitionsExec.outputPartitioning` returns
`UnknownPartitioning(0)` when `reducers.isDefined` and any input KP is flagged.
Note this is deliberately not `super.outputPartitioning` -- for a
`UnaryExecNode` that is the child's flagged keyed partitioning, which is
equally untrustworthy here since the regrouped data layout no longer matches
the child's declared keys. I did not build the end-to-end experiment (it needs
`allowCompatibleTransforms` plus the reducer to be picked for the flagged
side), so this is a conservative guard rather than a repro-tested fix -- happy
to add the experiment as a follow-up if you'd like.
##########
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:
Fixed in 355d38891f4 (the 'refuse the marker path unless both sides are the
same function' option): when either side is flagged, `areKeysCompatible` now
requires the partition expressions to match per position -- two leaf
attributes, or `TransformExpression.isSameFunction` -- and returns false
otherwise, so the subset comparison always happens within one key domain.
Unflagged pairs keep the existing behavior (the identity-vs-transform pair
stays admissible and `EnsureRequirements` reconciles the domains via reducers).
Unit-tested in `ShuffleSpecSuite`.
##########
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:
Inlined in 355d38891f4 -- the RIGHT OUTER block now stands on its own
alongside the FULL OUTER and LEFT OUTER blocks.
##########
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:
Done in 355d38891f4 -- all six new tests are now prefixed with
`SPARK-59050:`.
##########
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:
Added in 355d38891f4 -- the FULL OUTER block now also asserts
`collectGroupPartitions(...).isEmpty`, pinning the comment's claim.
--
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]
