peter-toth opened a new pull request, #58338: URL: https://github.com/apache/spark/pull/58338
### What changes were proposed in this pull request? `KeyedPartitioning.groupedSatisfies` refuses a narrowed, non-grouped partitioning unless `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` is set, because grouping it would merge partitions that held distinct keys in the original finer-grained partitioning. That guard lived inside the `requireAllClusterKeys = false` arm, so it never ran when `spark.sql.requireAllClusterKeysForDistribution` was enabled. This moves the guard above the `requireAllClusterKeys` check. The skew risk does not depend on which key sets count as matching, so the decision should not either. The config doc for `allowKeysSubsetOfPartitionKeys.enabled` is updated as well: it did not mention this second thing the config gates, which is not tied to `requireAllClusterKeysForDistribution` being false. ### Why are the changes needed? With `spark.sql.requireAllClusterKeysForDistribution = true` a narrowed partitioning was accepted, `EnsureRequirements` inserted a `GroupPartitionsExec`, and the skew exposure was taken with the opt-in config still off. Note this is the *stricter* of the two settings, which makes it the more surprising direction. Reproduced with a table partitioned by `(id, dept)` projected down to `id` -- keys collapsing to `[1, 1, 2]` -- and joined on `id`: with `requireAllClusterKeys = true` the plan gets a `GroupPartitionsExec` and no shuffle while `allowKeysSubsetOfPartitionKeys` is off, whereas with `requireAllClusterKeys = false` the same query correctly shuffles both sides. Please note that the guard's condition is imprecise, and the hoist makes that imprecision reachable for one more config value. `isNarrowed && !isGrouped` is a proxy for "the narrowing collapsed distinct keys", but `!isGrouped` has causes that have nothing to do with narrowing: a source that reports several splits per partition key, or a union whose children have distinct keys individually and repeat keys across children. Grouping such a partitioning merges only same-key partitions, which is what `GroupPartitionsExec` does for any non-narrowed partitioning and needs no opt-in, yet the condition refuses it. Until now that false refusal could only happen with `requireAllClusterKeysForDistribution = false`; after this change it can happen with either value. Tightening the condition to actual key collapse is a separate change, which I am working on as a follow-up; this PR keeps the condition as it is and only fixes where it is evaluated. ### Does this PR introduce _any_ user-facing change? Yes, a plan-level change. With `requireAllClusterKeysForDistribution` enabled, Spark previously coalesced partitions derived from a narrowed partitioning without `allowKeysSubsetOfPartitionKeys.enabled`, risking skewed partitions; it now inserts a shuffle unless that config is enabled. Query results are unchanged. No migration guide entry: the guard, and with it the bypass, arrived in 4.3.0, which is unreleased, so no released version behaves the old way. This goes to `branch-4.3` as well. ### How was this patch tested? * New unit test in `ProjectedOrderingAndPartitioningSuite` calling `groupedSatisfies` directly on a narrowed, ungrouped partitioning, for both values of `requireAllClusterKeys` and both values of the opt-in. It fails on master with `kp.groupedSatisfies(required) was true` for `requireAllClusterKeys = true`. * New end-to-end test in `KeyGroupedPartitioningSuite` asserting the guard behaves identically for both values of `requireAllClusterKeys`: no `GroupPartitionsExec`, and a shuffle instead. It fails on master for `requireAllClusterKeys = true`. It also varies `v2BucketingShuffleEnabled`, because that decides whether the refused side is laid out on the other side's declared partition keys (one shuffle) or both sides are shuffled (two). * That test also covers `allowKeysSubsetOfPartitionKeys = true` for both values of `requireAllClusterKeys`, since this is the first time the opt-in affects `groupedSatisfies` when `requireAllClusterKeysForDistribution` is enabled -- it must restore the coalescing and avoid the shuffles. Neither suite had any coverage of `requireAllClusterKeysForDistribution`, which is how the bug survived. * Also ran `DistributionSuite`, `KeyGroupedPartitioningSuite`, `EnsureRequirementsSuite`, `PlannerSuite`, `ProjectedOrderingAndPartitioningSuite`, `DataFrameSetOperationsSuite`, `AdaptiveQueryExecSuite`, `CoalesceShufflePartitionsSuite`, and the TPC-DS plan stability suites -- no golden file changed. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) -- 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]
