naveenp2708 opened a new pull request, #59317: URL: https://github.com/apache/spark/pull/59317
### What changes were proposed in this pull request? `GroupPartitionsExec.alignToExpectedKeys` spreads a key's splits across its expected partitions with `padTo`. `padTo` only pads, it never truncates, so if a key holds more splits than its expected count, the distribute branch emits more partitions than the count says, while `isGrouped` can still come back `true` (it only flips to `false` when the count is above 1). The layout then disagrees with what it reports. This adds a guard in the distribute branch, `assert(splits.size <= numSplits)`, next to the existing `assert(numSplits > 0)`. ### Why are the changes needed? Downstream code trusts `isGrouped`: `satisfies` for `ClusteredDistribution`, join pairing, and shuffle-spec construction all read it. So a wrong count doesn't fail where it's made, it fails somewhere else. In the SPARK-58996 review (finding 5, raised by @peter-toth) it showed up as a failed `PartitioningCollection` invariant reached through `SortMergeJoinExec.outputPartitioning`, even though the bad count came from `alignToExpectedKeys`. SPARK-58996 already removed the known sources of a wrong count, so nothing hits this today. The guard just makes sure that if some future change produces one, it fails right here at `GroupPartitionsExec` instead of further down. No behavior change on valid input. This is the fail-fast option from SPARK-59234. I started with it since it has no behavior change, but I'm glad to do the other option instead and derive `isGrouped` from the partitions actually emitted, if you'd rather. @peter-toth @ulysses-you, which do you prefer? ### Does this PR introduce any user-facing change? No. ### How was this patch tested? Added a `GroupPartitionsExecSuite` case: a child holds two splits for a key, the expected count for it is 1 under `distributePartitions`, and building the node fails fast. Checked that it fails without the guard. ### Was this patch authored or co-authored using generative AI tooling? No -- 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]
