peter-toth opened a new pull request, #58527: URL: https://github.com/apache/spark/pull/58527
### What changes were proposed in this pull request? `EnsureRequirements` stops asking a `ShuffleSpecCollection` for a single answer. It resolves the one member the matched children agreed on, preferring the finest when several qualify, and uses that member to build a re-shuffled child's partitioning. - `flattenSpec` replaces the head read. It recurses, because `ShuffledJoin.outputPartitioning` builds `PartitioningCollection.fromPartitionings(Seq(left, right))` for an inner join, so a chain of same-key joins nests collections. - The chosen member has to be compatible with every matched child. When no member is, there is no shared layout, so every child takes the ordinary shuffle. Reaching that needs three or more clustered children, and no operator has three today. - The `joinKeyPositions` pushed into a compatible child now come from that child's own matching member, because they index into that child's partition expressions. - `ShuffleSpecCollection.createPartitioning` is untouched. Its `require` stays as a guard, and a new unit test pins it. ### Why are the changes needed? Under `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled`, `KeyedPartitioning.createShuffleSpec` projects each member of a `PartitioningCollection` onto *its own* join-key subset and drops the duplicate keys that projection creates. The members of the resulting `ShuffleSpecCollection` can therefore end up with different `numPartitions`. `EnsureRequirements` then asks the collection for a shuffle template, and `ShuffleSpecCollection.createPartitioning` requires all members to agree: ``` java.lang.IllegalArgumentException: requirement failed: expected all specs in the collection to have the same number of partitions ``` so planning fails outright. With `items` partitioned by `[identity(id), identity(arrive_time)]`, one row per split, `purchases` unpartitioned, and `v2BucketingShuffleEnabled=true`, `partiallyClusteredDistribution=false`, `allowKeysSubsetOfPartitionKeys=true`: ```sql SELECT /*+ MERGE(i, p) */ id, t1, t2, i.price AS purchase_price, p.price AS sale_price FROM (SELECT id, arrive_time AS t1, arrive_time AS t2, price FROM testcat.ns.items) i JOIN testcat.ns.purchases p ON i.id = p.item_id AND i.t1 = p.time ``` Selecting `arrive_time` twice under two aliases makes the alias cross-product produce members that cover different numbers of join keys, which is where the counts diverge. The collection cannot answer that question locally. `isCompatibleWith` succeeds when *any* member matches, so the collection alone never said which member the two sides agreed on, and `createPartitioning` fell back to `specs.head`, whichever the alias cross-product enumerated first. Narrowing the collection to its finest members would satisfy the `require`, but it would still be a guess: the right member is the one the *other* side matched, and that is only visible in `EnsureRequirements`. ### Does this PR introduce _any_ user-facing change? Yes. The query above failed to plan and now runs, producing one shuffle and the right rows. The other two changes are latent, so they are not user-facing. No query reaches a wrong `joinKeyPositions` pushdown: a join goes through `checkKeyGroupCompatible`, which already pushes each side's own positions, and a cogroup's grouping key is synthesized so neither side stays keyed. No operator has three or more clustered children. ### How was this patch tested? Four new tests. Each was measured against the same commit with only the `EnsureRequirements` change reverted. | test | on base | |---|---| | `KeyGroupedPartitioningSuite`: both sides of the join land on the same collection member | fails with the `require` above | | `EnsureRequirementsSuite`: the re-shuffled side lands on the member the keyed side was matched on | fails with the `require` above | | `EnsureRequirementsSuite`: pushed-down positions index into the child's own partition expressions | fails | | `ShuffleSpecSuite`: a collection whose members cover different key subsets disagrees | passes, by design | The last one pins the guard rather than the fix. It asserts that the members disagree on purpose, that every one of them stays available for `isCompatibleWith`, and that asking the collection for a single partitioning throws. That turns the `require` from untested prose into a pinned contract, which matters now that the method has no production caller. Green: `ShuffleSpecSuite`, `EnsureRequirementsSuite`, `KeyGroupedPartitioningSuite`, 202 tests in all. `dev/lint-scala` is clean. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code -- 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]
