peter-toth opened a new pull request, #55901: URL: https://github.com/apache/spark/pull/55901
### What changes were proposed in this pull request? - Add a `require` in `PartitioningCollection` that all `KeyedPartitioning`s reachable from the collection share the same `partitionKeys` reference (`eq`) and have matching expression arity. The check walks the partitioning tree via `foreach` so nested collections are covered. - Add a smart factory `PartitioningCollection.fromPartitionings` that interns `partitionKeys` references across `KeyedPartitioning`s. Use this at sites that combine independently-computed partitionings (joins) where keys are structurally equal but not reference-equal. The factory uses manual recursion rather than `transformWithPruning` because `KeyedPartitioning.equals` compares `partitionKeys` element-wise, which would make `transformWithPruning` discard the rule's replacement as structurally-equal-to-input. - In `GroupPartitionsExec.outputPartitioning`, hoist `val partitionKeys = groupedPartitions.map(_._1)` above the `transform` so every rebuilt `KeyedPartitioning` shares the same `partitionKeys` reference. Drop the ad-hoc consistency assert (now enforced by `PartitioningCollection`). - Switch `ShuffledJoin` and `StreamingSymmetricHashJoinExec` to `PartitioningCollection.fromPartitionings` for their inner-join `outputPartitioning`. - Update affected tests to construct collections via `fromPartitionings`. Rewrite the `SPARK-46367` arity-mismatch test in `ProjectedOrderingAndPartitioningSuite` since the scenario is now rejected at `PartitioningCollection` construction rather than inside `AliasAwareOutputExpression`. ### Why are the changes needed? The "all `KeyedPartitioning`s in a collection must agree on `partitionKeys`" invariant already existed informally -- `GroupPartitionsExec.outputPartitioning` had a runtime assert checking `==`, `AliasAwareOutputExpression.projectKeyedPartitionings` asserted matching arity, and various consumers relied on the invariant being upheld. Consolidating the check into the `PartitioningCollection` constructor makes it load-bearing: any future construction site that violates it fails immediately rather than producing silently divergent state. The stronger `eq` check enables cheap reference comparisons downstream and is naturally achievable when collections are built through the smart factory. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing test suites (`EnsureRequirementsSuite`, `GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite`) updated to use `PartitioningCollection.fromPartitionings` where they previously constructed collections from independently-built `KeyedPartitioning`s. The `SPARK-46367` test was rewritten to assert that the invalid mixed-arity scenario is rejected at `PartitioningCollection` construction. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code 4.7 -- 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]
