[SPARK-58558][SQL] Make requireAllClusterKeysForCoPartition check key coverage instead of exact match for SPJ

[pan3793]

What changes were proposed in this pull request?

For storage-partitioned joins (V2 KeyedPartitioning), change the spark.sql.requireAllClusterKeysForCoPartition check in EnsureRequirements.createKeyedShuffleSpec from an exact match (same keys, same order, same count) to a coverage check: every join key must be covered by some partition key, ignoring key order and duplicated keys.

The config still gates HashPartitioning shuffle reuse (e.g., V1 bucketing) with the exact-match semantics; extending the coverage semantics to hash partitioning is deliberately out of scope for this PR. Its SQLConf doc, the SPJ tuning guide, and the migration guide are updated accordingly; the tuning guide also switches to the current name of spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled (the row used the deprecated alias).

Why are the changes needed?

The exact-match check was stricter than needed in two safe cases:

1. Duplicated join keys (e.g. ON t1.a = t2.a AND t1.b = t2.c AND t1.b = t2.c, with t1 partitioned by (a, b) and t2 by (a, c)) blocked SPJ because of the count mismatch, although the duplicated key is still covered by each side's partition keys and introduces no skew.

2. v2BucketingAllowKeysSubsetOfPartitionKeys=true additionally required requireAllClusterKeysForCoPartition=false even when the join keys are a strict subset of the partition keys, a case that always passes the coverage check -- the extra keys are partition keys, which coverage does not restrict. When the join keys only partially overlap the partition keys, requireAllClusterKeysForCoPartition=false is still required, as before.

The genuinely risky case -- partition keys covering only part of the join keys, where eliminating the shuffle drops the join's parallelism to the storage layout's partition count -- remains gated by the config, as before.

Does this PR introduce any user-facing change?

Yes. SPJ now works by default with duplicated join keys, and when the join keys are a subset of the partition keys, spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled=true no longer needs spark.sql.requireAllClusterKeysForCoPartition=false alongside it. As before, when the partition keys cover only part of the join keys, SPJ still requires spark.sql.requireAllClusterKeysForCoPartition=false.

How was this patch tested?

• EnsureRequirementsSuite: added a test for duplicated join keys under both config values; the partition-keys-cover-part-of-join-keys cases assert shuffle by default and SPJ with requireAllClusterKeysForCoPartition=false.
• KeyGroupedPartitioningSuite: removed the requireAllClusterKeysForCoPartition=false overrides that are no longer needed; kept them where the join keys are not fully covered by the partition keys.
• ShuffleSpecSuite: unchanged (covers HashShuffleSpec, which keeps the exact-match semantics).
• All three suites pass.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Fable 5)

[pan3793]

cc @peter-toth

[peter-toth]

Thanks for the PR, @pan3793!

Two of the problems you point at are real: a duplicated join key (ON a = a2 AND b = b2 AND b = b3) blocks SPJ today, and having to set requireAllClusterKeysForCoPartition=false alongside allowKeysSubsetOfPartitionKeys=true is a genuine annoyance. But dropping the gate isn't equivalent to fixing those. partitioning.satisfies(distribution) on a grouped KeyedPartitioning reduces to "every partition key is one of the join keys" (partitioning.scala:595), which is also true when the partition keys are a strict subset of the join keys - the coarser-than-the-join-key case the gate was added to block, and the one case allowKeysSubsetOfPartitionKeys does not cover (it guards the opposite direction). Meanwhile the key-order case the description leads with already works on master, because reorderJoinKeysRecursively rewrites the join keys into partition-key order before the gate ever runs. So the net effect of the patch as written is to make the skew case the default, with no config left to switch it off. Finding 1 has a smaller change that gets both of your wins without that.

Blocking

• 1. Dropping the gate turns on SPJ when the partition keys cover only part of the join keys: that is the skew case the config exists for, nothing else puts a floor under the join's parallelism, and after this change there is no targeted way back. Keeping the gate but making it order-insensitive ("every cluster key is covered by some partition key") gets you the duplicated-join-key case and the "had to set both configs" case without this one. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:785]
• 2. "Why are the changes needed" doesn't hold as written: join-key order is already handled by reorderJoinPredicates (EnsureRequirements.scala:904 -> :412-425, added in the same SPARK-37377 commit as the gate), so the gate never blocked it; and the gate isn't a no-op - your own test rename says so, EnsureRequirementsSuite.scala:866 is now called "KeyedPartitioning with subset of join keys" and its first case joins on a while neither side is partitioned on a. Extra join keys do change the outcome relative to the plan that would otherwise be picked (a shuffle on all the join keys), and that is where the skew comes from.

Non-blocking

• 3. Migration-guide entry states only the upside: it should name the parallelism trade-off and how to get the old plan back, since there is no longer a targeted switch. [inline: docs/sql-migration-guide.md:27]

Minor

• 4. Tuning-guide row still uses the deprecated config alias: allowJoinKeysSubsetOfPartitionKeys vs the current allowKeysSubsetOfPartitionKeys that your migration-guide entry uses. [inline: docs/sql-performance-tuning.md:543]

[peter-toth]

Finding 1. For a grouped KeyedPartitioning and a join's ClusteredDistribution (whose requireAllClusterKeys comes from requireAllClusterKeysForDistribution, default false), partitioning.satisfies(distribution) reduces to

attributes.forall(x => requiredClustering.exists(_.semanticEquals(x)))

(sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:595) - every partition key must be a join key, but not the other way round. So it is also true when the partition keys are a strict subset of the join keys, i.e. when the storage layout groups rows more coarsely than the join key does. That is exactly what the removed gate blocked, and what the config was added for: "This is to avoid data skews which can lead to significant performance regression if shuffles are eliminated" (SQLConf.scala:1105) - the same wording as the sibling check you keep for hash partitioning: "To avoid potential data skew, we don't allow HashShuffleSpec to create partitioning if the hash partition keys are not the full join keys" (partitioning.scala:1138).

v2BucketingAllowKeysSubsetOfPartitionKeys does not cover this. It guards the opposite direction - join keys being a subset of the partition keys (partitioning.scala:583-587).

Concretely, with two Iceberg tables PARTITIONED BY (days(ts)):

SELECT * FROM t JOIN s ON t.ts = s.ts AND t.id = s.id

on master both sides shuffle to spark.sql.shuffle.partitions; with this patch SPJ fires and the join runs with one task per day, each holding a whole day of rows. Nothing else puts a floor under it: shouldConsiderMinParallelism / defaultNumShufflePartitions only apply to the bestSpecOpt branch (EnsureRequirements.scala:192-200), which is skipped entirely once areChildrenCompatible is true (:246). And there is no targeted way back afterwards - only turning SPJ off completely (spark.sql.sources.v2.bucketing.enabled=false) or requireAllClusterKeysForDistribution=true, which also changes aggregate/window planning.

Your own test says as much: EnsureRequirementsSuite.scala:866 is now named "KeyedPartitioning with subset of join keys", and its first case joins on [a, b, c] while neither side is partitioned on a.

If the goal is to stop the check caring about key order, the smaller fix is to keep the gate but make it order-insensitive, which is what the config's name says anyway:

    def tryCreate(partitioning: KeyedPartitioning): Option[KeyedShuffleSpec] = {
      // 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 = {
        // 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
      }
    }

That keeps both wins from the description:

• the duplicated-join-key case works at defaults - a duplicated cluster key is still covered by a partition key - so your new test at :999 and the two "duplicated keys" cases in :866 still pass;
• allowKeysSubsetOfPartitionKeys=true no longer needs requireAllClusterKeysForCoPartition=false alongside it, because that config only ever adds extra partition keys, which the coverage check doesn't look at.

Only the first case of :866 (join key a present on neither side's partitioning) goes back to needing the config set to false, which is the case I'd argue should stay opt-in.

If you do want coarser-than-join-key SPJ on by default, that's a bigger call than "removing a redundant gate" - worth saying so plainly in the description and the migration guide, and I'd expect a config to switch it off.

[pan3793]

Adopted the suggested coverage check in bdafe4d. The partition-keys-cover-part-of-join-keys case stays gated by the config; duplicated join keys and allowKeysSubsetOfPartitionKeys=true no longer need requireAllClusterKeysForCoPartition=false.

[peter-toth]

Finding 3. This entry only states the upside. Two things a reader whose join got slower after upgrading would need:

• the trade-off: when the partition keys cover only part of the join keys, the join now runs with the storage layout's partition count instead of spark.sql.shuffle.partitions - a table partitioned on days(ts) and joined on (ts, id) loses parallelism and can get badly skewed (finding 1);
• how to get the old plan back. spark.sql.requireAllClusterKeysForCoPartition no longer does it, and the only remaining switches are spark.sql.sources.v2.bucketing.enabled=false (turns SPJ off completely) or spark.sql.requireAllClusterKeysForDistribution=true (also changes aggregate and window planning).

If finding 1 is addressed by keeping the gate order-insensitive instead, this entry should say that key order and duplicated join keys no longer matter, and that allowKeysSubsetOfPartitionKeys no longer needs requireAllClusterKeysForCoPartition=false alongside it.

[pan3793]

Rewrote the entry in bdafe4d: it states the coverage semantics and that the partition-keys-cover-part-of-join-keys case still requires setting the config to false.

[pan3793]

Thanks for the review, @peter-toth! All findings are addressed in bdafe4d: the gate is kept and changed to a coverage check per your suggestion, and the PR/JIRA titles are updated accordingly.

On finding 2: right, reorderJoinPredicates already handles key order, so the gate never blocked reordered keys. The description is rewritten around the two real cases: duplicated join keys, and allowKeysSubsetOfPartitionKeys=true no longer needing requireAllClusterKeysForCoPartition=false alongside it.

[pan3793]

a small doc fix, the config spark.sql.requireAllClusterKeysForCoPartition is introduced in 3.3.0

[peter-toth]

Re-checked through bdafe4d4 — findings 1, 2, 3, 4 resolved (gate kept and turned into a coverage check, so the partition-keys-cover-only-part-of-the-join-keys case stays opt-in; description, migration entry and config name restated around that), nothing regressed. The Since Version 3.3.0 fix matches SQLConf.scala:1111.

Edited after a closer look: finding 6 is rewritten (my first version named the wrong shape as the one this enables), findings 8 and 9 are added, and finding 7 is withdrawn — HashShuffleSpec keeping the positional match has no reachable consequence, since a hash partitioning can't repeat a column the way a V2 partitioning can and key order is already normalized by reorderJoinPredicates.

Blocking

• 5. SPARK-47094: Does not trigger when incompatible number of buckets stops testing bucket incompatibility (new): its tables are partitioned on store_id only while the join is on (store_id, dept_id), so the coverage check now rejects both sides before the bucket counts are ever compared — the assert(shuffles.nonEmpty) passes for the gate's reason instead. Keep the requireAllClusterKeysForCoPartition=false override here, as you did in the 7 tests you restored. [inline: sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:2245]
• 6. Neither case 1 of "Why are the changes needed?" nor the new SQLConf/docs wording describes a shape a query can produce (new): duplicated cluster keys can't reach the planner (BooleanSimplification dedups the conjunction before ExtractEquiJoinKeys builds the join keys) and an asymmetric duplicate dies on areKeysCompatible's equal-expression-count check, so the only reason the new test can assert otherwise is that it builds the SortMergeJoinExec by hand. Key order doesn't survive either: reorderJoinPredicates already normalizes it. What is left is the count mismatch in the other direction — a column partitioned by more than one transform — plus case 2, and since every one of those was already reachable with requireAllClusterKeysForCoPartition=false, the accurate motivation is that they no longer require switching a global skew guard off. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:1110]
• 8. The default-config half of the change has no test a query can reach (new): the 8 override removals that were load-bearing on base (:1829, :2305, :2426, :2488, :2699, :3200, :3277, :3470) all cover case 2, so that win is properly tested — but every one of them needs allowKeysSubsetOfPartitionKeys=true. Nothing covers the shape that newly works with no config: every partition attribute is a join key, but the partition expressions outnumber the join keys because a column is partitioned twice. Dropping the duplicated join key from the (years(a), bucket(4, b), days(a)) block at EnsureRequirementsSuite:917 turns it into exactly that, and it still fails on base. [inline: sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:1015]

Non-blocking

• 9. "How was this patch tested?" reads the override removals as if all of them follow from this change (new): 4 of the 13 (:1199, :1970, :2129, :2369) already passed base's positional check — same number of partition expressions as join keys, in order — so those overrides were dead before this PR and dropping them is unrelated cleanup, while :2242 went the other way and should have been kept (finding 5). Worth splitting the two groups in that bullet (breakdown in the finding-8 comment) so 13 removals aren't read as 13 demonstrations of the new semantics.

[peter-toth]

Finding 5. Both tables here are partitioned by bucket(N, store_id) only, and the join is ON t1.store_id = t2.store_id AND t1.dept_id = t2.dept_id, so dept_id is not covered by either side's partition keys: allClusterKeysCovered is false and createKeyedShuffleSpec returns None (EnsureRequirements.scala:795) before areKeysCompatible compares bucket(2, store_id) with bucket(3, store_id). V2_BUCKETING_SHUFFLE_ENABLED is off here, so KeyedShuffleSpec.canCreatePartitioning is false and the bestSpecOpt branch is skipped too — the bucket-count comparison the test is named for now runs nowhere. assert(shuffles.nonEmpty) still holds, so CI stays green, but the test would keep passing even if (2, 3) bucket counts became compatible.

Note this is specific to the coverage check, not to dropping the override: in the previous revision satisfies alone passed here and the bucket check was reached. This is the same partition-keys-cover-only-part-of-the-join-keys shape for which you restored the override in the 7 other tests, so it belongs here as well:

Suggested change

	SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
	SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
	SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",

[peter-toth]

Finding 6. (Replaces my earlier finding-7 comment on this line — that one is withdrawn, see the review body.)

"the check ignores key order and duplicated clustering keys" names two relaxations that no query reaches:

• Key order was already handled before this gate ran: reorderJoinPredicates runs first in the same transformUp case, and reorderJoinKeysRecursively has explicit KeyedPartitioning cases that permute the join keys into partition-key order (EnsureRequirements.scala:412-426). In the one situation it can't fix — the two sides' partition orders disagreeing relative to the join pairing — areKeysCompatible rejects the pair anyway, gate or no gate.
• Duplicated clustering keys don't survive the optimizer: BooleanSimplification splits the whole conjunction and dedups it through an ExpressionSet (expressions.scala:458-493, its own comment being (a && b) && a && (a && c) => a && b && c), so ON t1.a = t2.a AND t1.b = t2.c AND t1.b = t2.c collapses to two key pairs before ExtractEquiJoinKeys sees it. An asymmetric duplicate does survive the optimizer but dies in the planner: coverage is applied per side, so ON t1.a = t2.a AND t1.b = t2.b AND t1.b = t2.c needs c covered on t2, and areKeysCompatible then rejects 2 vs 3 partition expressions (partitioning.scala:1304).

What the check does relax is the partition side: a partitioning may now have more expressions than the join has key positions. Both sides PARTITIONED BY (bucket(8, id), truncate(4, id)) joined ON t1.id = t2.id gives cluster keys [id] against partition attrs [id, id] — the old attributes.length == clustering.length failed, coverage passes, and it is safe because the partition tuple is still a function of the join key. That, plus the allowKeysSubsetOfPartitionKeys direction, is the whole win.

One framing note while you're rewriting it. This is a pure relaxation of the config-true path, so nothing it enables was previously impossible — all of it was reachable by setting requireAllClusterKeysForCoPartition=false. What it buys is that you no longer have to switch that guard off globally, which would also admit the genuinely risky partition-keys-cover-part-of-the-join-keys case, to get two shapes that carry no skew risk: join keys a subset of the partition columns, and a join-key column partitioned by more than one transform. That reads as a stronger motivation than the duplicated-join-key story, and it is what the tests actually demonstrate.

The current wording also appears in docs/sql-migration-guide.md:27 ("duplicated join keys no longer prevent shuffle elimination"), docs/sql-performance-tuning.md:538 ("ignoring key order and duplicated keys") and case 1 of the PR description, so all four want the same correction. For this one (as a plain block, since the anchor can't reach line 1109):

      .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), 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.")

[peter-toth]

Sorry @pan3793 , I didn't review my latest comments thoroughly, let me iterate on it.

[peter-toth]

Finding 8. This key list can't come from a query: BooleanSimplification dedups the conjunction before ExtractEquiJoinKeys builds the join keys, so the pair (t1.b, t2.c) cannot appear twice (details in finding 6 on SQLConf.scala:1110).

The KeyGroupedPartitioningSuite removals do cover the other win properly. Going through all 13: 8 of them (:1829, :2305, :2426, :2488, :2699, :3200, :3277, :3470) have more partition expressions than join keys — (identity(id), identity(data)) joined on data alone, and so on — so base failed attributes.length == clustering.length and they genuinely needed the =false override; 4 (:1199, :1970, :2129, :2369) already matched positionally on base, so those overrides were dead and dropping them changes nothing either way (finding 9).

But all 8 also need allowKeysSubsetOfPartitionKeys=true, so nothing covers the shape that newly works with no config at all: every partition attribute is a join key, yet the partition expressions outnumber the join keys because a column is partitioned twice. That case is one edit away in this file — the block at :909-925 already partitions on (years(a), bucket(4, b), days(a)):

    smjExec = SortMergeJoinExec(
      exprA :: exprB :: Nil, exprA :: exprC :: Nil, Inner, None, plan1, plan2)

attrs [a, b, a] against cluster keys [a, b]: base fails the length check, coverage passes, the existing left.expressions/right.expressions assertions still hold, and reorderJoinPredicates doesn't interfere (reorder bails on the 3-vs-2 size mismatch). An end-to-end case is worth having too, since that is where reachability actually gets exercised: both sides PARTITIONED BY (bucket(8, id), identity(id)) joined ON t1.id = t2.id gives cluster keys [id] against partition attrs [id, id] — shuffle on base, SPJ with this patch, and both transforms are supported by InMemoryBaseTable.

[peter-toth]

@pan3793 , I've just updated my review comments, now they look good.