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 21:33:24 UTC
(spark) branch branch-4.3 updated: [SPARK-59025][SQL] Fix subset join key grouping when the keyed side reports a PartitioningCollection
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 7cc91948c378 [SPARK-59025][SQL] Fix subset join key grouping when the keyed side reports a PartitioningCollection
7cc91948c378 is described below
commit 7cc91948c378406bf133fade6facf811c59a75a3
Author: Dongjoon Hyun <do...@apache.org>
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 <do...@apache.org>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
(cherry picked from commit 25bae648b5b2a3a1e97c8011bc6be76cf01efd68)
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
.../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 8ca344e51194..d55c515ab782 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 7d335f285f50..f79e57ec1f87 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
@@ -2578,6 +2578,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: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org