[SPARK-59026][SQL] Propagate isNarrowed in KeyedPartitioning.toGrouped and KeyedShuffleSpec.createPartitioning - #58316
Open
dongjoon-hyun wants to merge 1 commit into
Open
Conversation
…d and KeyedShuffleSpec.createPartitioning
dongjoon-hyun
force-pushed
the
SPARK-59026
branch
from
August 26, 2026 16:24
e165d9f to
8c7e66d
Compare
Member
Author
|
Could you review this PR, @peter-toth ? |
Member
Author
|
CI is running here. |
Contributor
|
This overlaps with my not yet open SPARK-58974 fix. Let me think about how to proceed. |
Member
Author
|
Feel free to open your PR independently. Let's see the code difference, @peter-toth . |
Contributor
|
@dongjoon-hyun I filed SPARK-58974 for a related but separate defect in the same guard (#58338), and I want to be transparent about how the two relate. The I first had both fixes in one PR and then dropped yours from it, so #58338 is only the hoist. The reason is that I do not think this PR is the right fix. The deeper point, which I also wrote into #58338's description: the So I am following up with the contract fix: the flag will mean actual collapse - set only when a projection loses distinct keys, sticky from there on - which is decidable at the one site that projects the keys, since both key lists are in hand there. |
Member
Author
|
No problem at all from my side. Please proceed in your ways ASAP before RC1, @peter-toth . |
Contributor
|
The follow-up I promised is open: #58351 (SPARK-59057). It makes the flag mean actual key collapse -- a projection that mapped keys which were distinct in the input onto the same projected key -- instead of "positions were dropped", and renames it to Your finding is in it: all five producers that rebuild a |
peter-toth
added a commit
to peter-toth/spark
that referenced
this pull request
Aug 28, 2026
…collapse, and rename it to isCollapsed ### 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. It 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 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 collapse would let the gate through. Uniformity also means a consumer can read one representative instead of forcing the distinct key count on every member. Note 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. `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 apache#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. Propagating the flag through `createShuffleSpec` also makes the spec's projected partitioning compare equal to the child's in one more case -- opt-in on, every position selected, source already collapsed and grouped -- so `EnsureRequirements`' "child partitionings not modified" fast path fires where it previously did not. Same plan, one less rebuild. 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; 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; 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 apache#58316. * 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? Co-authored-by: Dongjoon Hyun <dongjoon@apache.org> Generated-by: Claude Code (Opus 5)
peter-toth
added a commit
to peter-toth/spark
that referenced
this pull request
Aug 28, 2026
…collapse, and rename it to isCollapsed ### 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. It 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 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 collapse would let the gate through. Uniformity also means a consumer can read one representative instead of forcing the distinct key count on every member. Note 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. `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 apache#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. Propagating the flag through `createShuffleSpec` also makes the spec's projected partitioning compare equal to the child's in one more case -- opt-in on, every position selected, source already collapsed and grouped -- so `EnsureRequirements`' "child partitionings not modified" fast path fires where it previously did not. Same plan, one less rebuild. 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; 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; 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 apache#58316. * 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? Co-authored-by: Dongjoon Hyun <dongjoon@apache.org> Generated-by: Claude Code (Opus 5)
peter-toth
added a commit
to peter-toth/spark
that referenced
this pull request
Aug 28, 2026
…collapse, and rename it to isCollapsed ### 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. It 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 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 collapse would let the gate through. Uniformity also means a consumer can read one representative instead of forcing the distinct key count on every member. Note 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. `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 apache#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 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** -- `toGrouped`, `createShuffleSpec`'s projected partitioning, `KeyedShuffleSpec.createPartitioning`, `GroupPartitionsExec` -- so the gate cannot see the risk even when it exists. This is the defect apache#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 -- a real collapse -- but 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` uses `nonGrouped.find(_.groupedSatisfies(distribution))`, which accepts when any single member does, so a collection holding one collapsed and one plain member reached a `GroupPartitionsExec` 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 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. Propagating the flag through `createShuffleSpec` also makes the spec's projected partitioning compare equal to the child's in one more case -- opt-in on, every position selected, source already collapsed and grouped -- so `EnsureRequirements`' "child partitionings not modified" fast path fires where it previously did not. Same plan, one less rebuild. 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 for the shuffle-template chain from apache#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 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; 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; 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 apache#58316. * 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? Co-authored-by: Dongjoon Hyun <dongjoon@apache.org> Generated-by: Claude Code (Opus 5)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR fixes two spots in
partitioning.scalawhere theKeyedPartitioning.isNarrowedflagwas silently dropped:
KeyedShuffleSpec.createPartitioningbuilt the shuffled side'sKeyedPartitioningwith the3-argument constructor, so
isNarrowedfell back to the defaultfalseeven when thetemplate partitioning was narrowed.
KeyedPartitioning.toGroupedlikewise dropped the flag when building the grouped variant.Both now propagate
isNarrowedfrom the source partitioning.Why are the changes needed?
isNarrowedmarks aKeyedPartitioningthat was derived from a finer-grained one by droppingkey positions.
AliasAwareOutputExpressiondocuments that this flag must be sticky: anarrowed, non-grouped partitioning must not satisfy
ClusteredDistributionvia groupingunless
spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabledis on,because
GroupPartitionsExecwould merge partitions that held distinct keys in the originalpartitioning, carrying the same skew risk.
KeyedShuffleSpec.createPartitioningbroke that stickiness. When a narrowed-but-groupedpartitioning is used as the shuffle template for the other join leg (with
spark.sql.sources.v2.bucketing.shuffle.enabled), the shuffled side mirrors the narrowedside's partition keys but declared
isNarrowed=false. If an outer join then exposes only theshuffled side and
UnionExecmerges it with an overlapping-key sibling into an ungroupedKeyedPartitioning, theisNarrowed && !isGroupedguard ingroupedSatisfiesnever fires,and a downstream operator can coalesce the narrowed keys with
GroupPartitionsExecwithoutthe opt-in config — bypassing the skew protection. The new end-to-end test reproduces this
chain and fails without the fix.
The flag was lost in the refactoring of SPARK-55535/SPARK-55092 (a1c62dd).
KeyGroupedPartitioningand Storage Partition Join #54330Does this PR introduce any user-facing change?
Yes, a narrow plan-level fix. In the multi-stage scenario above, Spark previously coalesced
partitions derived from a narrowed partitioning without
allowJoinKeysSubsetOfPartitionKeys.enabled, risking skewed partitions; it now inserts ashuffle unless the config is enabled. Query results are unchanged.
How was this patch tested?
KeyGroupedPartitioningSuitereproducing the full chain(narrowing projection -> shuffle from the narrowed template -> right outer join -> union
merge -> aggregate), verified to fail without the fix.
DistributionSuitepinning flag propagation throughtoGroupedandKeyedShuffleSpec.createPartitioning.KeyGroupedPartitioningSuite(100 tests),DistributionSuite,ProjectedOrderingAndPartitioningSuite, andEnsureRequirementsSuite.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Fable 5