peter-toth opened a new pull request, #58482: URL: https://github.com/apache/spark/pull/58482
### What changes were proposed in this pull request? This builds on #58335, which makes a reduced key-grouped partitioning report a transform that describes its keys for the two reducer shapes where one exists. When both sides of the join reduce there is none. The keys become `r1(f1(x))` = `r2(f2(x))`, a third key space that neither side's transform names. `bucket(12, id)` joined to `bucket(8, id)` is the flagship example. `BucketFunction.reducer` hands both sides `BucketReducer(4)`, the keys become `id % 4`, and both sides keep reporting the transforms they were built from. That is the gap #58335 names and leaves open, and this PR closes it by saying so rather than by inventing an expression. `TransformExpression` gains a fourth field, `reducedWith: Option[TransformFunctionId]`, which names the transform this one's keys were reduced together with. `KeyedShuffleSpec.reducersBothWays` is the only producer. In its both-sides-reduce branch it now reports `e1.reducedTogetherWith(e2)` instead of the bare `e1`. The marker holds no `Expression`, only a canonical name and a bucket count, so it is safe in a field canonicalization does not descend into. Putting it on the expression rather than on `KeyedPartitioning` is what keeps the change small. Every site that derives a partitioning already carries the expressions along. `AliasAwareOutputExpression` projects them, `GroupPartitionsExec.outputPartitioning` re-reports them, and `TransformExpression.withReference` re-targets them, so the marker is inherited, and dropped with the position it belongs to, without a line of new plumbing. `GroupPartitionsExec` needed no change at all. Four sites then refuse to reason about such keys, and a fifth refusal falls out of expression equality: - `KeyedShuffleSpec.isExpressionCompatible` does not compare marked keys by transform. Two of them are compatible when the same pair was reduced together, which is the pair the join produced. Anything else has to shuffle. - `KeyedShuffleSpec.canCreatePartitioning` does not shuffle another child onto marked keys, because that evaluates the reported expressions per row. - `KeyedShuffleSpec.reducersBothWays` does not reduce marked keys a second time. - `keysSatisfy` does not let marked keys satisfy an `OrderedDistribution`. An ordering is a claim about the key *values*, and nothing makes a reducer order-preserving. This one is a local guard. A marked position always carries a transform, and a SQL `ORDER BY` cannot name one. - `UnionExec.comparePartitioning` compares the children's expressions with `semanticEquals`, so a marked child no longer merges with an unmarked sibling reporting the same transform. No change was needed there. Clustering is the one thing a marked partitioning still satisfies, and that is sound. The keys remain a function of the same attributes, which is all `ClusteredDistribution` asks. Two things promised in the review of SPARK-59120 (#58420) are now delivered, since the marker answers the question they approximated: `KeyedPartitioning.expressionsDescribeKeyShape` and its use in `canCreatePartitioning` are replaced by `expressionsDescribeKeys`, and the three-reader list in the `keyDataTypes` scaladoc is gone. Keeping the type check as a second conjunct was considered and dropped. It is not sufficient, since `bucket(12)` and `bucket(8)` reducing onto `bucket(4)` keep their `IntegerType` and pass it. It is not necessary either. After a one-side reduce the reported expression is the target transform, and `EnsureRequirements` already refuses a reducer whose `resultType()` disagrees with the other side's key types, which are that transform's, with `STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES`. So the keys and the expressions agree at every reachable one-side reduce, and a both-sides reduce is marked. The unit test that pinned the type proxy is replaced by one that pins the marker, and it keeps the proxy's struct case. A partitioning whose key row was built at `struct<f>` while it declares `struct<g>`, which is what `createPartitioning` produces, is still accepted. Re-adding any type comparison to the gate fails that assertion. Two smaller things came with it. `TransformExpression.resolvedFunction` refuses a marked expression, so `eval` throws on one instead of computing the un-reduced transform and misrouting the row. That is a local gate rather than a live check, because every consumer of a reduced partitioning refuses it first and the write path never sees one. It also replaces the copy of the rule that prepends the bucket count as a literal argument in `DistributionAndOrderingUtils`, and both copies had to be touched here anyway, since the extractor grew a field. The `keyDataTypes` scaladoc also records where its no-key fallback stops being truthful. A marked partitioning can end up with no key, for instance when `v2BucketingPartitionFilterEnabled` intersects two sides that hold disjoint keys, and it then reports the un-reduced transform's type while the other leg of the same pairing reports the reducer's. SPARK-59176 tracks that, with the repro and two ways to fix it. The same query fails on a `ClassCastException` without this marker, so nothing regresses here. ### Why are the changes needed? Four failures, each one a test here, all measured on this PR's base. All of them need `spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled`, which is off by default, since that is what admits a reducer at all. 1. Two reduces onto different key spaces look co-partitioned. Four tables holding ids 0 until 24, bucketed by 12 and 8 on one side and by 12 and 18 on the other. The first pair reduces onto `id % 4`, the second onto `id % 6`, and both keep reporting `bucket(12, id)`. Joining the two returns **8 of 24 rows** with no shuffle. This is also what the pairing in the marker is for. Marking the keys without recording which pair produced them fixes everything else here and still returns those 8 rows. 2. Reducing keys twice. Three tables bucketed by 12, 8 and 6: the first join reduces onto `id % 4`, then the second derives a `bucket(6)` reducer from the reported `bucket(12, id)` and applies it to keys that are already reduced. `(id % 4) % 6` leaves the left keys alone while the right side moves to `id % 6`, so the query loses rows with no shuffle. 3. Reducing keys twice, the other way. Two `days`/`years` joins that each reduce both sides onto one year space, then joined to each other with different key sets on the two legs. The outer join derives reducers again and planning throws `ClassCastException: Long cannot be cast to Integer`. 4. Merging a reduced partitioning in a union. `(bucket12 JOIN bucket8) UNION ALL bucket12` reports `bucket(12, id)` on both sides of the union, so the two are merged although the join side's keys are `id % 4`. A `GROUP BY id` above it returns each id twice. The fifth refusal, `canCreatePartitioning`, guards a hazard the first one opens up. Once a marked partitioning is no longer compatible with anything but its own pairing, `EnsureRequirements` takes the one-side-shuffle path instead, and there it would happily shuffle the other child onto the reduced keys, which places rows by evaluating `bucket(12, id)` against partitions laid out by `id % 4`. Measured with tables bucketed by 12, 8 and 2 and `v2BucketingShuffleEnabled`, with the clause removed the query returns **8 of 12 rows**. ### Does this PR introduce _any_ user-facing change? Yes, it fixes the wrong results and the crash above. A query that used to reach one of them now shuffles instead, so its plan changes. It also refuses two plans that happen to be correct today, In both cases nothing can tell the two apart from the outside. The first is a second join onto an already reduced space. `bucket(12) JOIN bucket(8)` lands on `id % 4`. Meeting a `bucket(4, id)` or `bucket(2, id)` table, `BucketReducer` would compose correctly, since those counts divide 4. Meeting a `bucket(6, id)` table it would not, and that is failure 2 above. The `Reducer` API cannot say which of the two it is, so both now shuffle: `bucket12 JOIN bucket8 JOIN bucket4` goes from 0 shuffles to 2, with the same rows. This only arises where both sides reduced, i.e. where `gcd(a, b) < min(a, b)`. A chain in which one bucket count divides the other takes the one-side path and is unaffected. The second is two reduces that land on the same space through different pairings, e.g. `bucket(12)` with `bucket(8)` and `bucket(12)` with `bucket(20)`, both of which give `id % 4`. They are treated as different spaces and the join shuffles. Both cost a shuffle, not a wrong answer. A follow-up can give them back by letting a `ReducibleFunction` name the transform it reduces onto, which makes this whole shape disappear rather than be refused. Queries whose keys were not reduced on both sides are unaffected, and a join of two sides that were reduced together keeps its plan. ### How was this patch tested? Ten new tests. Six query tests in `KeyGroupedPartitioningSuite`, and one each in `TransformExpressionSuite`, `ShuffleSpecSuite`, `GroupPartitionsExecSuite` and `ProjectedOrderingAndPartitioningSuite`. One existing test changed, see below. Four of the six query tests fail on the base commit with wrong rows or a crash, listed as failures 1 to 4 above. A fifth, `another side is not shuffled onto reduced keys`, returns the right rows there with no shuffle. This PR costs it one shuffle, and it is what pins the `canCreatePartitioning` clause. Each refusal was ablated, and each ablation fails exactly the tests written for it: - `canCreatePartitioning`'s clause removed: `another side is not shuffled onto reduced keys` returns 8 of 12 rows. - the same-pairing allowance in `isExpressionCompatible` made unconditional, i.e. the refusal taken too far. `two sides reduced onto the same keys still join without a shuffle` fails, and so does the existing `SPARK-56164: Reducers with different result types to original keys`. - the pairing thrown away, so that any two marked keys count as one space. `two reduced partitionings are not compatible by their transforms` returns 8 of 24 rows. That is the ablation the test exists for, and it is why the marker records the pair rather than a bit. - the `reducersBothWays` guard removed. `two sides reduced together are not reduced a second time` throws the `ClassCastException`. No other test in the suite reaches that guard, and instrumenting it to throw on entry showed every one of them misses it. That is why the test is built the way it is, with different key sets on the two legs so that the join computes reducers at all. `SPARK-56164` is the existing test for the both-sides-reduce shape, and it gained one assertion, that the join's own requirements still validate, i.e. that the two sides which were reduced together are still co-partitioned. That is the no-regression half of this change, and it belongs on the test that already describes the shape rather than in a copy of it. The four unit tests cover what a query test states only indirectly. They are the pairing relation itself, the `canCreatePartitioning` refusal, that a node which reduces nothing (or reduces another position) inherits the marker from its child, and that a projection drops the marker with the position it belongs to. The last two hold by construction under this representation, which is exactly why they are worth keeping, since an earlier design of this fix carried the marker on `KeyedPartitioning` and got both wrong. Green: `ShuffleSpecSuite`, `DistributionSuite`, `TransformExpressionSuite`, `KeyGroupedPartitioningSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `ProjectedOrderingAndPartitioningSuite`, `ValidateRequirementsSuite`, `WriteDistributionAndOrderingSuite`, `PlannerSuite`, `UnionSuite`, `DataSourceV2Suite`, 488 tests in all. `dev/lint-scala` is clean. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code #### Backport to branch-4.2 The production change is the same. Four tailorings, all because the surrounding code arrived after this branch: - The `PartitioningCollection` invariant paragraph is left out. `checkKeyedPartitioningInvariant` came with SPARK-59057, so there is no invariant list here to document. - `V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS` is the name the config still has on this branch. - The unit test for `hasSameReducedKeys` is left out. `TransformExpressionSuite` does not exist here, because SPARK-58769 was never backported. The relation stays covered end to end by `two reduced partitionings are not compatible by their transforms`. - The projection test is left out. It needs `DummyLeafExecWithPartitioning`, which came with SPARK-46367. Run against the branch tip, four of the six new query tests fail. `a union does not merge an already reduced partitioning` passes here rather than failing, because this branch's `UnionExec` does not merge key-grouped partitionings at all. The test is kept, so the refusal stays covered if that ever lands here. Green here: `ShuffleSpecSuite`, `DistributionSuite`, `KeyGroupedPartitioningSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `ProjectedOrderingAndPartitioningSuite`, `ValidateRequirementsSuite`, `WriteDistributionAndOrderingSuite`, 312 tests in all. `dev/lint-scala` is clean. -- 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]
