You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by do...@apache.org on 2026/08/26 19:18:58 UTC
(spark) branch branch-4.3 updated: [SPARK-59022][SQL] Keyed shuffle must follow the declared partition key order
This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new 568b2c2f7947 [SPARK-59022][SQL] Keyed shuffle must follow the declared partition key order
568b2c2f7947 is described below
commit 568b2c2f7947bedb0002a51f83d5057cc53954dd
Author: Peter Toth <pe...@gmail.com>
AuthorDate: Wed Aug 26 12:18:16 2026 -0700
[SPARK-59022][SQL] Keyed shuffle must follow the declared partition key order
### What changes were proposed in this pull request?
`ShuffleExchangeExec.getPartitioner` built its `KeyGroupedPartitioner` from `KeyedPartitioning.toGrouped`, i.e. from the *sorted* partition keys. It now builds it from `partitionKeys` in the order the partitioning declares, and asserts the keys are unique.
`KeyedShuffleSpec.canCreatePartitioning` additionally refuses an ungrouped partitioning, so a spec whose keys are not unique is never chosen as the template for shuffling the other child.
### Why are the changes needed?
`partitionKeys` is a physical layout indexed by partition id: partition `i` holds key `partitionKeys(i)`. With `spark.sql.sources.v2.bucketing.shuffle.enabled`, `EnsureRequirements` shuffles only the non-keyed side of a storage-partitioned join and reuses the keyed side's `KeyedPartitioning` as the target (`KeyedShuffleSpec.createPartitioning`). Re-deriving the order in `getPartitioner` therefore contradicts the side the shuffle is supposed to co-partition with, and the join reads row [...]
Those keys are not always sorted. Two producers report them unsorted:
- `UnionExec` concatenates its children's keys in child order, so children partitioned by `identity(id)` holding `[3, 4]` and `[1, 2]` merge to `[3, 4, 1, 2]`.
- a narrowing projection through `PartitioningPreservingUnaryExecNode` projects the keys onto a subset of key positions, which preserves neither sortedness nor uniqueness.
Where uniqueness is lost the result no longer satisfies a `ClusteredDistribution`, so `EnsureRequirements` interposes a `GroupPartitionsExec` that re-groups and sorts. Where only sortedness is lost the partitioning still satisfies, nothing is interposed, and the unsorted keys are the real layout of that side.
The result is silently wrong. Nothing catches it: `createPartitioning` passes the keyed side's `partitionKeys` *reference* through unchanged, so `PartitioningCollection.fromPartitionings` interns it on `eq` and never runs the `require` added by SPARK-56877, and `ValidateRequirements` compares the same declared keys on both sides.
To reproduce, two v2 tables partitioned by `identity(id)` that report one input partition per key, one holding ids `3, 4` and the other `1, 2`:
```sql
CREATE TABLE testcat.ns.nt1 (id BIGINT, data STRING) PARTITIONED BY (id);
CREATE TABLE testcat.ns.nt2 (id BIGINT, data STRING) PARTITIONED BY (id);
INSERT INTO testcat.ns.nt1 VALUES (3, 'c'), (4, 'd');
INSERT INTO testcat.ns.nt2 VALUES (1, 'a'), (2, 'b');
CREATE TABLE t1 (id BIGINT, x STRING) USING parquet;
INSERT INTO t1 VALUES (1, 'x'), (2, 'x'), (3, 'x'), (4, 'x');
SET spark.sql.sources.v2.bucketing.shuffle.enabled=true;
SELECT u.id, u.data, t1.x
FROM (SELECT * FROM testcat.ns.nt1 UNION ALL SELECT * FROM testcat.ns.nt2) u
JOIN t1 ON u.id = t1.id;
```
keyed side (not re-grouped) partition 0 -> id 3 | 1 -> 4 | 2 -> 1 | 3 -> 2
shuffled side, before partition 0 -> id 1 | 1 -> 2 | 2 -> 3 | 3 -> 4
shuffled side, after partition 0 -> id 3 | 1 -> 4 | 2 -> 1 | 3 -> 2
Before the change the query returns no rows; the correct answer is four.
The affected branches are `master`, `branch-4.x` and `branch-4.3`. `branch-4.2` sorts in `getPartitioner` too but has no producer of unsorted keys, and `branch-4.1` and `branch-4.0` do not sort at all -- their `uniquePartitionValues` only deduplicates.
SPARK-52246 fixed the same symptom for a different trigger, where the join keys are a subset of the partition keys. This one needs no subset keys.
It is also distinct from SPARK-58988, which fixed the projected branch of `createShuffleSpec`. There a `GroupPartitionsExec` *is* interposed on the keyed side and sorts it, so the shuffled side's declared keys have to be sorted to match, and the mismatch surfaced as a loud `require` failure rather than wrong rows. Both are needed: reverting either one leaves its own case broken.
### Does this PR introduce _any_ user-facing change?
Yes, it fixes wrong results, but only for queries that opt into `spark.sql.sources.v2.bucketing.shuffle.enabled`, which is off by default.
### How was this patch tested?
Added two tests to `KeyGroupedPartitioningSuite`, one per producer of unsorted keys: the union case above, and a narrowing projection that drops a partition column. Both fail without the change, with `Correct Answer - 4` against `Spark Answer - 0` and `Correct Answer - 2` against `Spark Answer - 0` respectively. The plan shape is identical either way (one shuffle, no `GroupPartitionsExec`), which the tests assert, so `checkAnswer` is what discriminates. The union test also re-checks t [...]
Added a `ShuffleSpecSuite` test for the `canCreatePartitioning` gate, covering a grouped spec with unsorted keys (accepted) and an ungrouped one (refused); it fails if the new clause is removed.
Ran `ShuffleSpecSuite`, `KeyGroupedPartitioningSuite`, `EnsureRequirementsSuite`, `GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite`, `DataSourceV2Suite`, `WriteDistributionAndOrderingSuite`, `ValidateRequirementsSuite`, `PlannerSuite` and `AdaptiveQueryExecSuite`: 595 tests, 0 failures. Also verified with a temporary assertion that no covered scenario reaches `getPartitioner` with duplicate partition keys, and that a duplicating narrowing projection is handled by a ` [...]
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #58311 from peter-toth/SPARK-59022-keyed-shuffle-declared-key-order.
Authored-by: Peter Toth <pe...@gmail.com>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
(cherry picked from commit 99aed2a98525f6cb564e7b0a4063e87a9b4af6d9)
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
.../sql/catalyst/plans/physical/partitioning.scala | 26 +++++--
.../spark/sql/catalyst/ShuffleSpecSuite.scala | 20 +++++
.../execution/exchange/ShuffleExchangeExec.scala | 15 +++-
.../connector/KeyGroupedPartitioningSuite.scala | 89 ++++++++++++++++++++++
4 files changed, 140 insertions(+), 10 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 7376c0ff4466..31861559f6ec 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
@@ -446,19 +446,26 @@ case class CoalescedNullAwareHashPartitioning(
* because data sources produce them that way and `GroupPartitionsExec` sorts while grouping.
* Sorted order is not a hard requirement, but it is a useful property: when both sides of a
* storage-partitioned join report sorted keys, `EnsureRequirements` can often match them
- * without inserting an additional `GroupPartitionsExec`. After a narrowing projection through
- * `PartitioningPreservingUnaryExecNode`, the projected keys may no longer be sorted; this is
- * acceptable because `EnsureRequirements` can always reconcile both sides via
- * `GroupPartitionsExec` with `expectedPartitionKeys`.
+ * without inserting an additional `GroupPartitionsExec`. The keys may no longer be sorted after
+ * a narrowing projection through `PartitioningPreservingUnaryExecNode` or after `UnionExec`
+ * concatenates its children's keys; `EnsureRequirements` reconciles both sides either via
+ * `GroupPartitionsExec` with `expectedPartitionKeys`, or -- when the chosen shuffle spec is a
+ * `KeyedShuffleSpec` and the other child's spec is not compatible with it -- by shuffling that
+ * other child onto these keys in the order given here.
*
* 2. '''In KeyedShuffleSpec''': When used within `KeyedShuffleSpec`, the `partitionKeys` may not
- * be in sorted order. `EnsureRequirements` handles this by building a common ordered set of
- * keys and pushing them down to `GroupPartitionsExec` on both sides.
+ * be in sorted order, and consumers must not assume otherwise.
*
* == Partition Keys ==
* - `partitionKeys`: The partition keys, one per partition. May contain duplicates initially
* (ungrouped state), but becomes unique after `GroupPartitionsExec` applies grouping.
*
+ * `partitionKeys` is a physical layout indexed by partition id, not a set: partition `i` holds key
+ * `partitionKeys(i)`. A consumer must therefore treat the given order as authoritative rather than
+ * re-derive one. In particular `ShuffleExchangeExec` builds its `KeyGroupedPartitioner` from this
+ * order, so that a side shuffled onto a `KeyedPartitioning` lands in the same partitions as the
+ * side that declared it.
+ *
* == Grouping State ==
* A KeyedPartitioning can be in two states:
*
@@ -1377,6 +1384,13 @@ case class KeyedShuffleSpec(
override def canCreatePartitioning: Boolean =
SQLConf.get.v2BucketingShuffleEnabled &&
!SQLConf.get.v2BucketingPartiallyClusteredDistributionEnabled &&
+ // Shuffling another child onto these partition keys assigns each key a single partition, so
+ // an ungrouped partitioning cannot be reproduced: its duplicate keys live in more than one
+ // partition. This is the local gate. Such a spec is not reachable today for a non-local
+ // reason -- `EnsureRequirements` wraps a child whose `KeyedPartitioning` does not satisfy the
+ // distribution in a `GroupPartitionsExec` before it builds any spec -- so do not read this
+ // clause as redundant.
+ partitioning.isGrouped &&
partitioning.expressions.forall { e =>
e.isInstanceOf[AttributeReference] || e.isInstanceOf[TransformExpression]
}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
index cb5d77d44512..35f06a722f19 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
@@ -424,6 +424,26 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
assert(!RangeShuffleSpec(10, distribution).canCreatePartitioning)
}
+ test("SPARK-59022: canCreatePartitioning: KeyedShuffleSpec requires grouped partition keys") {
+ val a = $"a".int
+ val distribution = ClusteredDistribution(Seq(a))
+ def keyedSpec(keys: Seq[Int]): KeyedShuffleSpec = KeyedShuffleSpec(
+ KeyedPartitioning(Seq(a), keys.map(k => InternalRow(k))), distribution)
+
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ val grouped = keyedSpec(Seq(2, 1))
+ assert(grouped.partitioning.isGrouped)
+ assert(grouped.canCreatePartitioning,
+ "unsorted keys are fine, the shuffle follows the declared order")
+
+ // Duplicate keys mean one key spans several partitions, which a KeyGroupedPartitioner cannot
+ // reproduce, so this spec must not be the target for shuffling the other child.
+ val ungrouped = keyedSpec(Seq(1, 1, 2))
+ assert(!ungrouped.partitioning.isGrouped)
+ assert(!ungrouped.canCreatePartitioning)
+ }
+ }
+
test("createPartitioning: HashShuffleSpec") {
checkCreatePartitioning(
HashShuffleSpec(HashPartitioning(Seq($"a"), 10), ClusteredDistribution(Seq($"a", $"b"))),
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
index 0059473b3c65..0f3bb973000c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
@@ -399,11 +399,18 @@ object ShuffleExchangeExec {
samplePointsPerPartitionHint = SQLConf.get.rangeExchangeSampleSizePerPartition)
case SinglePartition => new ConstantPartitioner
case k: KeyedPartitioning =>
- val keyGroupedPartitioning = k.toGrouped
- val valueMap = keyGroupedPartitioning.partitionKeys.zipWithIndex.map {
- case (key, index) => (key.row.toSeq(keyGroupedPartitioning.expressionDataTypes), index)
+ // `partitionKeys` is the physical layout its producer declared: partition `i` holds key
+ // `partitionKeys(i)`. Keep that order, whatever it is -- it need not be sorted, and
+ // re-deriving one here would disagree with the side this shuffle co-partitions with.
+ // Unique keys are a precondition, since each key gets exactly one partition;
+ // `KeyedShuffleSpec.canCreatePartitioning` is what refuses an ungrouped partitioning.
+ assert(k.isGrouped,
+ s"Expected a grouped KeyedPartitioning on ${k.expressions}, but got ${k.numPartitions} " +
+ "partition keys with duplicates among them")
+ val valueMap = k.partitionKeys.zipWithIndex.map {
+ case (key, index) => (key.row.toSeq(k.expressionDataTypes), index)
}.toMap
- new KeyGroupedPartitioner(mutable.Map.from(valueMap), keyGroupedPartitioning.numPartitions)
+ new KeyGroupedPartitioner(mutable.Map.from(valueMap), k.numPartitions)
case _ => throw SparkException.internalError(s"Exchange not implemented for $newPartitioning")
// TODO: Handle BroadcastPartitioning.
}
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 f5e2b70725bd..7d335f285f50 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
@@ -4256,6 +4256,95 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
}
}
+ test("SPARK-59022: keyed shuffle follows the declared partition key order") {
+ val cols = Array(
+ Column.create("id", LongType),
+ Column.create("dt", StringType))
+ val partitions = Array[Transform](identity("dt"), identity("id"))
+
+ createTable("nt", cols, partitions)
+ // The scan reports its keys sorted on the full key: [('2020', 2), ('2021', 1)]. Narrowing them
+ // to `[id]` gives [2, 1], which is not sorted -- projecting a sorted sequence onto a subset of
+ // key positions does not preserve sortedness.
+ sql("INSERT INTO testcat.ns.nt VALUES (2, '2020'), (1, '2021')")
+
+ withTable("t1") {
+ sql("CREATE TABLE t1 (id BIGINT, data STRING) USING parquet")
+ sql("INSERT INTO t1 VALUES (1, 'a'), (2, 'b')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ // `length(dt) > 2` is not pushed down, so `dt` reaches the scan while the projection above
+ // it drops the column, narrowing the KeyedPartitioning to `[identity(id)]`.
+ val df = sql(
+ """
+ |SELECT nt.id, t1.data
+ |FROM (SELECT id FROM testcat.ns.nt WHERE length(dt) > 2) nt
+ |JOIN t1 ON nt.id = t1.id
+ |""".stripMargin)
+ val plan = df.queryExecution.executedPlan
+ // Only the v1 side is shuffled, onto the narrowed KeyedPartitioning of the v2 side, and
+ // nothing re-groups the v2 side -- so the shuffle is the only thing that can align them.
+ assert(collectShuffles(plan).length == 1)
+ assert(collectGroupPartitions(plan).isEmpty)
+ checkAnswer(df, Seq(Row(1L, "a"), Row(2L, "b")))
+ }
+ }
+ }
+
+ test("SPARK-59022: keyed shuffle follows the declared partition key order over a union") {
+ val cols = Array(
+ Column.create("id", LongType),
+ Column.create("data", StringType))
+ val partitions = Array[Transform](identity("id"))
+
+ createTable("nt1", cols, partitions)
+ createTable("nt2", cols, partitions)
+ // `UnionExec` concatenates its children's partition keys in child order, so the merged keys are
+ // [3, 4] ++ [1, 2]. They are unique, hence grouped, and no projection is involved, hence not
+ // narrowed -- but they are not sorted.
+ sql("INSERT INTO testcat.ns.nt1 VALUES (3, 'c'), (4, 'd')")
+ sql("INSERT INTO testcat.ns.nt2 VALUES (1, 'a'), (2, 'b')")
+
+ withTable("t1") {
+ sql("CREATE TABLE t1 (id BIGINT, x STRING) USING parquet")
+ sql("INSERT INTO t1 VALUES (1, 'x'), (2, 'x'), (3, 'x'), (4, 'x')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(
+ """
+ |SELECT u.id, u.data, t1.x
+ |FROM (SELECT * FROM testcat.ns.nt1 UNION ALL SELECT * FROM testcat.ns.nt2) u
+ |JOIN t1 ON u.id = t1.id
+ |""".stripMargin)
+ val plan = df.queryExecution.executedPlan
+ // Only the v1 side is shuffled, onto the union's KeyedPartitioning, and nothing re-groups
+ // the union -- so the shuffle is the only thing that can align them.
+ assert(collectShuffles(plan).length == 1)
+ assert(collectGroupPartitions(plan).isEmpty)
+ checkAnswer(df, Seq(
+ Row(1L, "a", "x"), Row(2L, "b", "x"), Row(3L, "c", "x"), Row(4L, "d", "x")))
+ }
+
+ // The plan-shape assertions above need AQE off, but the wrong results were live in the
+ // default configuration, so pin the answer with AQE on as well.
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ checkAnswer(
+ sql(
+ """
+ |SELECT u.id, u.data, t1.x
+ |FROM (SELECT * FROM testcat.ns.nt1 UNION ALL SELECT * FROM testcat.ns.nt2) u
+ |JOIN t1 ON u.id = t1.id
+ |""".stripMargin),
+ Seq(Row(1L, "a", "x"), Row(2L, "b", "x"), Row(3L, "c", "x"), Row(4L, "d", "x")))
+ }
+ }
+ }
+
test("SPARK-57881: storage-partitioned join leverages union output KeyedPartitioning to " +
"avoid shuffle") {
val cols = Array(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org