ammarchalifah opened a new pull request, #60058: URL: https://github.com/apache/spark/pull/60058
### What changes were proposed in this pull request? On `branch-4.1`, an outer storage-partitioned join pushes its merged partition values into the scans under an inner join. When that inner join shuffled one side (`spark.sql.sources.v2.bucketing.shuffle.enabled=true`), the exchange kept the partition count it was planned with. The partitioned scan beside it grew to the merged keys, so the two children of the inner join no longer had the same number of partitions. Stopping the push at the exchange, which [SPARK-60232](https://issues.apache.org/jira/browse/SPARK-60232) / [SPARK-60552](https://issues.apache.org/jira/browse/SPARK-60552) do, does not fix this: the exchange still has the old partition count. This PR re-plans that exchange onto the merged values, and refuses the SPJ when it cannot. All changes are in `EnsureRequirements`. * `retargetKeyGroupedShuffle`: a `ShuffleExchangeExec` whose output is a `KeyGroupedPartitioning` is rewritten onto the pushed values when it evaluates the same keys. That means the same number of fields, every pushed split count is 1, no reducer, and no partial clustering. A reducer changes the value domain, and partial clustering replicates partitions; `KeyGroupedPartitioner` can follow neither. Values the shuffle already has are left as they are. * `nestedShufflesFollowPush`: before the merged values are pushed, every shuffle the push would reach must already have those values, or be a `ShuffleExchangeExec` that `retargetKeyGroupedShuffle` can rewrite. The walk stops where the push stops, at a broadcast and at a shuffle. A `ShuffleQueryStageExec` or a `ReusedExchangeExec` cannot be rewritten, so it counts only when it already has the values. If either child fails the check, `checkKeyGroupCompatible` returns `None` and the outer join shuffles. * `populateSpjParams` applies that rewrite when it is given common partition values, and still does not descend into the shuffle's input. A scan below the shuffle is not filtered down to the merged keys (SPARK-60552). * The positions-only path has the same hazard. Pushing join-key positions into a scan that sits beside a nested one-side shuffle would regroup that scan and leave the shuffle behind. `hasKeyGroupedShuffleUnder` keeps that child out of the matched side, so it is shuffled onto the best spec instead. A shuffle that is itself the child is not treated as nested. Targets `branch-4.1` only. `master` and `branch-4.2+` regroup above the inner join with `GroupPartitionsExec` ([SPARK-55535](https://issues.apache.org/jira/browse/SPARK-55535)) and are unaffected. Found while reviewing #59599. ### Why are the changes needed? [SPARK-60460](https://issues.apache.org/jira/browse/SPARK-60460). On 4.0 and 4.1, with `spark.sql.sources.v2.bucketing.shuffle.enabled=true`: ```sql -- pa and pc PARTITIONED BY (id), pd not partitioned -- pa: 1, 2, 3 pd: 1, 2, 3 pc: 1, 2, 3, 4, 5 SELECT /*+ MERGE(pa, pd, pc) */ pa.id, pd.id, pc.id FROM pa JOIN pd ON pa.id = pd.id JOIN pc ON pa.id = pc.id ``` The inner join shuffles `pd` onto `pa`'s three keys. The outer join merges in `pc`'s five keys and pushes them into `pa`'s scan. `pd`'s exchange stays at three partitions. * AQE off: `IllegalArgumentException: Can't zip RDDs with unequal numbers of partitions: List(5, 3)` * AQE on: `requirement failed: PartitioningCollection requires all of its partitionings have the same numPartitions` The same push fails when the inner shuffle is wider than the outer key, for example an inner shuffle on `(store, dept)` under an outer join on `dept`. Those values are not the shuffle's keys, so the exchange cannot be retargeted and the outer join shuffles instead. `branch-4.2` and `master` return the three rows with one shuffle. ### Does this PR introduce _any_ user-facing change? Yes, for SPJ on `branch-4.1` (4.1.x) with `spark.sql.sources.v2.bucketing.shuffle.enabled=true`: * The query above, which failed during planning or at runtime, now returns the three matching rows. It keeps the one shuffle, retargeted onto the five merged keys. * A nested one-side shuffle whose keys are wider than the outer join key falls back to a shuffle and returns the correct rows. `branch-4.1` failed for that shape too. No change on `master` or `branch-4.2+`. ### How was this patch tested? Added two tests to `KeyGroupedPartitioningSuite`, both named after SPARK-60460. Each runs with AQE on and off. Neither is ignored. * Retarget: the query above. Asserts one shuffle, of five partitions, and the three result rows. It fails on `branch-4.1` with the zip error (AQE off) or the `PartitioningCollection` error (AQE on). * Wider key: the inner one-side shuffle is on `(store, dept)` and the outer join is on `dept` only, so the shuffle cannot be retargeted. Asserts the two matching rows. It fails on `branch-4.1` the same way. `KeyGroupedPartitioningSuite` and `EnsureRequirementsSuite` pass: 155 tests, 0 failures. ### Was this patch authored or co-authored using generative AI tooling? Yes: Grok 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]
