[SPARK-58974][SQL] Apply the narrowed-partitioning skew guard regardless of requireAllClusterKeysForDistribution

[peter-toth]

What changes were proposed in this pull request?

KeyedPartitioning.groupedSatisfies refuses a narrowed, non-grouped partitioning unless spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled is set, because grouping it would merge partitions that held distinct keys in the original finer-grained partitioning. That guard lived inside the requireAllClusterKeys = false arm, so it never ran when spark.sql.requireAllClusterKeysForDistribution was enabled.

This moves the guard above the requireAllClusterKeys check. The skew risk does not depend on which key sets count as matching, so the decision should not either.

The config doc for allowKeysSubsetOfPartitionKeys.enabled is updated as well: it did not mention this second thing the config gates, which is not tied to requireAllClusterKeysForDistribution being false.

Why are the changes needed?

With spark.sql.requireAllClusterKeysForDistribution = true a narrowed partitioning was accepted, EnsureRequirements inserted a GroupPartitionsExec, and the skew exposure was taken with the opt-in config still off. Note this is the stricter of the two settings, which makes it the more surprising direction.

Reproduced with a table partitioned by (id, dept) projected down to id -- keys collapsing to [1, 1, 2] -- and joined on id: with requireAllClusterKeys = true the plan gets a GroupPartitionsExec and no shuffle while allowKeysSubsetOfPartitionKeys is off, whereas with requireAllClusterKeys = false the same query correctly shuffles both sides.

Please note that the guard's condition is imprecise, and the hoist makes that imprecision reachable for one more config value. isNarrowed && !isGrouped is a proxy for "the narrowing collapsed distinct keys", but !isGrouped has causes that have nothing to do with narrowing: a source that reports several splits per partition key, or a union whose children have distinct keys individually and repeat keys across children. Grouping such a partitioning merges only same-key partitions, which is what GroupPartitionsExec does for any non-narrowed partitioning and needs no opt-in, yet the condition refuses it. Until now that false refusal could only happen with requireAllClusterKeysForDistribution = false; after this change it can happen with either value. Tightening the condition to actual key collapse is a separate change, which I am working on as a follow-up; this PR keeps the condition as it is and only fixes where it is evaluated.

Does this PR introduce any user-facing change?

Yes, a plan-level change. With requireAllClusterKeysForDistribution enabled, Spark previously coalesced partitions derived from a narrowed partitioning without allowKeysSubsetOfPartitionKeys.enabled, risking skewed partitions; it now inserts a shuffle unless that config is enabled. Query results are unchanged.

No migration guide entry: the guard, and with it the bypass, arrived in 4.3.0, which is unreleased, so no released version behaves the old way. This goes to branch-4.3 as well.

How was this patch tested?

• New unit test in ProjectedOrderingAndPartitioningSuite calling groupedSatisfies directly on a narrowed, ungrouped partitioning, for both values of requireAllClusterKeys and both values of the opt-in. It fails on master with kp.groupedSatisfies(required) was true for requireAllClusterKeys = true.
• New end-to-end test in KeyGroupedPartitioningSuite asserting the guard behaves identically for both values of requireAllClusterKeys: no GroupPartitionsExec, and a shuffle instead. It fails on master for requireAllClusterKeys = true. It also varies v2BucketingShuffleEnabled, because that decides whether the refused side is laid out on the other side's declared partition keys (one shuffle) or both sides are shuffled (two).
• That test also covers allowKeysSubsetOfPartitionKeys = true for both values of requireAllClusterKeys, since this is the first time the opt-in affects groupedSatisfies when requireAllClusterKeysForDistribution is enabled -- it must restore the coalescing and avoid the shuffles. Neither suite had any coverage of requireAllClusterKeysForDistribution, which is how the bug survived.
• Also ran DistributionSuite, KeyGroupedPartitioningSuite, EnsureRequirementsSuite, PlannerSuite, ProjectedOrderingAndPartitioningSuite, DataFrameSetOperationsSuite, AdaptiveQueryExecSuite, CoalesceShufflePartitionsSuite, and the TPC-DS plan stability suites -- no golden file changed.

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

Generated-by: Claude Code (Opus 5)

[peter-toth]

cc @dongjoon-hyun

[peter-toth]

cc @ulysses-you

[ulysses-you]

0 blocking, 1 non-blocking.

Verified

Confirmed by tracing every call site and by running the tests locally:

• groupedSatisfies has exactly two callers: EnsureRequirements.scala:78 (nonGrouped.find(_.groupedSatisfies(distribution)), i.e. always invoked on a non-grouped KP) and KeyedPartitioning.satisfies0 (partitioning.scala:574, guarded by isGrouped &&). So the delta surface of this move is exactly one route: requireAllClusterKeys=true + narrowed + ungrouped.
• For spark.sql.requireAllClusterKeysForDistribution=false this is provably a no-op: the old check order was allowSubset? -> narrowed guard -> attributes.forall, the new order is guard(&& !allowSubset) -> allowSubset? -> attributes.forall. Conjoining the guard with !allowSubset makes the reorder equivalent.
• Via satisfies0 the hoisted !isGrouped clause can never fire (guarded upstream by isGrouped &&), so the OrderedDistribution arm and all other behavior are untouched.
• Ran the unit test against unpatched master: fails with kp.groupedSatisfies(required) was true for requireAllClusterKeys=true, matching the description verbatim; with the patch applied, both new tests pass. Also ran KeyGroupedPartitioningSuite, EnsureRequirementsSuite and ProjectedOrderingAndPartitioningSuite with the patch - no failure attributable to it.
• Neither direction introduces a wrong-results axis: accepting groups partitions by equal projected key values (rows still routed correctly); refusing falls back to a standard hash shuffle. So plan shape / skew only, as the description states.

Finding 1 (non-blocking): the imprecise condition widens a safe-case regression to requireAllClusterKeysForDistribution=true users

isNarrowed && !isGrouped is a proxy for "the narrowing collapsed distinct keys", but !isGrouped also holds when a source natively reports multiple splits per partition key (e.g. uncommitted Iceberg files). Concretely: table partitioned by (a, b, c) whose scan emits 2 splits per key, then a projection keeping exactly the operator's clustering (a, b) without collapsing any projected key. Grouping here would merge only same-key partitions - precisely what GroupPartitionsExec does for any un-narrowed ungrouped KP - yet the guard now refuses it under requireAllClusterKeysForDistribution=true and inserts an avoidable shuffle (results identical, plan worse).

Peer code showing the intended tolerance for duplicate keys: createKeyedShuffleSpec.allClusterKeysCovered in EnsureRequirements.scala (~line 794-809), whose comment says "Key order and duplicated cluster keys don't matter" for the same skew concern. After this change, groupedSatisfies tolerates duplicated keys strictly less than that adjacent co-partitioning gate does.

Agreed on deferring the tightened condition to a follow-up rather than expanding this diff - but the follow-up should reconcile with allClusterKeysCovered too, so both SPJ skew gates end up sharing one notion of "collapsed keys". Non-blocking given results are unaffected; just flagging that the follow-up shouldn't be dropped.

[ulysses-you]

[Non-blocking] isNarrowed && !isGrouped is a proxy for "the narrowing collapsed distinct keys", but !isGrouped also holds when a source natively reports multiple splits per partition key (e.g. uncommitted Iceberg files under an (a, b, c) partitioning). In that case a projection keeping exactly the operator's clustering (a, b) need not collapse any projected key, and grouping would merge only same-key partitions - precisely what GroupPartitionsExec does for any un-narrowed ungrouped KP - yet this guard now refuses it even for requireAllClusterKeysForDistribution=true, inserting an avoidable shuffle (results unchanged).

Peer code showing the intended tolerance: createKeyedShuffleSpec.allClusterKeysCovered in EnsureRequirements.scala (~lines 794-809) says "Key order and duplicated cluster keys don't matter" for the same skew concern; after this change groupedSatisfies tolerates duplicated keys strictly less than that adjacent co-partitioning gate.

OK to defer tightening to the follow-up, but please make sure it reconciles with allClusterKeysCovered so both SPJ skew gates share one notion of "collapsed keys".

[peter-toth]

Agreed, and thanks for the allClusterKeysCovered pointer - its comment ("Key order and duplicated cluster keys don't matter") is exactly the tolerance I want this guard to have. The follow-up is already written: the flag will mean actual key collapse - a projection mapping keys that were distinct in the input onto the same projected key - instead of "positions were dropped", so a source that reports several splits per partition key no longer trips it. Measured on your shape, with the opt-in off: before, no GroupPartitionsExec and 2 shuffles; after, 1 GroupPartitionsExec and no shuffle. I will check allClusterKeysCovered against the tightened notion explicitly and report what I find on that PR, and link it here.

[dongjoon-hyun]

[Non-blocking] The !isGrouped conjunct here is load-bearing in a way the class-level doc contradicts. The doc above ("groupedSatisfies(): called on non-grouped KPs ...") predates this change, but satisfies0 also calls groupedSatisfies on grouped KPs (isGrouped && groupedSatisfies(required)), and grouped-but-narrowed KPs are constructible: PartitioningPreservingUnaryExecNode recomputes isGrouped from the projected keys while isNarrowed stays sticky, which is exactly the state exercised by the existing test "SPARK-46367: narrowing projection with distinct projected keys does not require allowKeysSubsetOfPartitionKeys".

If a later cleanup trusts the class doc and drops the "redundant" !isGrouped, grouped narrowed KPs would stop satisfying ClusteredDistribution with the config off and pick up unnecessary shuffles (that test would fail). Could you touch up the class-level doc (or add a short note here) so the conjunct's purpose is recorded?

[peter-toth]

You are right, and this was the most useful comment on the PR - I had convinced myself the conjunct was redundant. It is not: nonGroupedSatisfies is the default Partitioning implementation, so a grouped KP reaches a ClusteredDistribution only through groupedSatisfies, and a narrowing projection that keeps its keys distinct is precisely a grouped narrowed KP. Dropping the term would cost it a shuffle, and the test you name pins that. Fixed the class-level doc to describe both callers and to record the reason, and added a note at the guard itself.

[dongjoon-hyun]

[Nit] "GroupPartitionsExec will merge partitions that held distinct keys in the original finer-grained partitioning" states the risky case as fact, but when the source was already ungrouped (multiple splits per identical key) and the narrowing collapsed nothing, the merged partitions held the same original key. Since isNarrowed && !isGrouped is a proxy (as the PR description itself notes), "may merge" would be accurate here and in the isNarrowed scaladoc. As written, the docs make the guard look exact, which could make the promised follow-up tightening look unnecessary to a future reader.

[peter-toth]

Done, in the guard comment and in the isNarrowed scaladoc, together with the reason for "may": the duplicate keys can also come from a source reporting several splits per partition key.

[dongjoon-hyun]

[Nit] SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys is now read twice in this match arm (here and in the else branch below). Hoisting it into a local val at the top of the case c @ ClusteredDistribution(...) arm would make it explicit that both branches are gated by the same switch, which is the core of this fix, and would keep a future edit from decoupling the two reads.

[peter-toth]

Done - one allowKeysSubsetOfPartitionKeys val per ClusteredDistribution arm, so both branches visibly read the same switch.

[dongjoon-hyun]

[Nit] The directly analogous test "SPARK-46367: narrowing projection requires allowKeysSubsetOfPartitionKeys" builds this same shape from the suite's standing items/itemsColumns + purchases/purchasesColumns fixtures and the selectWithMergeJoinHint helper (it only generates the hint prefix from alias strings, so the subquery alias u works with it too -- the SPARK-46367 test itself passes "sub"). Reusing those here would keep the pre-fix and post-fix narrowing tests directly comparable instead of introducing a parallel t1/t2 schema.

[peter-toth]

Done. The test now builds on items/purchases and selectWithMergeJoinHint, with the same data as the neighbouring SPARK-46367: narrowing projection requires allowKeysSubsetOfPartitionKeys test (id=1 mapping to two partitions), so the two are directly comparable.

[dongjoon-hyun]

[Nit] The requireAll = false iteration exactly duplicates Scenario 1 of the existing test "SPARK-46367: narrowing projection with duplicate keys requires allowKeysSubsetOfPartitionKeys to satisfy ClusteredDistribution" (same keys, same ProjectExec fixture, same groupedSatisfies asserts; ClusteredDistribution's requireAllClusterKeys defaults to false). If the explicit both-values contrast pair is intentional, fine -- otherwise this test could cover only requireAll = true, or the loop could be folded into the existing test so the two don't have to be kept in lockstep.

[peter-toth]

Kept both values: "the guard answers the same either way" is the claim this test makes, and the false iteration is its control. I added a comment saying so, so the overlap with the SPARK-46367 scenario does not read as accidental.

[dongjoon-hyun]

+1, LGTM

[dongjoon-hyun]

Let's merge this, @peter-toth ! 😄

[peter-toth]

Merge Summary:

• merged into master 9d33e01
• merged into branch-4.x 7f888ce
• merged into branch-4.3 ed82c61

Posted by merge_spark_pr.py

[peter-toth]

Thank you @dongjoon-hyun and @ulysses-you for the reivew.

[peter-toth]

The allClusterKeysCovered reconciliation I promised here is reported on the follow-up: #58351 (comment)