peter-toth opened a new pull request, #59165: URL: https://github.com/apache/spark/pull/59165
### What changes were proposed in this pull request? A storage-partitioned join now refuses to pair a partition transform whose argument is not a bare column, i.e. is an expression or a struct field: - `TransformExpression.isSameFunction` and `TransformExpression.reducers` say no when either transform has an argument that is not an `Attribute`. `isCompatible` now goes through `reducers`, so the one check covers the reducer path too. - `KeyedShuffleSpec.canCreatePartitioning` asks the same question, so such a layout is never the one other children are shuffled onto. `createPartitioning` replaces a transform's argument with the other child's cluster key, which drops whatever surrounds the key. - The `keyPositions` doc now says what its reference-based mapping does not tell apart. ### Why are the changes needed? With `spark.sql.sources.v2.bucketing.shuffle.enabled` on, this query returns 0 rows instead of 7. `t1(id)` and `t3(x)` are partitioned by `bucket(4, ...)`, `plain(b)` is not partitioned, and each holds the values 0 to 7: ```sql SELECT t1.id, p.b, t3.x FROM t1 JOIN plain p ON t1.id = p.b + 1 JOIN t3 ON p.b = t3.x ``` 1. The first join shuffles `plain` onto `t1`'s partitioning. `KeyedShuffleSpec.createPartitioning` builds that partitioning from `plain`'s join key, so it is `bucket(4, b + 1)`. This join is correct. 2. The second join clusters `plain` on the bare `b`. `KeyedShuffleSpec.keyPositions` maps a partition expression to a cluster key through its reference, so `bucket(4, b + 1)` counts as a function of `b`, the counterpart of `x`. 3. `isSameFunction` compares only the function name and the bucket count. So `bucket(4, b + 1)` is the same as `t3`'s `bucket(4, x)`, and the join pairs the partitions as they stand. A row with `b = x` sits in bucket `(b + 1) % 4` on one side and in `x % 4` on the other. The comparison ignores which column an argument is, and relies on `keyPositions` to pair the columns up. That is sound only when the argument is the column itself. `bucket(4, b + 1)` is a function of `b`, but not the same function of it as `bucket(4, x)` is of `x`. A scan never reports such a transform, so only a one-side shuffle can produce one. The same problem shows up in three more shapes, all measured on master: - **The reducer path.** Under `spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled`, a `bucket(8, x)` third table loses the rows the same way. - **A struct field.** A join key `p.s.a` gives `bucket(4, s.a)`, which `keyPositions` pairs through `s`. Joining two such sides on `s` pairs `bucket(4, s.a)` with `bucket(4, s.b)` and returns 0 of 8 rows. Shuffling another side onto such a layout builds `bucket(4, s)`, which fails on executors with a `ClassCastException`. - **Shuffling onto the layout.** In `plain p JOIN t1 ON p.b + 1 = t1.id JOIN xy q ON t1.id = q.x AND p.b = q.y`, the second join shuffles `xy` onto the first join's `bucket(4, b + 1)` layout. That builds `bucket(4, y)` for `xy`, drops the `+ 1`, and returns 0 of 8 rows. The one-side shuffle of transform expressions came with SPARK-48012, in 4.0.0. ### Does this PR introduce _any_ user-facing change? Yes, it fixes the wrong results and the crash above. Such a query now shuffles the side instead of pairing it or shuffling onto it. That costs shuffles in a few shapes whose plan was already correct: - When two sides carry the same expression shape, e.g. `bucket(4, b + 1)` and `bucket(4, c + 1)` joined on `b = c`, both sides are shuffled. - When a join's output offers such a layout next to a usable one, e.g. `t1 JOIN plain p ON t1.id = p.b + 1 JOIN xy q ON t1.id = q.x AND p.b = q.y`, the next join takes 3 shuffles instead of 2. `ShuffleSpecCollection.canCreatePartitioning` requires every member to be usable, so the usable `bucket(4, id)` layout is ruled out too. The shuffling-onto-the-layout test below pays the same. A follow-up can give these back by describing a transform's argument relative to its column, so that `b + 1` and `c + 1` compare equal while `b + 1` and `x` do not. ### How was this patch tested? New tests, each of which fails on master: - `KeyGroupedPartitioningSuite`, "SPARK-59887: a side shuffled onto a transform of an expression is not paired as it is". It runs the query above against a `bucket(4)` third table, and with `allowCompatibleTransforms` against a `bucket(8)` one that the join reduces, and asserts two shuffles. Master returns 0 of 7 rows. - `KeyGroupedPartitioningSuite`, "SPARK-59887: a side shuffled onto a transform of a struct field is not paired as it is". Master returns 0 of 8 rows. - `KeyGroupedPartitioningSuite`, "SPARK-59887: a side is not shuffled onto a transform of a struct field". Master fails with the `ClassCastException`. - `KeyGroupedPartitioningSuite`, "SPARK-59887: a side is not shuffled onto a transform of an expression". Master returns 0 of 8 rows. - `TransformExpressionSuite`, "SPARK-59887: a transform of an expression is not the same as any transform". It covers `isSameFunction`, `reducers` in both directions and `isCompatible` for transforms over `a + 1` and over a struct field, with a column-argument control. Each of the three clauses was removed on its own: - without the `isSameFunction` one, the `bucket(4)` case returns 0 of 7 rows and the struct field test 0 of 8; - without the `reducers` one, the `bucket(8)` case returns 0 of 7 rows; - without the `canCreatePartitioning` one, both struct field tests fail with the `ClassCastException`, and the first test fails on the shuffle count. The third test passes with any one of them removed. Both the `isSameFunction` clause and the `canCreatePartitioning` clause keep `xy` off that layout, the first because the layout then does not pair with its own side. Also ran `TransformExpressionSuite`, `ShuffleSpecSuite`, `DistributionSuite`, `KeyGroupedPartitioningSuite` (and the other `KeyGroupedPartitioning*` suites), `WriteDistributionAndOrderingSuite`, `PlannerSuite`, `ProjectedOrderingAndPartitioningSuite`, `GroupPartitionsExecSuite` and `EnsureRequirementsSuite`, 583 tests in all, plus `dev/lint-scala`. ### 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]
