peter-toth opened a new pull request, #58963:
URL: https://github.com/apache/spark/pull/58963
### 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.
#### Backport to branch-4.3
The fix itself ports verbatim: `areKeysCompatible`,
`isExpressionCompatible`, `canReduceKeys`,
`TransformExpression.hasReducedKeys` and
`KeyedPartitioning.expressionsDescribeKeys` are all the same
shape on this branch. Three tailorings, all because `KeyLayout` is
master-only:
- `isCompatibleWith` keeps this branch's own trailing clause,
`partitioning.partitionKeys == otherPartitioning.partitionKeys`, in place
of
`layout.describesSameKeys`. That carries the type clause too, since
`InternalRowComparableWrapper.equals` compares the data types before the
rows.
- The `ShuffleSpecSuite` control that pins the fixture's premise compares
`partitionKeys` for the same
reason.
- `EnsureRequirements` has no `agreeingPairs` loop here. The lenient caller
is the single
`leftSpec.areKeysCompatible(rightSpec)` inside the push branch, and that
is the one call that names
`allowReduce = true`.
Nothing was omitted: all three tests came across, and **all three fail on
this branch without the fix**
(the end-to-end one returns 0 rows of 2, the spec-level one inverts, the
cogroup one stops shuffling),
which is the measurement that says the bug is live here.
`build/sbt 'catalyst/testOnly *ShuffleSpecSuite'` 24 pass.
`build/sbt 'sql/testOnly *KeyGroupedPartitioningSuite
*EnsureRequirementsSuite *ValidateRequirementsSuite
*ProjectedOrderingAndPartitioningSuite *GroupPartitionsExecSuite'` 292 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]
