[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, meaning "a projection dropped key positions, or my input already had some dropped". The gate's own comment describes something else, a 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. A collapse is a projection or a reduction mapping keys that were distinct in the input onto the same key, so one partition here stands for several of the original ones. Dropping key positions no longer sets the flag 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 still rests on the same two conditions, but each now sits where it belongs. isCollapsed says the collapse happened, and it gates mayGroupToSatisfy, the question EnsureRequirements asks before it inserts a GroupPartitionsExec. The second condition is that there is still something left to merge, and that one is the caller. EnsureRequirements asks the question only of partitionings that are not grouped. Once the keys are unique, grouping merges nothing and there is nothing left to gate, however coarse the partitioning already is, so satisfies needs no gate at all.

That splits the old groupedSatisfies in two, because its two callers were asking different questions through one method:

• keysSatisfy(required): whether the partition keys match what the distribution asks for, ignoring duplicate keys. satisfies is this plus isGrouped.
• mayGroupToSatisfy(required): keysSatisfy plus a check that coalescing the duplicates is allowed. EnsureRequirements asks this of non-grouped partitionings.

That check is a conjunct rather than a case inside the key matching. SPARK-58974 established it there by hoisting it above the requireAllClusterKeys branch, since the risk it guards does not depend on which key sets count as matching. OrderedDistribution stays exempt, because GroupPartitionsExec pads that path out to the expected split counts instead of coalescing.

Producers:

• PartitioningPreservingUnaryExecNode goes through a new KeyedPartitioning.project, which computes isGrouped and isCollapsed together from the projected keys. The flag is inherited, or set when the projection leaves fewer distinct keys than the input had. A projection that keeps every position returns the partitioning unchanged, with no pass over the keys.
• UnionExec's keyed merge moves to KeyedPartitioning.concat, which joins apply, project and toGrouped as the fourth rule for how the flags travel. The rule itself is unchanged: it ORs the children's flags, so a child that really did collapse its keys marks the union. What changes for a union is upstream of it, since the children's flags are now precise.
• GroupPartitionsExec, KeyedPartitioning.toGrouped, KeyedPartitioning.createShuffleSpec and KeyedShuffleSpec.createPartitioning all propagate the flag. Each of them used to build a partitioning through the 3-argument constructor, which defaulted the flag to false. That default is gone from the parameter, so every producer now has to state it.
• createShuffleSpec projects onto the operation keys, so it computes its own collapse with project. GroupPartitionsExec decides both answers from the key groups it keeps, by asking whether a group spans more than one of its own child's partition keys. It does not compare against the aligned key list, because that list is the one both join sides agreed on. Such a list can be missing keys this side had (partition filtering) or repeat them (padding), and neither is a collapse. GroupPartitionsExec computes the flag once, from one member of the child's partitioning.

PartitioningCollection normalizes the flag across its members by OR, alongside the partitionKeys interning it already did, and its invariant check enforces the result. 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 collapse would let the gate through. Uniformity also means a consumer can read one representative instead of every member. The coverage there is structural. Once the flag is uniform, the mixed collection that could slip through cannot be built, so the test asserts the invariant rather than a plan shape.

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 taken. Its table has unique ids, so dropping the second key column keeps every key distinct, and that is exactly the scenario this change reclassifies as no collapse. The two cannot both land as written.

The class doc is reorganized around the new meaning. It leads with what a partition key is and what EnsureRequirements uses the keys for, and it gains a "Key Collapse" section that tells a collapse apart from a grouping. That section also carries why a collapsed partitioning is still reported rather than dropped to UnknownPartitioning, which was nowhere in the code.

Sidenote: both the method split and project point the same way as item 5 of #58262 (comment), which is that the SPJ decisions in EnsureRequirements are worth untangling on their own. project settles one half of that here. Applying a projection to a partitioning now happens in one place instead of once per producer. The other half is left to the follow-up. Which positions to keep is still asked in two places, KeyedShuffleSpec.keyPositions (where createShuffleSpec takes its joinKeyPositions from) and the overlap test in keysSatisfy. A comment there now records that the two are the same test, and that the answer only becomes true once the projection is applied.

Why are the changes needed?

Provenance leaves the gate open in three ways, which is why this is filed as a bug:

• The flag is dropped wherever a partitioning is rebuilt. That happens in toGrouped, in createShuffleSpec's projected partitioning, in KeyedShuffleSpec.createPartitioning and in GroupPartitionsExec, so the gate cannot see the risk even when it exists. This is the defect [SPARK-59026][SQL] Propagate isNarrowed in KeyedPartitioning.toGrouped and KeyedShuffleSpec.createPartitioning #58316 (SPARK-59026) reports, and it is fixed here.
• A reduction never sets the flag at all. Under allowCompatibleTransforms, a bucket(4, id) side reduced onto a bucket(2, id) join maps four distinct keys onto two, which is a real collapse. It drops no key position, so provenance misses the whole class of them and the opt-in is bypassed.
• A PartitioningCollection whose members disagreed could be entered through the plain one.EnsureRequirements looks for any single member that a GroupPartitionsExec would make satisfy the distribution, so a collection holding one collapsed and one plain member reached one regardless. Normalizing the flag across members closes that.

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

• 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. That is what GroupPartitionsExec does for any partitioning that never went through a projection, and it 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) has two splits for the same (1, 'x') value, is projected down to id, and is 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 is needed, because the flag and its gate arrived in 4.3.0 (SPARK-46367), which is unreleased, so no released version behaves the old way.

Propagating the flag through createShuffleSpec also makes the spec's projected partitioning compare equal to the child's in one more case. That case is the opt-in on, every position selected, and the source already collapsed and grouped. EnsureRequirements' "child partitionings not modified" fast path then fires where it previously did not. The partitions are the same, since that path is only reached when both sides' keys already match, but the existing GroupPartitionsExec keeps expectedPartitionKeys = None, so explain loses its ExpectedPartitionKeys line.

Reducers are a second case. With allowCompatibleTransforms, a bucket(4, id) side reduced onto a bucket(2, id) join really does collapse its keys, 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 for the collapse test. The partitioning is projected ten times over, each projection dropping one key position, which is what a plan with stacked Projects asks for. Every projection drops a position, so the identity fast path never fires, and no projection collapses a key, so the walk never stops early. Over a 50k-split, 25k-distinct-key partitioning with 12-position keys, 20 evaluations of that chain took 345 ms for the provenance formula and 1291 ms for this one, i.e. about 5 ms per projection that drops a position. The obvious two-pass form of the same test, one distinct over the projected keys and another over the source keys, took 2137 ms in the same harness, which is why project walks the two key lists in lockstep instead. A projection that keeps every position pays nothing, since it cannot collapse anything. The trade is deliberate: a few milliseconds of planning buys the removal of a runtime shuffle in the cases above, and it also closes the three holes that let the gate through.

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 mayGroupToSatisfy 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 collapsed member anywhere in a PartitioningCollection marks every member of it, including through a nested collection, so the outcome does not depend on which join side the collapse came from.
• New operator-level test in GroupPartitionsExecSuite that the output flag reports what the node's grouping merges. It covers the replicating and the distributing mode, and the case where the join's agreed keys drop the merged one.
• New end-to-end test for the shuffle-template chain from [SPARK-59026][SQL] Propagate isNarrowed in KeyedPartitioning.toGrouped and KeyedShuffleSpec.createPartitioning #58316, with the expectation these semantics call for. Its items has unique ids, so dropping name collapses nothing, and the final aggregate may group the union's overlapping keys with the opt-in off. That keeps the chain covered while pinning the reclassification.
• New end-to-end test that reducing keys onto a coarser transform collapses them. With allowCompatibleTransforms, an identity(item_id) side reduced onto bucket(4, id) maps several ids onto one bucket, 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 a collapse. 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 four of them: the multi-split unit test, its end-to-end counterpart, the flipped SPARK-46367 scenario and the shuffle-template chain. Deciding the flag from a key count taken after the alignment, instead of from the key groups kept, fails the partition-filter test. Dropping the collection normalization fails its unit test. Dropping the GroupPartitionsExec term fails the reducer test.
• DistributionSuite covers the propagation through toGrouped and KeyedShuffleSpec.createPartitioning. That test comes from [SPARK-59026][SQL] Propagate isNarrowed in KeyedPartitioning.toGrouped and KeyedShuffleSpec.createPartitioning #58316.
• Also ran DistributionSuite, ShuffleSpecSuite, KeyGroupedPartitioningSuite, ProjectedOrderingAndPartitioningSuite, GroupPartitionsExecSuite, EnsureRequirementsSuite, PlannerSuite, WriteDistributionAndOrderingSuite, 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 -- 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:

valisCollapsed= 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:

valnonGroupedSatisfiesWhenGrouped= 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

[peter-toth]

Good point, and dongjoon made the same one about the section title -- "coarsened" was a second word for what isCollapsed already names. Unified on collapse, and the word is gone from all the touched files.

The section is now == Key Collapse == and opens with your two entries, with the examples:

• Key collapse: a projection or a reduction maps two keys that were distinct onto the same new key. [(1, 'a'), (1, 'b'), (2, 'c')] projected onto the first position gives [1, 1, 2]: three distinct keys became two. isCollapsed records this.
• Grouping: GroupPartitionsExec physically combines the partitions that share a key. [1, 1, 2] becomes [1, 2]. isGrouped says the keys are unique, however they got that way -- a source with natively unique keys reports it too.

Then one paragraph on why the pair matters: grouping after a collapse is what produces a partition larger than any the source declared, which is what needs the config; grouping without one only merges partitions that already shared a key, and needs no opt-in.

Writing it out also caught a real bug in the config doc, which was still phrased the old way: it claimed the config gates "a partitioning that was narrowed to a subset of its keys and whose keys are no longer distinct". Both halves are false now, and two of this PR's own tests prove it -- a subset projection that keeps every distinct key is not gated, and a reduction that drops no key position at all is. Fixed in the same wording as the class doc.

[dongjoon-hyun]

Please rebase this once more to resolve the conflicts, @peter-toth .

[peter-toth]

Rebased on master and squashed to a single commit, so #58338 is out of the diff now that it has landed. The description is updated to match.

Three other things went in with the rebase, each from a review pass over the change:

• The config doc for allowKeysSubsetOfPartitionKeys was still phrased in the old vocabulary and had become wrong -- it claimed the config gates "a partitioning that was narrowed to a subset of its keys and whose keys are no longer distinct", while under these semantics a subset projection that keeps every distinct key is not gated and a reduction that drops no key position is. Two of this PR's tests pin both directions.
• GroupPartitionsExec read the flag from every member of the child's partitioning, which forced distinctKeyCount -- a pass over the keys -- once per member for an answer they share. It reads one representative now.
• I dropped the unit test that asserted the flag is taken from every member of a PartitioningCollection: once the collection normalizes the flag, head and exists agree by invariant, so that test could not fail. The normalization test covers the invariant instead, and the description's ablation list now names only the ablations that actually bite.

[peter-toth]

Two corrections after a closer look at how this relates to #58316.

Reclassified SPARK-59057 as a Bug (it was an Improvement), affected versions 4.3.0 / 4.4.0 / 5.0.0. Provenance leaves the gate open in three ways, and the description now leads with them rather than with the false refusals:

• The flag is dropped wherever a partitioning is rebuilt, so the gate cannot see a real risk. That is SPARK-59026's defect, fixed here.
• A reduction never sets the flag at all: under allowCompatibleTransforms, a bucket(4, id) side reduced onto a bucket(2, id) join maps four distinct keys onto two, but drops no key position, so provenance misses the entire class and the opt-in is bypassed.
• A PartitioningCollection whose members disagreed could be entered through the plain one, since EnsureRequirements uses find.

Took the shuffle-template chain from #58316's end-to-end test, with the expectation these semantics call for. I had left it out because its items has unique ids, so dropping name collapses nothing and all three of its assertions flip -- but that left the chain uncovered. It is now a positive test for the reclassification: the union's duplicate keys come from the two children holding the same ids rather than from the projection, so the final aggregate groups them with the opt-in off. Verified it fails on the old formula.

@dongjoon-hyun your DistributionSuite test was already in, adapted to the new flag name.

[dongjoon-hyun]

Findings from a review pass, most severe first: one behavioral issue (perf-only, never wrong results) and a few efficiency/maintainability notes. Details inline.

[dongjoon-hyun]

projectedDistinctKeyCount is taken from keyToPartitionIndices.sizebeforealignToExpectedKeys prunes to the join-agreed key set, so a collapse confined entirely to keys that the intersection then filters out still marks the output collapsed, even though every surviving partition maps 1:1 to a source partition.

Reachable on the reducer path with allowCompatibleTransforms + partitionFilter on and the subset opt-in off: ids [0, 4, 5] reduced onto buckets [0, 0, 1] with bucket 0 pruned by the other side gives 2 < 3, flagging an output whose only surviving partition was never merged. Since the flag is sticky, a later union + GROUP BY then hits the isCollapsed && !isGrouped gate and pays a shuffle the pre-PR code planned shuffle-free.

The new "filtering partition keys out is not a key collapse" test covers pruning without a collapse, but not this variant where the pruned keys are themselves the collapsed ones. Perf-only, never wrong results, but it contradicts the flag's documented meaning ("one partition here can stand for several of the original ones").

[peter-toth]

Good catch, and it took two goes to get right.

The flag no longer compares key counts at all. groupedPartitionsTuple now asks the question directly, of the key groups it keeps: does any of them cover more than one of the child's own partition keys? Counting the child's keys rather than its partitions is what tells a collapse from a source reporting several splits per key, and asking it of the kept groups is what tells it from filtering. Your shape is a test now: identity(id) with ids 0, 4, 5 reduced onto buckets [0, 0, 1], bucket 0 dropped because the other side has no rows for it, so the only surviving key covers id 5 alone.

The second go: my first version read the answer off the partitions alignToExpectedKeys emits, which is blind under distributePartitions -- that branch spreads a group's splits over one partition each, so no partition ever holds two, and the whole mode reported nothing collapsed. That was a regression against the state you reviewed. Asking the key groups is right for both modes, and there is a second test for that direction, so the two tests are each other's control.

[dongjoon-hyun]

outputPartitioning is a def, so this block re-runs on every call: flatten allocates a Seq just to collectFirst one KP, and k.distinctKeyCount is forced on a fresh child KP instance each time (DataSourceV2ScanExecBase.outputPartitioning rebuilds its KeyedPartitioning per call), so the lazy val never amortizes -- an O(#splits) distinct pass per consultation during EnsureRequirements/validation/AQE.

Computing the flag once inside the memoized groupedPartitionsTuple -- which already collectFirsts the same child KP and holds keyToPartitionIndices.size -- is semantically identical, and would also remove the unreachable .getOrElse(false) (groupedPartitions, forced just above via partitionKeys, throws when no KP exists) and the second, divergent first-KP lookup idiom in this file (line 148 uses plain collectFirst).

[peter-toth]

Done. The flag is computed inside groupedPartitionsTuple, so outputPartitioning no longer forces distinctKeyCount on a child partitioning that DataSourceV2ScanExecBase rebuilt on every call. The child lookup is now a shared childKeyedPartitioning used by both the grouping and the flag, which removes the second idiom and the unreachable getOrElse(false) as well.

[dongjoon-hyun]

Since isCollapsed is a constructor param, it participates in case-class equality and survives canonicalization. Two otherwise-identical exchanges -- one templated from a collapsed partitioning, one from an equivalent non-collapsed one (two independent joins over a shared subplan; the collection OR-normalization doesn't reach across them) -- no longer compare equal under sameResult, so ReuseExchangeAndSubquery shuffles the shared subplan twice where it previously reused the exchange.

Narrow configuration, and arguably defensible since the metadata genuinely differs (reusing a false-flag exchange where a true one was planned would re-launder the protection under AQE re-planning) -- but it deserves a conscious decision.

[peter-toth]

Conscious decision: keep it. The metadata genuinely differs, and reusing an exchange planned from a non-collapsed partitioning where a collapsed one was expected would re-launder the protection under AQE re-planning, which is the failure this change exists to close. The cost needs the opt-in plus two independent joins over a shared subplan, and it is a lost reuse rather than a wrong answer.

[dongjoon-hyun]

The inherited flag is also stamped on join types that expose only the shuffled side (LeftOuter/LeftSemi/LeftAnti/LeftExistence in ShuffledJoin.outputPartitioning); for semi/anti the output carries zero collapsed-side rows, yet inherits the flag forever. The created KP is grouped (canCreatePartitioning requires isGrouped), so the flag only bites after a union reintroduces duplicate keys -- exactly the case the class doc blesses as groupable without opt-in -- and the flag's documented contract is false for these partitionings.

The spec can't know the consuming join type, so if this is worth tightening, the place would be ShuffledJoin.outputPartitioning (clear the flag for LeftExistence). Fine to keep as the documented safe-direction conservatism, but worth noting the cost and the labeling mismatch.

[peter-toth]

I would keep it, and I do not think the labelling is wrong. The flag describes key granularity, not which rows survive: a semi-join output laid out on the template's collapsed key space really does have partitions covering several of the source's keys, whichever side's rows flow through. Clearing it for LeftExistence would claim the layout is finer than it is, and the first union that reintroduces duplicate keys would then group genuinely coarse partitions without the opt-in.

[dongjoon-hyun]

Keeping isCollapsed: Boolean = false as a defaulted parameter leaves the laundering vector this PR closes open for the next producer: a future operator that derives a KeyedPartitioning and calls the 3-arg constructor silently gets false and bypasses the gate with no compiler or test signal -- the same failure mode as SPARK-59026, and the wrong code looks identical to correct pre-existing code.

Only the companion apply (a provably-fresh source partitioning) actually relies on the default; dropping the default and passing isCollapsed = false explicitly there would make the compiler force every future producer to decide the flag.

[peter-toth]

Done, the default is gone. In main code only the companion apply relied on it, and it now passes isCollapsed = false explicitly with a comment saying why a fresh source partitioning is the layout everything else is compared against.

It also caught a live instance of exactly the laundering you describe, in test code: DistributionAndOrderingSuiteBase.resolvePartitioning destructured the flag away with _ and rebuilt the partitioning without it. That is the concrete payoff for the compile-time change, so it is in the commit message.

[dongjoon-hyun]

This inherited || projected distinct < source distinct predicate is now hand-rolled at three producer sites: here, KeyedPartitioning.createShuffleSpec, and GroupPartitionsExec.outputPartitioning, each with site-specific inputs and caveats. Given the PR's own observation that one producer laundering the flag is how the protection went missing, a shared helper on KeyedPartitioning, e.g.

defcollapsedAfterProjection(projectedDistinctKeyCount: Int):Boolean=
isCollapsed || projectedDistinctKeyCount < distinctKeyCount

would keep the correctness rule in one place instead of three files that must stay in sync.

[peter-toth]

Added, as KeyedPartitioning.collapsesOnProjection. One correction to the shape you sketched: the isCollapsed || disjunct would be dead at both call sites, because each already carries the inherited flag outside it -- and AliasAwareOutputExpression has to, since it reads the flag from every input rather than only from the one whose keys it counts. So the helper is the count comparison alone, and its scaladoc says that the callers own the inherited term.

GroupPartitionsExec deliberately does not use it: it can answer the question exactly from the key groups it keeps, and the scaladoc points at that as the reference definition.

[peter-toth]

Update after a rewrite, because the helper I named here no longer exists.

collapsesOnProjection is gone, and the rule it held moved inside KeyedPartitioning.project, which both producers now call for the whole projection. So the rule is in one place rather than the three you started from, and there is no helper left that a future producer could forget to call.

GroupPartitionsExec still answers the question its own way, from the key groups it keeps, for the reason in my earlier reply.

[dongjoon-hyun]

Two O(n) distinct passes over the same list: projectedKeys.distinct.length here, then .toGrouped on the next statement re-runs partitionKeys.distinct.sorted over the same projectedKeys; for an ungrouped, uncollapsed source, forcing distinctKeyCount adds a third pass over the source keys. The block also runs when joinKeyPositions selects every position, where the comparison is tautologically false.

Materializing val d = projectedKeys.distinct once and building the grouped KP directly (same ordering source as toGrouped), and/or guarding on positions actually dropped, is semantically identical and cheaper. Planning-time only, so minor -- but it runs per createShuffleSpec call from EnsureRequirements/ValidateRequirements.

[peter-toth]

Done. The count now comes off the grouped partitioning that toGrouped already builds, so there is one distinct where there were two, and the comparison is skipped entirely when joinKeyPositions selects every position -- which, as you note, cannot collapse anything.

[peter-toth]

Update, and the intermediate state was worse than what I told you here. Extracting the projection into KeyedPartitioning.project brought the second distinct back for a while, because project counted the projected keys and toGrouped then deduped the same list again.

It is one pass now, and not a distinct at all. project walks the projected keys alongside the keys they came from and fills a projectedKey -> sourceKey map. A second, different source key on an existing projected key is the collapse, and it also means the projected keys are not unique, so the walk stops there. The source's own distinct count is not computed any more, and the source keys are never hashed, only compared where a projected key repeats.

Measured on the worst case for that test, a 50k-split partitioning with 25k distinct keys and 12-position keys, projected ten times over, 20 evaluations: 345 ms for the old provenance formula, 1291 ms for this one, 2137 ms for the two-distinct form. toGrouped after it also skips its own dedup when the keys are already unique.

[dongjoon-hyun]

LGTM, thanks for the thorough follow-up. The group-based collapse detection resolves the pre-alignment counting concern cleanly in both directions -- the two new tests (the pruned-collapse shape and the distributePartitions mode) pin exactly the states that matter -- and removing the constructor default immediately catching a live laundering instance in DistributionAndOrderingSuiteBase is a nice payoff for the compile-time change. The keep decisions on exchange reuse and the shuffled-side flag inheritance both come with sound rationale.

Two non-blocking nits inline, plus one note that has no diff line to attach to: the reworked collapsesKeys branches (kept-groups path, getOrElse(Seq.empty) miss path, childIsCollapsed passthrough) are covered only end-to-end through SPJ planning. An operator-level test in GroupPartitionsExecSuite pinning the tuple's third element against a hand-built child partitioning would keep them covered if planner routing ever changes. Fine as a follow-up or not at all.

[dongjoon-hyun]

Nit, non-blocking: this re-derives the kept groups with its own keyToPartitionIndices.getOrElse lookups, duplicating the key matching alignToExpectedKeys just performed -- if that method's matching ever changes (normalization, emitting instead of dropping unexpected keys), the emitted partitions and the flag would silently be computed from different group selections. Computing the collapse bit inside alignToExpectedKeys in the same pass would keep one source of truth; short of that, fusing the map into expectedPartitionKeys.get.exists { case (key, _) => coversSeveralChildKeys(keyToPartitionIndices.getOrElse(key, Seq.empty)) } at least drops the intermediate Seq and short-circuits.

[peter-toth]

This is your second form now, the exists over expectedPartitionKeys, so there is no intermediate Seq and it short-circuits.

I went your first way to begin with and computed the collapse bit inside alignToExpectedKeys, then moved it back out. Inside, the aligner had to grow a mutable buffer and a third tuple element that it never used itself, and its doc had to explain a return value that belonged to the caller.

What the current form gives up is the single lookup site you were after. keyToPartitionIndices.get(key) now appears twice in the file, and if the aligner's key matching changed, the two could drift. One thing did get better with the move: both branches of the scan take their groups from keyToPartitionIndices, so a padded expected key cannot produce an empty group at all, and the getOrElse(key, Seq.empty) plus its guard are gone.

[dongjoon-hyun]

Nit, non-blocking: this copies even when projectedCollapsed == grouped.isCollapsed (both false in the common no-collapse case); if (projectedCollapsed) grouped.copy(isCollapsed = true) else grouped avoids the redundant allocation per spec creation.

[peter-toth]

Moot now. createShuffleSpec does not copy at all any more, it is project(joinKeyPositions).toGrouped, and project decides the flag in the constructor call.

[szehon-ho]

Could we state the two cases explicitly here?

• Collapsed and ungrouped: duplicate keys remain, so ClusteredDistribution is refused when the config is off.
• Collapsed and grouped: keys are already unique, so it is accepted regardless of the config.

The current paragraph mixes these states, which makes “a collapsed partitioning” ambiguous.

[peter-toth]

Done, in your words:

• Collapsed and ungrouped: duplicate keys remain, so grouping them would merge partitions the source held apart. mayGroupToSatisfy() refuses a ClusteredDistribution unless the config is on.
• Collapsed and grouped: the keys are already unique, so grouping merges nothing and satisfies() accepts a ClusteredDistribution whatever the config says. A partitioning reaches this state by being grouped with the config on, or by having its keys reduced onto a coarser transform.

The paragraph that mixed the two states is gone. What follows the bullets is only what is common to both, which is why a collapsed partitioning is still reported rather than dropped to UnknownPartitioning.

[peter-toth]

Rebased and pushed. SPARK-58974 merged, so this is no longer stacked and the PR is one commit against master. The description is regenerated from it. The flag's contract and every producer are as reviewed, and two things around them changed shape.

groupedSatisfies split in two. Its !isGrouped term was never a property of the partitioning, it was a caller discriminator: satisfies0 only ever called the method with isGrouped = true, and EnsureRequirements only ever with false. So there are now keysSatisfy (the key matching, private) and mayGroupToSatisfy (keysSatisfy plus the permission to coalesce), and EnsureRequirements asks the second one of its non-grouped list. The permission is checked first, as it was before the split, so the gated case pays nothing for the key matching.

One projection.KeyedPartitioning.project applies a projection and decides both isGrouped and isCollapsed, and PartitioningPreservingUnaryExecNode and createShuffleSpec both go through it. UnionExec's merge rule moved to KeyedPartitioning.concat for the same reason, so the four rules for how the flag travels are now all on the type: from scratch, projection, grouping, concatenation.

Sidenote: both point the same way as item 5 of the follow-up list on #58262. The remaining half is that which positions are operation keys is still derived in more than one place, and a comment in keysSatisfy records it.

Two decisions from your review are unchanged and stated in the threads: ReuseExchange keeps the flag in equality, and LeftExistence keeps inheriting it. The commit message no longer carries the Co-authored-by line, to match what you did to the description.

[peter-toth]

@ulysses-you the report I promised on #58338, on how the tightened notion reconciles with allClusterKeysCovered.

Your shape is the one this PR fixes. A source reporting several splits per partition key, projected onto exactly the operator's clustering, no longer trips the gate, because the flag means an actual collapse instead of "positions were dropped". Measured on it with the opt-in off: before, no GroupPartitionsExec and 2 shuffles; after, 1 GroupPartitionsExec and 0 shuffles. It is a test in this PR, several splits per partition key are grouped without allowKeysSubsetOfPartitionKeys, plus a unit test on the partitioning itself.

So the mismatch you named is gone. Duplicated keys on their own no longer refuse anything, which is the tolerance allClusterKeysCovered's comment states. The gate now refuses only when two keys that were distinct in the source ended up on one key, and it is insensitive to key order and to repeated cluster keys for the same reason that gate is: it compares distinct key values, which no permutation or repetition of the key expressions changes.

On the two gates sharing one notion, I ended up thinking they should not. They ask different questions and neither subsumes the other. allClusterKeysCovered is a set-coverage test on expressions, "is every cluster key among the partition keys", and it guards against joining on keys coarser than the join keys. The collapse gate is a test on the key values, and it guards a coarsening that already happened upstream in a projection or a reduction. Coverage says nothing about whether those keys collapsed, so with requireAllClusterKeysForCoPartition on the collapse gate still has work to do. What they now share is the tolerance, not the notion.

[peter-toth]

@dongjoon-hyun, @szehon-ho, I've updated the PR with splitting groupedSatisfies() and introducing KeyedPartitioning.project() and rephrasing KeyedPartitioning documentation. Can you please take another look when you get a chance?

@HeartSaVioR, I would like to include this PR into 4.3.0 if possible, as besides this is a bugfix, it changes the partition narrowing semantics added in this version (SPARK-46367 / #55519). So if we make this change now, then we don't need a migration guide addition later.

[HeartSaVioR]

I'm OK with merging this in 4.3. Whether this has to go in 4.3 RC1 is the main question, but I don't know whether this is an issue since we know RC1 tends to fail, so I don't think it is a big deal.

[szehon-ho]

This claim is not true for reductions: a reducer can map distinct keys onto a coarser transform without the partitioning ever going through a projection, which is exactly the case covered by the new reducing keys onto a coarser transform collapses keys test. Could we remove this sentence or qualify it so the class doc does not contradict the collapse definition above?

[peter-toth]

Thanks, you are right, and the doc contradicted itself: the collapse definition thirteen lines above names both a projection and a reduction, and this bullet only excluded the projection. Fixed in 6d2aa2b by naming both producers, so the reducing keys onto a coarser transform collapses keys case is no longer contradicted.

[szehon-ho]

thanks this looks much clearer. left one comment about comment correctness, so its quite minor, feel free to apply when convenient.

[peter-toth]

Merge Summary:

• merged into master 8fd856c
• merged into branch-4.x 1ff0ea7
• merged into branch-4.3 7d44b56

Posted by merge_spark_pr.py

[peter-toth]

Thank you @dongjoon-hyun and @szehon-ho for the review.