peter-toth opened a new pull request, #58943: URL: https://github.com/apache/spark/pull/58943
### What changes were proposed in this pull request? `KeyedShuffleSpec.isCompatibleWith` answers whether two key-grouped children are lined up as they stand, so that a join can pair their partitions up index by index with no shuffle and no `GroupPartitionsExec`. It ends in `KeyLayout.describesSameKeys`, which compares the two sides' partition key rows and their types. That comparison only says what it looks like it says when both sides' keys are values of the same thing. The predicate it consults first, `areKeysCompatible`, is deliberately looser than that. Under `spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled` its `isExpressionCompatible` admits an `AttributeReference` against a `TransformExpression`, and two different but reducible transforms, because `EnsureRequirements` reduces such a pair onto one key space before it pairs the partitions up. So `isCompatibleWith` was reading equal key rows as "lined up" for two sides whose keys are not in one space yet. This PR gives `areKeysCompatible` an `allowReduce` parameter, threaded into `isExpressionCompatible`, and has `isCompatibleWith` ask with `allowReduce = false`: - `EnsureRequirements` keeps the loose question where it selects the member pair to plan on (`agreeingPairs`), because it is the caller that runs the reduce. - `isCompatibleWith` gets the strict one, so a pair that still needs reconciling is not reported as compatible as it stands. `compatibleAsIs` is then false, and the join takes the push branch that computes the reducers and regroups both sides onto the merged keys. - A pair an earlier join reduced together is unaffected: those keys are in one space already, and `isExpressionCompatible`'s reduced-keys arm answers for them through `hasSameReducedKeys` whatever `allowReduce` says. - The unknown-partition-keys branch inside `areKeysCompatible` already required this same single-key-space property before comparing key subsets, through a hand-written matcher with the same reasoning in its comment. It now calls the shared predicate with `allowReduce = false` rather than repeating it. The predicate's reduced-keys arm cannot fire there, and the comment says why. ### Why are the changes needed? Wrong results. Take a table partitioned by `identity(id)` and one partitioned by a connector transform that reorders its key space, say `flip_low_bit(id) = id ^ 1`, with ids 0 and 1 in both: ```sql SELECT t1.id, t1.data, t2.data FROM t1 JOIN t2 ON t1.id = t2.id ``` Both scans report the partition key list `[0, 1]` of `LongType`, so `describesSameKeys` holds, but the rows behind a key differ: the identity side's key 0 holds `id = 0` while the transform side's holds `id = 1`. The join pairs partition 0 with partition 0 and finds no match in either pair, so it returns **0 rows instead of 2**, with no shuffle and no `GroupPartitionsExec` in the plan. The new end-to-end test measures exactly this on `master`. The `Reducer` contract is what the fast path implicitly assumed more of. It says `r(f1(x)) = f2(x)`, and nothing about `r` leaving alone the keys it is applied to, so on a one-side reduce key `k` can belong with key `r(k)` on the other side. A modulo-style reducer, which is the shape a bucket count reduce takes, does leave a coinciding key alone, which is why this has stayed latent: for `bucket(8)` against `bucket(4)` the key lists can only coincide on ids below 4, where the reduce is the identity. It takes a reducer that permutes its key space to turn the dishonest answer into wrong rows. The identity-versus-transform arm came in with SPARK-56182, which is on `branch-4.2`, `branch-4.3`, `branch-4.x` and `master`, so this is not a `master`-only bug. ### Does this PR introduce _any_ user-facing change? Yes, it is a bug fix. The query above returns its 2 rows instead of 0, and it still runs without a shuffle: the pair goes through the reduce, so the plan gains a `GroupPartitionsExec` on each side and the identity side's keys are reduced onto `flip_low_bit`. No configuration changes, and no public API changes. `areKeysCompatible` is `private[sql]`-reachable catalyst code and the new parameter has a default, so every existing caller is unchanged. Two plan changes are worth stating, both on pairs whose key lists coincide: - Where the coincidence was benign, the join now reduces rather than pairing raw keys, and the reducing side reports the target transform instead of its own. A downstream operator therefore sees the coarser claim, for instance `bucket(4, id)` where it used to see `bucket(8, id)`. - A co-partitioned operator that is not a sort-merge or shuffled-hash join, a cogroup for instance, has no reduce branch at all: `pickCoPartitionTarget` pairs children on `isCompatibleWith` alone. Such a pair is now shuffled onto one side instead of being read as it stands. That is the right answer where the reducer permutes the key space and a lost optimization where the coincidence is benign, and the `Reducer` API gives Spark nothing to tell the two apart with. ### How was this patch tested? Two new tests, and both fail on a revert of `partitioning.scala` alone: - `KeyGroupedPartitioningSuite`, end to end: `identity(id)` against `flip_low_bit(id)` with ids 0 and 1 in both tables. On the base it returns 0 of 2 rows; with the fix it returns both, with 0 shuffles, one `GroupPartitionsExec` per side, exactly one of them carrying a reducer, and `ValidateRequirements.validate` accepting the plan. `flip_low_bit` is a new test fixture, the smallest transform that is type-preserving and not the identity on values, which is what lets its keys coincide with a raw column's while the rows behind a key differ. No existing fixture has that shape: `bucket`, `days`, `years` and `signed_zeros` all change the type, so such a pair is refused on the key types alone, and `string_self` and `truncate` are the identity on any data whose key lists coincide. - `ShuffleSpecSuite`, at the spec level: both shapes the parameter gates. An identity side against that transform, and two bucket counts a reducer reconciles, are each admitted by `areKeysCompatible` and refused by `isCompatibleWith`, in both directions. Two positive controls keep it honest: the two layouts do describe the same keys, and one function over one key list is still compatible with itself. A fourth case pins that a pair reduced together stays compatible, which is what keeps a chained storage-partitioned join from shuffling. The suite's `FakeBucket`, a real `ReducibleFunction`, is hoisted to the suite so both tests use it. `build/sbt 'catalyst/testOnly *ShuffleSpecSuite'` 31 pass. `build/sbt 'sql/testOnly *KeyGroupedPartitioningSuite *EnsureRequirementsSuite *ValidateRequirementsSuite *ProjectedOrderingAndPartitioningSuite *KeyGroupedPartitioningRuntimeFilterSuite *KeyGroupedPartitioningCatalystRuntimeFilterSuite'` 318 pass. `dev/lint-scala` clean. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) -- 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]
