This is an automated email from the ASF dual-hosted git repository.
peter-toth pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.2 by this push:
new 40484f28c253 [SPARK-58988][4.2][SQL] Fix PartitioningCollection
invariant violation exception when allowKeysSubsetOfPartitionKeys
40484f28c253 is described below
commit 40484f28c253ee470ca05bf92e6d5fe0dbb44dff
Author: Peter Toth <[email protected]>
AuthorDate: Tue Oct 27 20:28:40 2026 +0100
[SPARK-58988][4.2][SQL] Fix PartitioningCollection invariant violation
exception when allowKeysSubsetOfPartitionKeys
### Differences from #58273
This is the `branch-4.2` backport of #58273. The description below is the
original one.
- **The problem on 4.2 is wrong rows, not a planning failure.** 4.2 has no
`PartitioningCollection.fromPartitionings` check (SPARK-56877), so nothing
fails. `ShuffleExchangeExec` routes the re-shuffled side by the sorted keys,
while `createShuffleSpec` declares them in first-occurrence order.
- With AQE off, the keyed side's `GroupPartitionsExec` sorts its keys
too, so the two sides line up.
- With AQE on, the join is planned again over the materialized shuffle.
The `GroupPartitionsExec` planned above that shuffle reads it in the declared
order, so the two sides pair different keys. An inner join returns no rows, and
a LEFT JOIN returns nulls for the shuffled side.
- Measured on `v4.2.0` and on `branch-4.2`. `branch-4.0`, `branch-4.1`
and `master` return the right rows.
- **Code.** 4.2's `toGrouped` builds a new `KeyedPartitioning` and drops
`mayContainUnknownPartitionKeys`, so the backport takes only its sorted keys
and keeps the marker. On 4.3 and `master`, `toGrouped` keeps it through `copy`.
The code comment describes the 4.2 symptom instead of the
`PartitioningCollection` check.
- **Tests.**
- The original test runs on 4.2 and is kept. Its subset conf is
`V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS` here, and its comment
no longer mentions the collection check. It passes on 4.2 without the fix too,
since that failure needs the check.
- A new test, `SPARK-58988: one-side shuffle onto subset join keys whose
projected keys are unsorted`, covers the 4.2 symptom with an inner and a LEFT
JOIN, AQE on and off. It fails on `branch-4.2` without the fix and passes with
it.
- **Conflicts.**
- The projected branch of `createShuffleSpec` conflicted. On 4.3 and
`master`, this change landed before SPARK-59050, which added
`mayContainUnknownPartitionKeys`. On 4.2, the SPARK-59050 backport is already
in, so the lines this change replaces already pass the marker. The resolution
sorts the keys and keeps the marker, which is what 4.3 does today with both
changes in.
- The test conflicted only in its position, since its neighbour on 4.3
(SPARK-46367) is not on 4.2. Both tests are at the end of the suite.
- **What I ran on `branch-4.2`.**
- `ShuffleSpecSuite` and `DistributionSuite`: 30 tests.
- `KeyGroupedPartitioningSuite`, `EnsureRequirementsSuite`,
`GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite`,
`WriteDistributionAndOrderingSuite`, `PlannerSuite` and `SparkPlannerSuite`:
434 tests.
- `AdaptiveQueryExecSuite`, `DataSourceV2Suite`, `FileDataSourceV2Suite`,
`NestedDataSourceV2Suite` and `StreamingDataSourceV2Suite`: 211 tests.
- `dev/lint-scala`.
- This backport was prepared with Claude Code (Claude Opus 5.5).
### What changes were proposed in this pull request?
`KeyedPartitioning.createShuffleSpec` now sorts the distinct projected
partition keys (via
`toGrouped`) so they follow the same natural ascending ordering as
`GroupPartitionsExec`.
### Why are the changes needed?
SPARK-56877 added a check in `PartitioningCollection.fromPartitionings`
requiring all
`KeyedPartitioning`s to share equal `partitionKeys`. In a
storage-partitioned join whose join keys
are a subset of the partition keys (e.g. a v2 table partitioned by `[dt,
bucket(16, c1)]` joined
on `c1`), with `spark.sql.sources.v2.bucketing.shuffle.enabled` enabled so
only the non-keyed side
is re-shuffled, the keyed side's projected keys are sorted by
`GroupPartitionsExec` while
`createShuffleSpec` kept them in first-occurrence order. The two sides then
carry the same keys in
different orders and the query fails with:
```
java.lang.IllegalArgumentException: requirement failed: All
KeyedPartitionings in a
PartitioningCollection must have equal partitionKeys
```
### Does this PR introduce _any_ user-facing change?
No by default.
### How was this patch tested?
Added a regression test in `KeyGroupedPartitioningSuite`
(`SPARK-58988: v2 bucketed table with subset join keys joining v1 table`)
that reproduces the
failure and passes with the fix. Also ran `KeyGroupedPartitioningSuite`,
`EnsureRequirementsSuite`,
`GroupPartitionsExecSuite`, and `ProjectedOrderingAndPartitioningSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #59832 from peter-toth/SPARK-58988-4.2.
Lead-authored-by: Peter Toth <[email protected]>
Co-authored-by: Xiduo You <[email protected]>
Signed-off-by: Peter Toth <[email protected]>
---
.../sql/catalyst/plans/physical/partitioning.scala | 11 +++-
.../connector/KeyGroupedPartitioningSuite.scala | 72 ++++++++++++++++++++++
2 files changed, 81 insertions(+), 2 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
index 76e930395ea2..0e51227cb497 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
@@ -643,9 +643,16 @@ case class KeyedPartitioning(
// the narrowing one.
val projectedExpressions = joinKeyPositions.map(expressions)
val projectedKeys = projectKeys(joinKeyPositions)._2
- val distinctProjectedKeys = projectedKeys.distinct
+ // Sort the distinct projected keys the same way `GroupPartitionsExec`
does. When only the
+ // keyed side is grouped, the other side is re-shuffled using this spec,
and
+ // `ShuffleExchangeExec` routes it by the sorted keys. So the keys
declared here have to be
+ // sorted too. Otherwise a `GroupPartitionsExec` that AQE plans over
that shuffle reads its
+ // partitions in the wrong order. `toGrouped` drops the unknown-keys
marker, so only its keys
+ // are taken.
+ val groupedProjectedKeys = new KeyedPartitioning(
+ projectedExpressions, projectedKeys, isGrouped =
false).toGrouped.partitionKeys
val projectedPartitioning = KeyedPartitioning(
- projectedExpressions, distinctProjectedKeys,
+ projectedExpressions, groupedProjectedKeys,
isGrouped = true, mayContainUnknownPartitionKeys =
mayContainUnknownPartitionKeys)
result.copy(partitioning = projectedPartitioning, joinKeyPositions =
Some(joinKeyPositions))
} else {
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
index 7b9f21d6d1eb..9d91e4f3bcfa 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -6999,4 +6999,76 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
}
}
+ test("SPARK-58988: v2 bucketed table with subset join keys joining v1
table") {
+ // The v2 table is partitioned by an extra identity key `dt` plus
`bucket(16, c1)`, while the
+ // join is only on `c1`. allowJoinKeysSubsetOfPartitionKeys lets the join
key `c1` be a subset
+ // of the partition keys `[dt, bucket(16, c1)]`, so EnsureRequirements
projects the keyed side
+ // to `[bucket(16, c1)]`. v2BucketingShuffleEnabled then re-shuffles only
the v1 side using that
+ // projected KeyedPartitioning. The v2 side's keys are sorted by
GroupPartitionsExec, while the
+ // keys re-used for the v1 side kept their first-occurrence order from
createShuffleSpec.
+ val cols = Array(
+ Column.create("c1", LongType),
+ Column.create("c2", StringType),
+ Column.create("dt", StringType))
+ val partitions = Array(identity("dt"), bucket(16, "c1"))
+
+ createTable("iceberg_t2", cols, partitions)
+ sql("INSERT INTO testcat.ns.iceberg_t2 VALUES (2, 'cc', '2020'), (1, 'aa',
'2021')")
+
+ withTable("t1") {
+ sql("CREATE TABLE t1 (c1 BIGINT, c2 STRING) USING parquet")
+ sql("INSERT INTO t1 VALUES (1, 'aa'), (2, 'cc')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql("SELECT * FROM testcat.ns.iceberg_t2 t0 JOIN t1 ON t0.c1
= t1.c1")
+ val plan = df.queryExecution.executedPlan
+ // Only the v1 side is re-shuffled; the v2 side is regrouped onto the
join key instead.
+ assert(collectShuffles(plan).length == 1)
+ assert(collectGroupPartitions(plan).length == 1)
+ checkAnswer(df, Seq(
+ Row(1L, "aa", "2021", 1L, "aa"),
+ Row(2L, "cc", "2020", 2L, "cc")))
+ }
+ }
+ }
+
+
+ test("SPARK-58988: one-side shuffle onto subset join keys whose projected
keys are unsorted") {
+ // `sa` is partitioned by (store, dept) and joined on dept only. Projected
onto dept, its keys
+ // come out as 20, 10. `sb` is re-shuffled onto that projected spec, and
`ShuffleExchangeExec`
+ // routes it by the sorted keys. With AQE, the join is planned again over
the materialized
+ // shuffle, and the `GroupPartitionsExec` planned there reads its
partitions in the declared
+ // order. Both joins returned wrong rows with AQE on before the declared
keys were sorted.
+ val cols = Array(
+ Column.create("store", IntegerType),
+ Column.create("dept", IntegerType),
+ Column.create("data", StringType))
+ createTable("sa", cols, Array(identity("store"), identity("dept")))
+ sql("INSERT INTO testcat.ns.sa VALUES (1, 20, 'a1'), (2, 10, 'a2')")
+ createTable("sb", cols, Array.empty)
+ sql("INSERT INTO testcat.ns.sb VALUES (0, 10, 'b10'), (0, 20, 'b20')")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ Seq("JOIN", "LEFT JOIN").foreach { joinType =>
+ // Selecting `sa.store` keeps `sa`'s keyed partitioning from being
pruned away.
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(sa, sb) */ sa.dept, sa.data, sb.data, sa.store
+ |FROM testcat.ns.sa $joinType testcat.ns.sb ON sa.dept = sb.dept
+ |""".stripMargin)
+ assert(collectAllShuffles(df.queryExecution.executedPlan).length ==
1,
+ s"only sb should be shuffled for $joinType")
+ checkAnswer(df, Seq(Row(10, "a2", "b10", 2), Row(20, "a1", "b20",
1)))
+ }
+ }
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]