This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 99aed2a98525 [SPARK-59022][SQL] Keyed shuffle must follow the declared
partition key order
99aed2a98525 is described below
commit 99aed2a98525f6cb564e7b0a4063e87a9b4af6d9
Author: Peter Toth <[email protected]>
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 <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
.../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 426ad8e658fb..5e529a37f8bc 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
@@ -4246,6 +4246,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: [email protected]
For additional commands, e-mail: [email protected]
