[SPARK-59057][SQL] Make KeyedPartitioning.isNarrowed mean actual key collapse, and rename it to isCollapsed

[peter-toth]

What changes were proposed in this pull request?

KeyedPartitioning carries a flag that gates whether GroupPartitionsExec may coalesce its duplicate partition keys without spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled. Today the flag records provenance -- "a projection dropped key positions, or my input already had some dropped" -- while the gate's own comment describes collapse: "partitions that held distinct keys in the original finer-grained partitioning". The two are not the same, and the gate reads the wrong one.

This changes the flag to mean what the gate needs, and renames it from isNarrowed to isCollapsed: the partitioning is coarser than the layout it was derived from, because a projection mapped keys that were distinct in the input onto the same projected key. Dropping key positions no longer sets it on its own; the projected keys have to actually lose distinctness. It stays sticky, since neither grouping nor a further projection can make a partitioning finer again.

The gate keeps its two terms, isCollapsed && !isGrouped, and the code now says why: isCollapsed states that the coarsening happened, !isGrouped states that there is still something left for GroupPartitionsExec to merge. Once the keys are unique, grouping merges nothing and there is no further risk to gate, however coarse the partitioning already is.

Producers:

• PartitioningPreservingUnaryExecNode computes it as keySource.isCollapsed || (positions were dropped && projected distinct keys < input distinct keys). The two cheap terms come first so the distinct count is skipped for an inherited flag or a pass-through projection.
• UnionExec's keyed merge is unchanged apart from the rename -- it ORs the children's flags. It stops refusing the case where the children's keys merely overlap, because the children's flags are now precise; a child that really is coarsened still marks the union, as before.
• GroupPartitionsExec, KeyedPartitioning.toGrouped, KeyedPartitioning.createShuffleSpec and KeyedShuffleSpec.createPartitioning all propagate the flag instead of defaulting it to false through the 3-argument constructor. GroupPartitionsExec and createShuffleSpec project onto the operation keys, so they also compute their own coarsening. GroupPartitionsExec compares against the key count of its own side after projection and reduction, not against the aligned key list: that list is the one both join sides agreed on, so it can be missing keys this side had (partition filtering) or repeat them (padding), and neither is a collapse.

PartitioningCollection normalizes the flag across its members by OR, alongside the partitionKeys interning it already did, and its invariant check enforces that. The flag describes the shared physical layout rather than one member's naming of it, and consumers read it off a single member -- satisfies0 and EnsureRequirements accept when any one member satisfies the distribution -- so a member that under-reported the coarsening would let the gate through.

KeyedPartitioning gains a distinctKeyCount lazy val for this: free when the partitioning is grouped, and computed on demand otherwise. It is lazy so that the members of a PartitioningCollection, which share one key list, do not each force it when a consumer only needs one.

This supersedes #58316 (SPARK-59026), which restores the flag in two of the producers above. Propagating it is that PR's finding, and its unit test is taken here with credit. Its end-to-end test is not: its table has unique ids, so dropping the second key column keeps every key distinct, and that scenario is exactly what this change reclassifies as not coarsened. The two cannot both land as written.

The class doc also gains a section on coarsened partitionings, including why such a partitioning is kept rather than dropped to UnknownPartitioning -- that rationale was nowhere in the code.

Why are the changes needed?

Provenance over-refuses, and it does so on one of the commonest shapes. !isGrouped has causes that have nothing to do with narrowing:

• A data source reports one partition key per input split, so a table with several splits for the same partition value already has duplicate keys before any projection.
• UnionExec computes isGrouped over the concatenation of its children's keys, so two children that are individually key-distinct make it false by overlapping with each other.

In both cases grouping merges only partitions that already shared a key, which is what GroupPartitionsExec does for any partitioning that never went through a projection, and needs no opt-in. Today the mere presence of a narrowing projection anywhere upstream turns that into a refusal, so a plan gets a shuffle that protects nothing.

Measured on the multi-split shape: a table partitioned by (id, dept) with two splits for the same (1, 'x') value, projected down to id and joined on it, with the opt-in off. Before: no GroupPartitionsExec and 2 shuffles, and the projected partitioning reports the flag set. After: 1 GroupPartitionsExec, 0 shuffles, flag clear.

Does this PR introduce any user-facing change?

It should reach branch-4.3 and branch-4.x as well as master, since 4.3.0 is where the flag first ships.

Yes, a plan-level change: storage-partitioned operations now proceed without allowKeysSubsetOfPartitionKeys.enabled in the cases above, where they previously fell back to a shuffle. Query results are unchanged. No migration guide entry: the flag and its gate arrived in 4.3.0 (SPARK-46367), which is unreleased, so no released version behaves the old way.

Reducers are a second case: with allowCompatibleTransforms, a bucket(4, id) side reduced onto a bucket(2, id) join really does end up coarser than its source, so its GroupPartitionsExec output now carries the flag where it did not before, and a coalescing further up the plan needs the opt-in. That direction is a narrowing, not a widening, and it is what the flag is supposed to say.

Planning cost was measured on the worst case: an ungrouped source (so the distinct count is a real pass), every hop dropping a position (so the cheap term does not short-circuit) and no hop collapsing a key (so the inherited flag does not either). 20 evaluations of a 10-hop chain over a 50k-split, 25k-key partitioning: 1020 ms before, 1738 ms after, i.e. about 3.6 ms per narrowing hop on top of the distinct pass isGrouped already needs. A pass-through hop pays nothing, since it cannot coarsen anything.

How was this patch tested?

• New unit test in ProjectedOrderingAndPartitioningSuite for the multi-split shape: a source with duplicate keys, projected down, is ungrouped but not collapsed, and groupedSatisfies accepts it with the opt-in off.
• New end-to-end test in KeyGroupedPartitioningSuite for the same shape through a real plan, asserting the grouping happens and the shuffles disappear.
• New unit test that a coarsened member anywhere in a PartitioningCollection marks every projected partitioning, so the outcome does not depend on which join side the coarsening came from.
• New end-to-end test that reducing keys onto a coarser transform reports the coarsening: with allowCompatibleTransforms, an identity(item_id) side reduced onto bucket(4, id) really is coarser than its source, and the flag now says so.
• New end-to-end test that filtering partition keys out is not a collapse: with partitionFilter on, an inner join plans both sides on the intersection of their keys, and the side that lost a key must not report itself coarsened -- otherwise the sticky flag costs a shuffle above a later union.
• One existing expectation flipped, which is the contract change: in SPARK-46367: narrowing projection with duplicate keys ..., the scenario whose projected keys stay distinct now asserts the flag is clear.
• Every new expectation is guarded by an ablation, verified one at a time: the old provenance formula fails the multi-split unit test, its end-to-end counterpart and the flipped SPARK-46367 scenario; the head-only inherited flag fails the PartitioningCollection test; comparing against the aligned key list instead of this side's own count fails the partition-filter test, where the shuffle count goes from 0 to 1.
• Also ran DistributionSuite, KeyGroupedPartitioningSuite, ProjectedOrderingAndPartitioningSuite, EnsureRequirementsSuite, PlannerSuite, 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)

Closes #58316

[peter-toth]

cc @dongjoon-hyun -- this is the follow-up to #58316 I mentioned there.

One note on reading the diff: this PR is based on #58338 (SPARK-58974), which is still open, so the first two commits here belong to that PR. Only the last commit belongs to this one. I will rebase onto master once #58338 lands.

[dongjoon-hyun]

Thanks for the detailed writeup. I read through the main changes and the tests, and I think the core direction is right: the gate's comment described collapse while the flag recorded provenance, and distinct-after < distinct-before is exactly the right predicate -- it is true iff two keys that were distinct in the input map onto the same projected key. Filling in the producers that were laundering the flag (toGrouped, createShuffleSpec, KeyedShuffleSpec.createPartitioning, GroupPartitionsExec) is a good pickup from #58316 too.

One property worth noting, because it makes the new class doc read consistently: a projection preserves the partition count, so distinct-after < distinct-before <= numPartitions means a freshly computed isCollapsed always implies !isGrouped. The !isGrouped term in the gate therefore only does work for a partitioning that inherited a sticky isCollapsed and was then re-grouped -- by GroupPartitionsExec or by reducing keys onto a coarser transform -- which is precisely what the new doc says. Good.

A few things below.

1. One of this PR's own producers breaks the doc rule the PR adds

The class doc gains:

A producer that projects or merges the members has to take isCollapsed from all of them, not from one

but GroupPartitionsExec.outputPartitioning projects the members and reads the flag per member:

val isCollapsed = k.isCollapsed || projectedDistinctKeyCount < k.distinctKeyCount

The comment right above it says "There can be multiple KeyedPartitionings in an output partitioning of a join", so a collection here is an anticipated input. Meanwhile PartitioningPreservingUnaryExecNode does kps.exists(_.isCollapsed) in the same situation. The two producers disagree, and the new unit test ("isCollapsed is taken from every member of a PartitioningCollection") pins only the ProjectExec side.

2. isCollapsed looks like a property of the physical layout, not of a member

Members of a PartitioningCollection describe the same physical partitioning and share the partitionKeys reference. If the left leg was coarsened, an output partition really does cover several partitions of the finer left layout, whichever side's expressions you name it by. So members disagreeing on the flag seems like the anomaly rather than something to preserve.

It also leaves the gate bypassable. EnsureRequirements uses:

val nonGroupedSatisfiesWhenGrouped = nonGrouped.find(_.groupedSatisfies(distribution))

find, so an ungrouped collection holding one coarsened and one plain member is accepted via the plain one and gets a GroupPartitionsExec. The reachable path is narrow (a join over a join, made ungrouped by a padding GroupPartitionsExec), so I do not think it blocks this PR -- but the doc presents it as intended ("the satisfaction path is separate and accepts when any single member does"), and to me it reads more like a remaining hole than a design choice.

Suggestion: normalize isCollapsed by OR across members in PartitioningCollection.fromPartitionings / checkKeyedPartitioningInvariant, the same way partitionKeys references are interned. That would remove the producer disagreement in (1) and this bypass at once, and let the doc drop both caveats.

3. KeyedShuffleSpec.createPartitioning -- question

The shuffled side inherits partitioning.isCollapsed. Under the new definition ("coarser than the layout it was derived from") the shuffled side has no finer layout it was derived from: its key-1 partition holding all of its key-1 rows is just what a HashPartitioning would give. The direction is safe (it can only add shuffles), but the flag is sticky, so it can cost a refusal further up with nothing behind it. Is this deliberate conservatism, or is the intent to widen the definition? Either way it would help to say so at the call site, since the comment there argues from the coarsened side's keys rather than from this side's history.

4. Comment volume (nit)

The 11-line comment above the single false in groupedSatisfies overlaps heavily with the new "Coarsened Partitionings" section in the class doc. Trimming the inline one to a few lines and leaving the rest in the doc would be easier to read. The doc section itself is worth having -- especially the rationale for keeping a coarsened partitioning instead of dropping it to UnknownPartitioning, which was nowhere in the code before.

Checked, no issues

• No isNarrowed references left; every producer goes through the 4-arg constructor.
• The isGrouped short-circuit in distinctKeyCount, and the evaluation order in PartitioningPreservingUnaryExecNode (inherited flag -> position dropped -> distinct count), are both correct.
• Comparing against this side's own key count rather than the aligned key list is well argued (filtering is pruning, padding is repetition), and the new end-to-end test pins it.
• The doc claim that OrderedDistribution is not gated checks out: that path goes through distributePartitions = true, where padTo gives one partition per split and nothing is coalesced.

[dongjoon-hyun]

BTW,

• I removed Co-authored-by: Dongjoon Hyun <dongjoon@apache.org> from this PR. :)
• Also, added Closes #58316 explicitly.

[peter-toth]

Thanks for the review. You are right on (2), and it reverses a call I made earlier: I had left the members free to disagree, reasoning that each records its own side's history. But they describe one physical layout and share the key list, so if any side is coarse an output partition really does cover several finer ones, whichever member's expressions name it. And the find in EnsureRequirements makes that concrete -- a mixed collection reaching a GroupPartitionsExec through the plain member is a hole, not a design choice.

So PartitioningCollection.fromPartitionings now normalizes isCollapsed by OR, alongside the partitionKeys interning, and checkKeyedPartitioningInvariant enforces it. Both stay O(members) per level: one representative per member is enough, since every collection agrees internally by the same construction, and a nested collection is only descended into when it disagrees -- the property your interning code deliberately protects for linearly-nested same-key joins.

That answers (1) too: the flag is now computed once in GroupPartitionsExec.outputPartitioning, outside the transform, so the two producers no longer disagree. Both doc caveats are gone.

On (3): deliberate conservatism, and I have written that at the call site. The shuffled side has no finer layout of its own -- its partitions are what a hash partitioning would give -- but the two sides are co-located on one key set, and a later grouping of that key set carries the coarsened side's risk. It can only add shuffles, never remove one.

(4) done: the gate comment is trimmed to what the class doc does not already carry.

Also added a unit test that a coarsened member marks the whole collection, including through a nested one. Pushed as a separate commit.

[szehon-ho]

this is just collapsed right? is another new term necessary? is 'grouping' ok?

On that note, wdyt we have a quick glossary somewhere for the terms we are defining here?

• Key collapse: A projection or reduction maps two distinct old keys to the same new key.
  (1,A), (1,B) -> 1, 1
• Grouping: Physically combines partitions that now share the same key.
  1, 1, 2 -> 1, 2

Maybe its me but the javadoc is a bit hard to read , i think example is worth 1000 words