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 25bae648b5b2 [SPARK-59025][SQL] Fix subset join key grouping when the
keyed side reports a PartitioningCollection
25bae648b5b2 is described below
commit 25bae648b5b2a3a1e97c8011bc6be76cf01efd68
Author: Dongjoon Hyun <[email protected]>
AuthorDate: Wed Aug 26 14:28:39 2026 -0700
[SPARK-59025][SQL] Fix subset join key grouping when the keyed side reports
a PartitioningCollection
### What changes were proposed in this pull request?
Fix the post-shuffle regrouping in `EnsureRequirements` for
storage-partitioned joins with
`spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys=true`. The
unwrap that pushes join key positions down to the keyed side only matches a
bare `KeyedShuffleSpec`:
```scala
bestSpecOpt match {
case Some(KeyedShuffleSpec(_, _, Some(joinKeyPositions))) =>
withJoinKeyPositions(child, joinKeyPositions)
case _ => child
}
```
When the keyed child's `outputPartitioning` is a `PartitioningCollection`
(e.g. an inner SPJ join result, or a projection duplicating a partition column
under multiple aliases like `SELECT a AS a1, a AS a2, b`), its spec is a
`ShuffleSpecCollection` wrapping the `KeyedShuffleSpec`, so the match falls
through and no `GroupPartitionsExec` is inserted. This PR unwraps a (possibly
nested) `ShuffleSpecCollection` to its head spec — the same spec
`ShuffleSpecCollection.createPartitioning` d [...]
### Why are the changes needed?
Without the fix, the keyed side keeps its N ungrouped partitions while the
other side is shuffled into the projected layout with M partitions, and
planning fails with:
```
java.lang.IllegalArgumentException: requirement failed:
All KeyedPartitionings in a PartitioningCollection must have equal
partitionKeys
```
### Does this PR introduce _any_ user-facing change?
The affected queries used to fail with the error above; they now run
correctly, shuffling only the non-keyed side.
### How was this patch tested?
Pass the CIs.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Fable 5
Closes #58315 from dongjoon-hyun/SPARK-59025.
Authored-by: Dongjoon Hyun <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
.../execution/exchange/EnsureRequirements.scala | 16 ++++++-
.../connector/KeyGroupedPartitioningSuite.scala | 53 ++++++++++++++++++++++
2 files changed, 67 insertions(+), 2 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
index 52a761090b0b..c2051ac91b06 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
@@ -17,6 +17,7 @@
package org.apache.spark.sql.execution.exchange
+import scala.annotation.tailrec
import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
@@ -248,13 +249,16 @@ case class EnsureRequirements(
child
case ((child, dist), idx) =>
if (bestSpecOpt.isDefined &&
bestSpecOpt.get.isCompatibleWith(specs(idx))) {
- bestSpecOpt match {
+ // If the child's partitioning is a `PartitioningCollection`, its
spec is a
+ // `ShuffleSpecCollection` whose `createPartitioning` delegates to
the head spec,
+ // so unwrap to the head spec to stay aligned with the re-shuffled
side below.
+ unwrapSpecCollection(bestSpecOpt.get) match {
// If `areChildrenCompatible` is false, we can still perform SPJ
// by shuffling the other side based on join keys (see the else
case below).
// Hence we need to ensure that after this call, the
outputPartitioning of the
// partitioned side's BatchScanExec is grouped by join keys to
match,
// and we do that by pushing down the join keys
- case Some(KeyedShuffleSpec(_, _, Some(joinKeyPositions))) =>
+ case KeyedShuffleSpec(_, _, Some(joinKeyPositions)) =>
withJoinKeyPositions(child, joinKeyPositions)
case _ => child
}
@@ -760,6 +764,14 @@ case class EnsureRequirements(
}
}
+ // Unwraps a `ShuffleSpecCollection` (possibly nested) to the spec that its
+ // `createPartitioning` delegates to, i.e. the head spec.
+ @tailrec
+ private def unwrapSpecCollection(spec: ShuffleSpec): ShuffleSpec = spec
match {
+ case ShuffleSpecCollection(specs) => unwrapSpecCollection(specs.head)
+ case other => other
+ }
+
/**
* Applies join key positions to a plan by wrapping or updating
GroupPartitionsExec.
*/
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 5e529a37f8bc..64c0d0f2bc90 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
@@ -2572,6 +2572,59 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
}
}
+ test("SPARK-59025: shuffle one side and join keys are less than partition
keys " +
+ "when the keyed side reports a PartitioningCollection") {
+ val items_partitions = Array(identity("id"), identity("name"))
+ createTable(items, itemsColumns, items_partitions)
+
+ // 4 distinct (id, name) partition keys but only 3 distinct ids, so
grouping by the join
+ // key must reduce the keyed side from 4 partitions to 3.
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+ "(1, 'ab', 30.0, cast('2020-01-02' as timestamp)), " +
+ "(3, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
+ "(4, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+
+ createTable(purchases, purchasesColumns, Array.empty)
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ "(1, 42.0, cast('2020-01-01' as timestamp)), " +
+ "(1, 89.0, cast('2020-01-03' as timestamp)), " +
+ "(3, 19.5, cast('2020-02-01' as timestamp)), " +
+ "(5, 26.0, cast('2023-01-01' as timestamp))")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ // Duplicating `id` under two aliases makes the projection report a
+ // `PartitioningCollection` of `KeyedPartitioning`s, so the keyed side's
shuffle spec
+ // is a `ShuffleSpecCollection` wrapping a `KeyedShuffleSpec` with join
key positions.
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("i", "p")}
+ |id1, id2, name, i.price AS purchase_price, p.price AS sale_price
+ |FROM (SELECT id AS id1, id AS id2, name, price FROM
testcat.ns.$items) i
+ |JOIN testcat.ns.$purchases p ON i.id1 = p.item_id
+ |ORDER BY id1, purchase_price, sale_price
+ |""".stripMargin)
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.size == 1, "only the non-keyed side should be shuffled")
+ val groupPartitions =
collectGroupPartitions(df.queryExecution.executedPlan)
+ assert(groupPartitions.size == 1 &&
groupPartitions.head.joinKeyPositions.isDefined,
+ "the keyed side should be grouped by the join keys")
+ assert(groupPartitions.head.outputPartitioning.numPartitions == 3,
+ "the keyed side should be grouped down to 3 partitions")
+ assert(shuffles.head.outputPartitioning.numPartitions == 3,
+ "the shuffled side should match the 3 grouped partitions")
+ checkAnswer(df, Seq(
+ Row(1, 1, "ab", 30.0, 42.0),
+ Row(1, 1, "ab", 30.0, 89.0),
+ Row(1, 1, "aa", 40.0, 42.0),
+ Row(1, 1, "aa", 40.0, 89.0),
+ Row(3, 3, "bb", 10.0, 19.5)))
+ }
+ }
+
test("SPARK-45652: SPJ should handle empty partition after dynamic
filtering") {
withSQLConf(
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
