From 266a19c58b5f31907663280860b4bee774ede58e Mon Sep 17 00:00:00 2001 From: Cheng Pan Date: Tue, 4 Aug 2026 21:11:24 +0800 Subject: [PATCH 1/4] [SPARK-58558][SQL] Remove requireAllClusterKeysForCoPartition as a gate for SPJ requireAllClusterKeysForCoPartition enforced an exact attribute-level match between V2 partition keys and join keys before allowing storage-partitioned join to skip shuffle. This was stricter than necessary: for value-based V2 partitioning, key order or extra join keys do not introduce skew. The real skew threshold (joining on a subset of partition keys) is already controlled by v2BucketingAllowKeysSubsetOfPartitionKeys. This PR replaces the config check in createKeyedShuffleSpec with a direct partitioning.satisfies(distribution) call, so SPJ is allowed by default whenever all partition keys appear in the join keys. The config still gates classic HashPartitioning shuffle reuse (V1 bucketing). Assisted-by: GLM 5.2 --- docs/sql-migration-guide.md | 4 ++ docs/sql-performance-tuning.md | 11 +-- .../apache/spark/sql/internal/SQLConf.scala | 4 +- .../exchange/EnsureRequirements.scala | 15 +--- .../KeyGroupedPartitioningSuite.scala | 20 ------ .../exchange/EnsureRequirementsSuite.scala | 68 ++++++++++--------- 6 files changed, 46 insertions(+), 76 deletions(-) diff --git a/docs/sql-migration-guide.md b/docs/sql-migration-guide.md index f25295ad5349e..252605ae030a7 100644 --- a/docs/sql-migration-guide.md +++ b/docs/sql-migration-guide.md @@ -22,6 +22,10 @@ license: | * Table of contents {:toc} +## Upgrading from Spark SQL 4.3 to 4.4 + +- Since Spark 4.4, `spark.sql.requireAllClusterKeysForCoPartition` no longer affects storage-partitioned joins (V2 data sources). A shuffle is now avoided whenever all partition keys appear in the join keys, regardless of order; joining on a subset of partition keys remains controlled by `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled`. Users who previously set `spark.sql.requireAllClusterKeysForCoPartition` to `false` solely to enable storage-partitioned joins no longer need to do so. The config still applies to hash-partitioned children (e.g., V1 bucketing). + ## Upgrading from Spark SQL 4.2 to 4.3 - Since Spark 4.3, [ASOF JOIN](sql-ref-syntax-qry-select-asof-join.html) is available as an opt-in SQL feature gated by `spark.sql.join.asofJoin.enabled` (default `false`). When disabled, `ASOF JOIN` fails at parse time with `UNSUPPORTED_FEATURE.ASOF_JOIN`. diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md index 6acb36413d937..fd3d2a37711d9 100644 --- a/docs/sql-performance-tuning.md +++ b/docs/sql-performance-tuning.md @@ -531,14 +531,6 @@ The following SQL properties enable Storage Partition Join in different join que 3.4.0 - - spark.sql.requireAllClusterKeysForCoPartition - true - - When true, require the join or MERGE keys to be same and in the same order as the partition keys to eliminate shuffle. Hence, set to false in this situation to eliminate shuffle. - - 3.4.0 - spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled false @@ -551,7 +543,7 @@ The following SQL properties enable Storage Partition Join in different join que spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled false - When enabled, try to avoid shuffle if join or MERGE condition does not include all partition columns. This config requires both spark.sql.sources.v2.bucketing.enabled and spark.sql.sources.v2.bucketing.pushPartValues.enabled to be true, and spark.sql.requireAllClusterKeysForCoPartition to be false. + When enabled, try to avoid shuffle if join or MERGE condition does not include all partition columns. This config requires both spark.sql.sources.v2.bucketing.enabled and spark.sql.sources.v2.bucketing.pushPartValues.enabled to be true. 4.0.0 @@ -607,7 +599,6 @@ ON t.dep = s.dep AND t.id = s.id SET 'spark.sql.sources.v2.bucketing.enabled' 'true' SET 'spark.sql.iceberg.planning.preserve-data-grouping' 'true' SET 'spark.sql.sources.v2.bucketing.pushPartValues.enabled' 'true' -SET 'spark.sql.requireAllClusterKeysForCoPartition' 'false' SET 'spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled' 'true' -- Plan with Storage Partition Join diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index d060d5f7a7f99..16d854f8a4da8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -1105,7 +1105,9 @@ object SQLConf { .doc("When true, the planner requires all the clustering keys as the hash partition keys " + "of the children, to eliminate the shuffles for the operator that needs its children to " + "be co-partitioned, such as JOIN node. This is to avoid data skews which can lead to " + - "significant performance regression if shuffles are eliminated.") + "significant performance regression if shuffles are eliminated. This only applies to " + + "hash partitioning (e.g., V1 bucketing); storage-partitioned joins over V2 data " + + "sources are not affected.") .version("3.3.0") .booleanConf .createWithDefault(true) 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 c632b3d841e61..57c48a1af0c9c 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 @@ -782,20 +782,7 @@ case class EnsureRequirements( partitioning: Partitioning, distribution: ClusteredDistribution): Option[KeyedShuffleSpec] = { def tryCreate(partitioning: KeyedPartitioning): Option[KeyedShuffleSpec] = { - // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one - // attribute per partition expression. - val attributes = partitioning.expressions.flatMap(_.references) - val clustering = distribution.clustering - - val satisfies = if (SQLConf.get.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION)) { - attributes.length == clustering.length && attributes.zip(clustering).forall { - case (l, r) => l.semanticEquals(r) - } - } else { - partitioning.satisfies(distribution) - } - - if (satisfies) { + if (partitioning.satisfies(distribution)) { Some(partitioning.createShuffleSpec(distribution).asInstanceOf[KeyedShuffleSpec]) } else { None 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 ceaac9729ddd1..8de4150c22bb9 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 @@ -1057,7 +1057,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq((true, true, 5), (false, true, 3), (true, false, 7), (false, false, 5)).foreach { case (partial, filter, expected) => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> false.toString, SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> filter.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> @@ -1107,7 +1106,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(("true", 5), ("false", 5)).foreach { case (enable, expected) => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> false.toString, SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { val df = createJoinTestDF( @@ -1153,7 +1151,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(("true", 5), ("false", 5)).foreach { case (enable, expected) => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> false.toString, SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { val df = createJoinTestDF( @@ -1199,7 +1196,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with enable <- Seq("true", "false") } yield { withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> false.toString, SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { // The left side uses a key-grouped partitioning to satisfy the WINDOW function's @@ -1767,7 +1763,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { pushDownValues => Seq(true, false).foreach { partiallyClustered => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> partiallyClustered.toString, @@ -1830,7 +1825,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { partiallyClustered => Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> partiallyClustered.toString, @@ -1972,7 +1966,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> @@ -2132,7 +2125,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> @@ -2246,7 +2238,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> @@ -2310,7 +2301,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", @@ -2375,7 +2365,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach{ allowPushDown => Seq(true, false).foreach{ partiallyClustered => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> allowPushDown.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> partiallyClustered.toString, @@ -2433,7 +2422,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { partiallyClustered => Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> @@ -2496,7 +2484,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> partiallyClustered.toString, @@ -2708,7 +2695,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with "(6, 50.0, cast('2023-02-01' as timestamp))") withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", @@ -3178,7 +3164,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with "(1, 10.0, cast('2021-05-20' as timestamp))") withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true") { val metrics = runAndFetchMetrics { @@ -3210,7 +3195,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with test("SPARK-55411: Fix ArrayIndexOutOfBoundsException when join keys " + "are less than cluster keys") { withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", @@ -3286,7 +3270,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with test("SPARK-55535: Multi table join granular partition grouping") { withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { val items_partitions = Array(identity("id"), years("arrive_time")) @@ -3482,7 +3465,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with createTable(purchases, purchasesColumns, purchases_partitions) sql(s"INSERT INTO testcat.ns.$purchases VALUES (2, 10.0, cast('2021-01-01' as timestamp))") withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { val df = sql( @@ -3825,7 +3807,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { preserveOrdering => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> preserveOrdering.toString) { @@ -3976,7 +3957,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with "(1, 20.0, cast('2022-03-10' as timestamp))") withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true" ) { val hashDf = sql( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 17d00ec055e07..94763e4fd2ca6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -35,7 +35,6 @@ import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoin import org.apache.spark.sql.execution.python.FlatMapCoGroupsInPandasExec import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{IntegerType, StructField, StructType} @@ -862,28 +861,9 @@ class EnsureRequirementsSuite extends SharedSparkSession { assert(right.expressions === Seq(bucket(4, exprA), years(exprC))) case other => fail(other.toString) } - - // by default spark.sql.requireAllClusterKeysForCoPartition is true, so when there isn't - // exact match on all partition keys, Spark will fallback to shuffle. - plan1 = new DummySparkPlanWithBatchScanChild( - outputPartitioning = KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprB) :: Nil, Seq.empty) - ) - plan2 = new DummySparkPlanWithBatchScanChild( - outputPartitioning = KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprC) :: Nil, Seq.empty) - ) - smjExec = SortMergeJoinExec( - exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - EnsureRequirements.apply(smjExec) match { - case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => - assert(left.expressions === Seq(exprA, exprB, exprB)) - assert(right.expressions === Seq(exprA, exprC, exprC)) - case other => fail(other.toString) - } } - test(s"KeyedPartitioning with ${REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key} = false") { + test("KeyedPartitioning with subset of join keys") { var plan1 = new DummySparkPlanWithBatchScanChild( outputPartitioning = KeyedPartitioning(bucket(4, exprB) :: years(exprC) :: Nil, Seq.empty) ) @@ -894,7 +874,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { // simple case var smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => @@ -912,7 +892,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => @@ -930,7 +910,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { KeyedPartitioning(years(exprA) :: bucket(4, exprC) :: days(exprA) :: Nil, Seq.empty)) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => @@ -949,7 +929,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprB :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => @@ -967,7 +947,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => @@ -985,7 +965,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => @@ -1006,7 +986,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => @@ -1016,6 +996,32 @@ class EnsureRequirementsSuite extends SharedSparkSession { } } + test("KeyedPartitioning: SPJ is not affected by requireAllClusterKeysForCoPartition") { + // requireAllClusterKeysForCoPartition only gates HashPartitioning shuffle reuse, + // not V2 KeyedPartitioning. SPJ is allowed with either value when partition keys + // are a subset of join keys. + val plan1 = new DummySparkPlanWithBatchScanChild( + outputPartitioning = + KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprB) :: Nil, Seq.empty)) + val plan2 = new DummySparkPlanWithBatchScanChild( + outputPartitioning = + KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprC) :: Nil, Seq.empty)) + val smjExec = SortMergeJoinExec( + exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) + Seq("true", "false").foreach { requireAllKeys => + withSQLConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> requireAllKeys) { + EnsureRequirements.apply(smjExec) match { + case SortMergeJoinExec(_, _, _, _, + SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), + SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => + assert(left.expressions === Seq(bucket(4, exprA), bucket(4, exprB))) + assert(right.expressions === Seq(bucket(4, exprA), bucket(4, exprC))) + case other => fail(s"Expected no shuffle, but got: $other") + } + } + } + } + test("SPARK-41413: check compatibility when partition values mismatch") { withSQLConf(SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true") { val leftPartValues = Seq(Array[Any](1, 1), Array[Any](2, 2)).map(new GenericInternalRow(_)) @@ -1034,7 +1040,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { // simple case var smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, GroupPartitionsExec(DummySparkPlan(_, _, left: KeyedPartitioning, _, _), @@ -1058,7 +1064,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, GroupPartitionsExec(DummySparkPlan(_, _, left: PartitioningCollection, _, _), @@ -1093,7 +1099,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, GroupPartitionsExec(DummySparkPlan(_, _, left: PartitioningCollection, _, _), From bdafe4d46b4ad3ea809b612bcb714d0b2d7daaa1 Mon Sep 17 00:00:00 2001 From: Cheng Pan Date: Tue, 4 Aug 2026 22:40:50 +0800 Subject: [PATCH 2/4] Check key coverage instead of removing the gate Assisted-by: Claude Fable 5 --- docs/sql-migration-guide.md | 2 +- docs/sql-performance-tuning.md | 10 +++++++- .../apache/spark/sql/internal/SQLConf.scala | 6 ++--- .../exchange/EnsureRequirements.scala | 14 ++++++++++- .../KeyGroupedPartitioningSuite.scala | 7 ++++++ .../exchange/EnsureRequirementsSuite.scala | 25 +++++++++++-------- 6 files changed, 48 insertions(+), 16 deletions(-) diff --git a/docs/sql-migration-guide.md b/docs/sql-migration-guide.md index 252605ae030a7..5809cec700c12 100644 --- a/docs/sql-migration-guide.md +++ b/docs/sql-migration-guide.md @@ -24,7 +24,7 @@ license: | ## Upgrading from Spark SQL 4.3 to 4.4 -- Since Spark 4.4, `spark.sql.requireAllClusterKeysForCoPartition` no longer affects storage-partitioned joins (V2 data sources). A shuffle is now avoided whenever all partition keys appear in the join keys, regardless of order; joining on a subset of partition keys remains controlled by `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled`. Users who previously set `spark.sql.requireAllClusterKeysForCoPartition` to `false` solely to enable storage-partitioned joins no longer need to do so. The config still applies to hash-partitioned children (e.g., V1 bucketing). +- Since Spark 4.4, for storage-partitioned joins (V2 data sources), `spark.sql.requireAllClusterKeysForCoPartition` no longer requires the join keys to exactly match the partition keys: duplicated join keys no longer prevent shuffle elimination, and `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` no longer additionally requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false`. As before, when the partition keys cover only part of the join keys, eliminating the shuffle still requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false`. ## Upgrading from Spark SQL 4.2 to 4.3 diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md index fd3d2a37711d9..8e544d44a5324 100644 --- a/docs/sql-performance-tuning.md +++ b/docs/sql-performance-tuning.md @@ -531,6 +531,14 @@ The following SQL properties enable Storage Partition Join in different join que 3.4.0 + + spark.sql.requireAllClusterKeysForCoPartition + true + + When true, storage-partitioned join requires every join or MERGE key to be covered by the partition keys, ignoring key order and duplicated keys, to eliminate shuffle. When the partition keys cover only part of the join or MERGE keys, set to false to eliminate shuffle, at the risk of data skew and reduced parallelism from the coarser storage partitioning. + + 3.3.0 + spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled false @@ -540,7 +548,7 @@ The following SQL properties enable Storage Partition Join in different join que 3.4.0 - spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled + spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled false When enabled, try to avoid shuffle if join or MERGE condition does not include all partition columns. This config requires both spark.sql.sources.v2.bucketing.enabled and spark.sql.sources.v2.bucketing.pushPartValues.enabled to be true. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 16d854f8a4da8..fde6f6cbb5a14 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -1105,9 +1105,9 @@ object SQLConf { .doc("When true, the planner requires all the clustering keys as the hash partition keys " + "of the children, to eliminate the shuffles for the operator that needs its children to " + "be co-partitioned, such as JOIN node. This is to avoid data skews which can lead to " + - "significant performance regression if shuffles are eliminated. This only applies to " + - "hash partitioning (e.g., V1 bucketing); storage-partitioned joins over V2 data " + - "sources are not affected.") + "significant performance regression if shuffles are eliminated. For V2 data source " + + "partitioning (storage-partitioned join), the check ignores key order and duplicated " + + "clustering keys: it requires every clustering key to be covered by the partition keys.") .version("3.3.0") .booleanConf .createWithDefault(true) 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 57c48a1af0c9c..bed87c7d9d60a 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 @@ -782,7 +782,19 @@ case class EnsureRequirements( partitioning: Partitioning, distribution: ClusteredDistribution): Option[KeyedShuffleSpec] = { def tryCreate(partitioning: KeyedPartitioning): Option[KeyedShuffleSpec] = { - if (partitioning.satisfies(distribution)) { + // The config requires all the cluster keys to be covered by the partition keys, to avoid + // the skew of joining on keys that are coarser than the join keys. Key order and duplicated + // cluster keys don't matter. + def allClusterKeysCovered: Boolean = { + // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one + // attribute per partition expression. + val attributes = partitioning.expressions.flatMap(_.references) + distribution.clustering.forall(c => attributes.exists(_.semanticEquals(c))) + } + + if (partitioning.satisfies(distribution) && + (!SQLConf.get.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION) || + allClusterKeysCovered)) { Some(partitioning.createShuffleSpec(distribution).asInstanceOf[KeyedShuffleSpec]) } else { None 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 8de4150c22bb9..ca3d53e2ea372 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 @@ -1057,6 +1057,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq((true, true, 5), (false, true, 3), (true, false, 7), (false, false, 5)).foreach { case (partial, filter, expected) => withSQLConf( + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> false.toString, SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> filter.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> @@ -1106,6 +1107,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(("true", 5), ("false", 5)).foreach { case (enable, expected) => withSQLConf( + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> false.toString, SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { val df = createJoinTestDF( @@ -1151,6 +1153,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(("true", 5), ("false", 5)).foreach { case (enable, expected) => withSQLConf( + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> false.toString, SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { val df = createJoinTestDF( @@ -1763,6 +1766,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { pushDownValues => Seq(true, false).foreach { partiallyClustered => withSQLConf( + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> partiallyClustered.toString, @@ -3164,6 +3168,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with "(1, 10.0, cast('2021-05-20' as timestamp))") withSQLConf( + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true") { val metrics = runAndFetchMetrics { @@ -3807,6 +3812,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { preserveOrdering => withSQLConf( + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> preserveOrdering.toString) { @@ -3957,6 +3963,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with "(1, 20.0, cast('2022-03-10' as timestamp))") withSQLConf( + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true" ) { val hashDf = sql( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 94763e4fd2ca6..8e78603f5a071 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -871,10 +871,15 @@ class EnsureRequirementsSuite extends SharedSparkSession { outputPartitioning = KeyedPartitioning(bucket(4, exprC) :: years(exprB) :: Nil, Seq.empty) ) - // simple case + // simple case: join key exprA is not covered by either side's partition keys, so by default + // the coverage check of requireAllClusterKeysForCoPartition falls back to shuffle to avoid + // joining on a partitioning coarser than the join keys var smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner, None, plan1, plan2) - EnsureRequirements.apply(smjExec) match { + assert(EnsureRequirements.apply(smjExec) + .collect { case s: ShuffleExchangeLike => s }.length == 2) + // with requireAllClusterKeysForCoPartition=false, SPJ is allowed + applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => @@ -929,7 +934,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprB :: exprC :: Nil, Inner, None, plan1, plan2) - EnsureRequirements.apply(smjExec) match { + applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => @@ -996,10 +1001,10 @@ class EnsureRequirementsSuite extends SharedSparkSession { } } - test("KeyedPartitioning: SPJ is not affected by requireAllClusterKeysForCoPartition") { - // requireAllClusterKeysForCoPartition only gates HashPartitioning shuffle reuse, - // not V2 KeyedPartitioning. SPJ is allowed with either value when partition keys - // are a subset of join keys. + test("KeyedPartitioning: duplicated join keys do not block SPJ") { + // The coverage check of requireAllClusterKeysForCoPartition ignores key order and + // duplicated cluster keys: join keys [a, b, b] are fully covered by partition keys + // on [a, b], so SPJ is allowed with either config value. val plan1 = new DummySparkPlanWithBatchScanChild( outputPartitioning = KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprB) :: Nil, Seq.empty)) @@ -1040,7 +1045,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { // simple case var smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner, None, plan1, plan2) - EnsureRequirements.apply(smjExec) match { + applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, GroupPartitionsExec(DummySparkPlan(_, _, left: KeyedPartitioning, _, _), @@ -1064,7 +1069,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner, None, plan1, plan2) - EnsureRequirements.apply(smjExec) match { + applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, GroupPartitionsExec(DummySparkPlan(_, _, left: PartitioningCollection, _, _), @@ -1099,7 +1104,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner, None, plan1, plan2) - EnsureRequirements.apply(smjExec) match { + applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, GroupPartitionsExec(DummySparkPlan(_, _, left: PartitioningCollection, _, _), From 25c15cb1a06614767aa3e76d4d2efe6abd3ca4df Mon Sep 17 00:00:00 2001 From: Cheng Pan Date: Tue, 4 Aug 2026 23:58:07 +0800 Subject: [PATCH 3/4] Restore incompatible-buckets test override, state V2-only coverage scope Assisted-by: Claude Fable 5 --- .../src/main/scala/org/apache/spark/sql/internal/SQLConf.scala | 3 ++- .../spark/sql/connector/KeyGroupedPartitioningSuite.scala | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index fde6f6cbb5a14..5154d15f65683 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -1107,7 +1107,8 @@ object SQLConf { "be co-partitioned, such as JOIN node. This is to avoid data skews which can lead to " + "significant performance regression if shuffles are eliminated. For V2 data source " + "partitioning (storage-partitioned join), the check ignores key order and duplicated " + - "clustering keys: it requires every clustering key to be covered by the partition keys.") + "clustering keys: it requires every clustering key to be covered by the partition keys; " + + "hash partitioning deliberately keeps the exact match (same keys in the same order).") .version("3.3.0") .booleanConf .createWithDefault(true) 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 ca3d53e2ea372..fdd3847b63d38 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 @@ -2242,6 +2242,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> From d3b42c82a80ff0d1d1245b2243b7a18d4e4a6b1c Mon Sep 17 00:00:00 2001 From: Cheng Pan Date: Wed, 5 Aug 2026 14:15:59 +0800 Subject: [PATCH 4/4] Reword docs around reachable shapes, add multi-transform SPJ tests Assisted-by: Claude Fable 5 --- docs/sql-migration-guide.md | 2 +- docs/sql-performance-tuning.md | 2 +- .../apache/spark/sql/internal/SQLConf.scala | 9 ++++---- .../KeyGroupedPartitioningSuite.scala | 22 ++++++++++++++++++ .../exchange/EnsureRequirementsSuite.scala | 23 +++++++++++++++---- 5 files changed, 48 insertions(+), 10 deletions(-) diff --git a/docs/sql-migration-guide.md b/docs/sql-migration-guide.md index 5809cec700c12..6e4b76ae2caf6 100644 --- a/docs/sql-migration-guide.md +++ b/docs/sql-migration-guide.md @@ -24,7 +24,7 @@ license: | ## Upgrading from Spark SQL 4.3 to 4.4 -- Since Spark 4.4, for storage-partitioned joins (V2 data sources), `spark.sql.requireAllClusterKeysForCoPartition` no longer requires the join keys to exactly match the partition keys: duplicated join keys no longer prevent shuffle elimination, and `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` no longer additionally requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false`. As before, when the partition keys cover only part of the join keys, eliminating the shuffle still requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false`. +- Since Spark 4.4, for storage-partitioned joins, `spark.sql.requireAllClusterKeysForCoPartition` requires every join key to be covered by some partition key instead of matching the partition keys positionally. As a result, a join-key column partitioned by more than one transform no longer prevents shuffle elimination, and `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` no longer additionally requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false` when the join keys are a subset of the partition keys. As before, when the partition keys cover only part of the join keys, eliminating the shuffle still requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false`. ## Upgrading from Spark SQL 4.2 to 4.3 diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md index 8e544d44a5324..cad32a1c77e8d 100644 --- a/docs/sql-performance-tuning.md +++ b/docs/sql-performance-tuning.md @@ -535,7 +535,7 @@ The following SQL properties enable Storage Partition Join in different join que spark.sql.requireAllClusterKeysForCoPartition true - When true, storage-partitioned join requires every join or MERGE key to be covered by the partition keys, ignoring key order and duplicated keys, to eliminate shuffle. When the partition keys cover only part of the join or MERGE keys, set to false to eliminate shuffle, at the risk of data skew and reduced parallelism from the coarser storage partitioning. + When true, storage-partitioned join requires every join or MERGE key to be covered by some partition key (rather than matching the partition keys positionally) to eliminate shuffle. When the partition keys cover only part of the join or MERGE keys, set to false to eliminate shuffle, at the risk of data skew and reduced parallelism from the coarser storage partitioning. 3.3.0 diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 5154d15f65683..6bfe147731696 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -1105,10 +1105,11 @@ object SQLConf { .doc("When true, the planner requires all the clustering keys as the hash partition keys " + "of the children, to eliminate the shuffles for the operator that needs its children to " + "be co-partitioned, such as JOIN node. This is to avoid data skews which can lead to " + - "significant performance regression if shuffles are eliminated. For V2 data source " + - "partitioning (storage-partitioned join), the check ignores key order and duplicated " + - "clustering keys: it requires every clustering key to be covered by the partition keys; " + - "hash partitioning deliberately keeps the exact match (same keys in the same order).") + "significant performance regression if shuffles are eliminated. For storage-partitioned " + + "join, every clustering key must be covered by some partition key, rather than matching " + + "the partition keys positionally, so a column partitioned by more than one transform " + + "does not prevent shuffle elimination; hash partitioning deliberately keeps the " + + "positional match.") .version("3.3.0") .booleanConf .createWithDefault(true) 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 fdd3847b63d38..a732294786b2c 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 @@ -4483,4 +4483,26 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } } + + test("SPARK-58558: SPJ on a join key column partitioned by multiple transforms") { + // Partition expressions outnumber the join keys because `id` is partitioned twice, but + // every join key is covered, so SPJ works with default configs. + val items_partitions = Array(bucket(8, "id"), identity("id")) + createTable(items, itemsColumns, items_partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + "(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + "(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") + + val purchases_partitions = Array(bucket(8, "item_id"), identity("item_id")) + createTable(purchases, purchasesColumns, purchases_partitions) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + "(1, 42.0, cast('2020-01-01' as timestamp)), " + + "(2, 19.5, cast('2020-02-01' as timestamp))") + + val df = createJoinTestDF(Seq("id" -> "item_id")) + val shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.isEmpty, "should not contain any shuffle") + checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(2, "bb", 10.0, 19.5))) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 8e78603f5a071..c6610c8c5c860 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -924,6 +924,19 @@ class EnsureRequirementsSuite extends SharedSparkSession { case other => fail(other.toString) } + // a column partitioned by more than one transform: partition expressions outnumber the + // join keys, but every join key is covered, so SPJ is allowed with default configs + smjExec = SortMergeJoinExec( + exprA :: exprB :: Nil, exprA :: exprC :: Nil, Inner, None, plan1, plan2) + EnsureRequirements.apply(smjExec) match { + case SortMergeJoinExec(_, _, _, _, + SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), + SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => + assert(left.expressions === Seq(years(exprA), bucket(4, exprB), days(exprA))) + assert(right.expressions === Seq(years(exprA), bucket(4, exprC), days(exprA))) + case other => fail(other.toString) + } + // invalid case: partitioning key positions don't match plan1 = new DummySparkPlanWithBatchScanChild( outputPartitioning = KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprB) :: Nil, Seq.empty) @@ -1001,10 +1014,12 @@ class EnsureRequirementsSuite extends SharedSparkSession { } } - test("KeyedPartitioning: duplicated join keys do not block SPJ") { - // The coverage check of requireAllClusterKeysForCoPartition ignores key order and - // duplicated cluster keys: join keys [a, b, b] are fully covered by partition keys - // on [a, b], so SPJ is allowed with either config value. + test("KeyedPartitioning: duplicated join keys in hand-built plans do not block SPJ") { + // Queries produce this key list only in unusual configurations: BooleanSimplification + // normally dedups the conjunction, but it is an excludable rule + // (spark.sql.optimizer.excludedRules), and EnsureRequirements must also stay robust + // for hand-built or rewritten plans. The coverage check treats duplicated cluster + // keys as covered, so SPJ is allowed with either config value. val plan1 = new DummySparkPlanWithBatchScanChild( outputPartitioning = KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprB) :: Nil, Seq.empty))