peter-toth opened a new pull request, #58351: URL: https://github.com/apache/spark/pull/58351
### 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. `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 https://github.com/apache/spark/pull/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? Co-authored-by: Dongjoon Hyun <[email protected]> Generated-by: Claude Code (Opus 5) This PR is based on #58338 (SPARK-58974), which is still open, so the first two commits here are that PR. Only the last commit belongs to this one. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
