[SPARK-59688][SQL] Fix wrong results when a storage-partitioned join skips the reduce on coinciding partition keys - #58943
peter-toth wants to merge 2 commits into
Conversation
…skips the reduce on coinciding partition keys `KeyedShuffleSpec.isCompatibleWith` read equal partition key rows as "the two sides are lined up" even for a pair whose keys are not in one key space yet, so the join paired the partitions up index by index and skipped the reduce that reconciles them. Ask `areKeysCompatible` with `allowReduce = false` there, so such a pair takes the push branch instead.
|
Thanks for the fix, @peter-toth. I went through this carefully and I agree with both the diagnosis and the place you landed the fix: One thing worth adding to the description: this also restores transitivity of Below are the things I'd like to see addressed, most important first. 1. The
|
…up path dongjoon-hyun's 15 items, then a simplify and a blind pass over the result. - `allowReduce` loses its default, so both call sites say which question they ask. - The unknown-keys branch asks the shared predicate rather than its own matcher, and asks for un-reduced keys instead of asserting them, so a producer this misses costs a shuffle rather than the query. - The end-to-end test validated an `AdaptiveSparkPlanExec`, which is a leaf, so it checked nothing. It validates the join subtree now. - New `EnsureRequirementsSuite` test for the cogroup path, which has no reduce branch. - `flip_low_bit` becomes a `SimpleFunction` in catalyst's test sources, so `InMemoryBaseTable` computes its keys by calling it and `ShuffleSpecSuite` uses the real function.
|
Thanks for going through it at that depth, @dongjoon-hyun. All 15 taken, and the transitivity point is in the description now. Pushed in 1. The assertion was validating an 2. Both failures are in the user-facing section, named, with what makes them reachable. 3. Kept the delegation, and the invariant is now part of the condition rather than an assert: 4. Rewritten: "must describe one key space: two bare references, the same transform function, or the two sides of one reduce". 5. Corrected in both places. The 6. Default dropped. 7. Said in one line, with the 8. New 9. The comment names 10. Done, one pass: 11. 12. The fixture moved to 13. 14. 15. Retitled Two things I did not take, both yours to call. A simplify pass argued the boolean should be two named methods instead, mirroring this file's The same pass wanted the invariant enforced at the producer, a |
…skips the reduce on coinciding partition keys ### 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. The parameter has no default, so every caller says which question it is asking. - `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. It now calls the shared predicate with the reduce disallowed, in the one pass the method already made. That comparison also asks that neither side's keys were reduced, which a marked layout never carries: asked rather than asserted, so that a producer the argument misses costs a shuffle rather than the query. The strict answer is also **transitive**, which the loose one is not: `bucket(12)` reduces onto `bucket(4)` and `bucket(4)` onto `bucket(8)`, while `bucket(12)` and `bucket(8)` have no reducer between them. `ValidateRequirements` compares every child against `specs.head` alone, so a three-child operator could be accepted on a chain that does not hold pairwise. `isSameFunction` is an equivalence, so that hole closes with this. ### Why are the changes needed? Wrong results. Take a table partitioned by `identity(id)` and one partitioned by a connector transform that permutes 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. What it takes to turn that into wrong rows is a transform that **permutes** its key space. A many-to-one transform cannot: the two key lists coincide only where it is the identity on them, which is why neither `bucket` nor `truncate` shows it, and why this has stayed latent. Note it takes no connector `Reducer` at all - the identity-versus-transform arm synthesizes one from the other side's ordinary transform - so a plain `ScalarFunction` is enough. 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. The blast radius is that arm, or two reducible transforms, plus coinciding key lists, plus `allowCompatibleTransforms.enabled`, which defaults to `false`: without it `canReduceKeys` is false, the two questions below agree at every position and this change is a no-op. ### 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 catalyst-internal, and catalyst is excluded from MiMa). Four consequences are worth stating, all on pairs whose key lists coincide and which now take the reduce instead of being read as they stand: - Where the coincidence was benign, the reducing side reports the target transform instead of its own, so a downstream operator 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 transform permutes the key space, and a lost optimization where the coincidence is benign; the `Reducer` API gives Spark nothing to tell the two apart with. For a cogroup the wrong answer was worse than for a join, since no equi-predicate drops the mispaired rows again, which is what the new `EnsureRequirementsSuite` test pins. - Two planning-time failures become reachable for a query that previously planned, both in place of wrong rows. `storagePartitionJoinIncompatibleReducedTypesError`, when a connector's `Reducer.resultType()` disagrees with the target transform's type, since the reduce now runs and compares them. And `cannotEvaluateExpressionError` for an identity-versus-transform pair whose transform is bound to something that is not a `ScalarFunction`, since `IdentityReducer` evaluates that transform and `TransformExpression.resolvedFunction` is empty there. ### How was this patch tested? Three new tests, each pinning a different layer, and all three fail on a revert of `partitioning.scala` alone: - `KeyGroupedPartitioningSuite`, end to end and the wrong-rows regression: `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 join subtree (the subtree, not the whole plan: the validator walks children and a query stage is a leaf, so validating an AQE plan checks nothing). - `ShuffleSpecSuite`, at the spec level, for 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, that the two layouts do describe the same keys and that one function over one key list is still compatible with itself, and a fourth case pins that a pair reduced together stays compatible, which is what keeps a chained storage-partitioned join from shuffling. - `EnsureRequirementsSuite`, for the path that has no reduce branch: a keyed cogroup over the same pair now lays the transform side out on the identity side's keys rather than reading the two as they stand, the rule stays idempotent over the result, and a control pins that two sides holding one key space are still read as they stand. `flip_low_bit` is a new fixture, and the minimal one for this: the transform has to keep its argument's type and to *permute* its key space, and no existing fixture does both (`bucket`, `days`, `years` and `signed_zeros` change the type, so such a pair is refused on the key types alone, while `string_self` is the identity and `truncate` is many-to-one). It is a `SimpleFunction`, so one object is both the unbound and the bound function, and it lives in catalyst's test sources so that `InMemoryBaseTable` computes its partition keys by calling `produceResult` rather than repeating the arithmetic, which is how the other four transforms in that file stay in step. `build/sbt 'catalyst/testOnly *ShuffleSpecSuite'` 31 pass. `build/sbt 'sql/testOnly *KeyGroupedPartitioningSuite *EnsureRequirementsSuite *ValidateRequirementsSuite *ProjectedOrderingAndPartitioningSuite *KeyGroupedPartitioningRuntimeFilterSuite *KeyGroupedPartitioningCatalystRuntimeFilterSuite *GroupPartitionsExecSuite'` 348 pass. `dev/lint-scala` clean. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) Closes #58943 from peter-toth/SPARK-59688-coinciding-partition-keys. Authored-by: Peter Toth <peter.toth@gmail.com> Signed-off-by: Dongjoon Hyun <dongjoon@apache.org> (cherry picked from commit 9daeea7) Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
|
Could you make a backporting PR to branch-4.3 and branch-4.2, @peter-toth ? There were some conflicts from branch-4.3. |
2 participants
What changes were proposed in this pull request?
KeyedShuffleSpec.isCompatibleWithanswers 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 noGroupPartitionsExec. It ends inKeyLayout.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. Underspark.sql.sources.v2.bucketing.allowCompatibleTransforms.enableditsisExpressionCompatibleadmits anAttributeReferenceagainst aTransformExpression, and two different but reducible transforms, becauseEnsureRequirementsreduces such a pair onto one key space before it pairs the partitions up. SoisCompatibleWithwas reading equal key rows as "lined up" for two sides whose keys are not in one space yet.This PR gives
areKeysCompatibleanallowReduceparameter, threaded intoisExpressionCompatible, and hasisCompatibleWithask withallowReduce = false:EnsureRequirementskeeps the loose question where it selects the member pair to plan on (agreeingPairs), because it is the caller that runs the reduce. The parameter has no default, so every caller says which question it is asking.isCompatibleWithgets the strict one, so a pair that still needs reconciling is not reported as compatible as it stands.compatibleAsIsis then false, and the join takes the push branch that computes the reducers and regroups both sides onto the merged keys.isExpressionCompatible's reduced-keys arm answers for them throughhasSameReducedKeyswhateverallowReducesays.areKeysCompatiblealready required this same single-key-space property before comparing key subsets, through a hand-written matcher. It now calls the shared predicate with the reduce disallowed, in the one pass the method already made. That comparison also asks that neither side's keys were reduced, which a marked layout never carries: asked rather than asserted, so that a producer the argument misses costs a shuffle rather than the query.The strict answer is also transitive, which the loose one is not:
bucket(12)reduces ontobucket(4)andbucket(4)ontobucket(8), whilebucket(12)andbucket(8)have no reducer between them.ValidateRequirementscompares every child againstspecs.headalone, so a three-child operator could be accepted on a chain that does not hold pairwise.isSameFunctionis an equivalence, so that hole closes with this.Why are the changes needed?
Wrong results. Take a table partitioned by
identity(id)and one partitioned by a connector transform that permutes its key space, sayflip_low_bit(id) = id ^ 1, with ids 0 and 1 in both:Both scans report the partition key list
[0, 1]ofLongType, sodescribesSameKeysholds, but the rows behind a key differ: the identity side's key 0 holdsid = 0while the transform side's holdsid = 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 noGroupPartitionsExecin the plan. The new end-to-end test measures exactly this onmaster.The
Reducercontract is what the fast path implicitly assumed more of. It saysr(f1(x)) = f2(x), and nothing aboutrleaving alone the keys it is applied to, so on a one-side reduce keykcan belong with keyr(k)on the other side. What it takes to turn that into wrong rows is a transform that permutes its key space. A many-to-one transform cannot: the two key lists coincide only where it is the identity on them, which is why neitherbucketnortruncateshows it, and why this has stayed latent. Note it takes no connectorReducerat all - the identity-versus-transform arm synthesizes one from the other side's ordinary transform - so a plainScalarFunctionis enough.The identity-versus-transform arm came in with SPARK-56182, which is on
branch-4.2,branch-4.3,branch-4.xandmaster, so this is not amaster-only bug. The blast radius is that arm, or two reducible transforms, plus coinciding key lists, plusallowCompatibleTransforms.enabled, which defaults tofalse: without itcanReduceKeysis false, the two questions below agree at every position and this change is a no-op.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
GroupPartitionsExecon each side and the identity side's keys are reduced ontoflip_low_bit. No configuration changes, and no public API changes (areKeysCompatibleis catalyst-internal, and catalyst is excluded from MiMa).Four consequences are worth stating, all on pairs whose key lists coincide and which now take the reduce instead of being read as they stand:
bucket(4, id)where it used to seebucket(8, id).pickCoPartitionTargetpairs children onisCompatibleWithalone. Such a pair is now shuffled onto one side instead of being read as it stands. That is the right answer where the transform permutes the key space, and a lost optimization where the coincidence is benign; theReducerAPI gives Spark nothing to tell the two apart with. For a cogroup the wrong answer was worse than for a join, since no equi-predicate drops the mispaired rows again, which is what the newEnsureRequirementsSuitetest pins.storagePartitionJoinIncompatibleReducedTypesError, when a connector'sReducer.resultType()disagrees with the target transform's type, since the reduce now runs and compares them. AndcannotEvaluateExpressionErrorfor an identity-versus-transform pair whose transform is bound to something that is not aScalarFunction, sinceIdentityReducerevaluates that transform andTransformExpression.resolvedFunctionis empty there.How was this patch tested?
Three new tests, each pinning a different layer, and all three fail on a revert of
partitioning.scalaalone:KeyGroupedPartitioningSuite, end to end and the wrong-rows regression:identity(id)againstflip_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, oneGroupPartitionsExecper side, exactly one of them carrying a reducer, andValidateRequirements.validateaccepting the join subtree (the subtree, not the whole plan: the validator walks children and a query stage is a leaf, so validating an AQE plan checks nothing).ShuffleSpecSuite, at the spec level, for both shapes the parameter gates: an identity side against that transform, and two bucket counts a reducer reconciles, are each admitted byareKeysCompatibleand refused byisCompatibleWith, in both directions. Two positive controls keep it honest, that the two layouts do describe the same keys and that one function over one key list is still compatible with itself, and a fourth case pins that a pair reduced together stays compatible, which is what keeps a chained storage-partitioned join from shuffling.EnsureRequirementsSuite, for the path that has no reduce branch: a keyed cogroup over the same pair now lays the transform side out on the identity side's keys rather than reading the two as they stand, the rule stays idempotent over the result, and a control pins that two sides holding one key space are still read as they stand.flip_low_bitis a new fixture, and the minimal one for this: the transform has to keep its argument's type and to permute its key space, and no existing fixture does both (bucket,days,yearsandsigned_zeroschange the type, so such a pair is refused on the key types alone, whilestring_selfis the identity andtruncateis many-to-one). It is aSimpleFunction, so one object is both the unbound and the bound function, and it lives in catalyst's test sources so thatInMemoryBaseTablecomputes its partition keys by callingproduceResultrather than repeating the arithmetic, which is how the other four transforms in that file stay in step.build/sbt 'catalyst/testOnly *ShuffleSpecSuite'31 pass.build/sbt 'sql/testOnly *KeyGroupedPartitioningSuite *EnsureRequirementsSuite *ValidateRequirementsSuite *ProjectedOrderingAndPartitioningSuite *KeyGroupedPartitioningRuntimeFilterSuite *KeyGroupedPartitioningCatalystRuntimeFilterSuite *GroupPartitionsExecSuite'348 pass.dev/lint-scalaclean.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)