peter-toth opened a new pull request, #58825: URL: https://github.com/apache/spark/pull/58825
### 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 `joinKeyPositions` half is user-facing too. I originally wrote here that it was latent, on the grounds that a cogroup's grouping key is synthesized so neither side stays keyed. @sunchao pointed out that this is only true of the Scala `CoGroupExec`, whose key comes from an `AppendColumns` that no `KeyedPartitioning` satisfies. A Pandas or Arrow cogroup groups on real columns, so two keyed children do reach the per-child branch, and this suite already had a `FlatMapCoGroupsInPandasExec` test with two of them. So: with two keyed children whose partition expressions are laid out differently, the second side is handed the first side's positions and ends up grouped on its other partition column. The plan test below reproduces it and fails without the fix, reporting `List(Some(List(1)), Some(List(1)))` where `List(Some(List(1)), Some(List(0)))` is right. I have not built an end-to-end query for it. The only part that stays latent is the three-or-more clustered children case, which no operator has. ### 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 (a Pandas cogroup over two keyed children) | 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. ### Backport notes Cherry-picked from `branch-4.2` [`c4b102465c7`](https://github.com/apache/spark/commit/c4b102465c7d77c3dfc7b755d673ea0e748e48fb), the lowest branch that has it. The `branch-4.1` backport is apache#58824. It replaces SPARK-59025's `unwrapSpecCollection`, which is on this branch as `26fb64f34f9`, the same way it did upstream. Three mechanical conflicts in `EnsureRequirements`: the `scala.annotation.tailrec` import goes with `unwrapSpecCollection`, and the names here are `KeyGroupedShuffleSpec` and `populateJoinKeyPositions`. The incoming `scala.collection.immutable.BitSet` import is deliberately not added, since the only uses on this branch are `mutable.BitSet`. All four tests are kept. Three of them read the resolved member off a `GroupPartitionsExec`, which does not exist here, so they read `BatchScanExec.spjParams.joinKeyPositions` instead, which is where the positions are recorded on this branch. The `ShuffleSpecSuite` test builds `KeyGroupedPartitioning(exprs, n, keys)` rather than `KeyedPartitioning(exprs, keys)`, and uses the suite's own `$"id".int` style so it needs no new import. The two `EnsureRequirementsSuite` tests also need `DummySparkPlanWithBatchScanChild`, a seven-line private test class that came with SPARK-53322. That ticket is on `branch-4.1` but not here, so the class is copied in with a comment saying where it came from. Dropping the two tests was the alternative, and they are the ones that cover the fix most directly. **Measured here.** With only the `EnsureRequirements` change reverted, the `KeyGroupedPartitioningSuite` test and both `EnsureRequirementsSuite` tests fail, the first two with the exact `require` this fixes. The `ShuffleSpecSuite` test passes either way, by design, since it pins the guard rather than the fix. Green: `ShuffleSpecSuite` 13, `KeyGroupedPartitioningSuite` 86, `EnsureRequirementsSuite` 21. `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]
