dongjoon-hyun commented on PR #58943:
URL: https://github.com/apache/spark/pull/58943#issuecomment-5763022823
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: `isCompatibleWith`
really is the site that asks "are these lined up *as they stand*", and
`areKeysCompatible` answers a different question. I also verified the change is
monotonically stricter (`isSameFunction ⊆ isCompatible`, `canReduce=false ⊆
true`), that `canReduceKeys ⇒ v2BucketingPushPartValuesEnabled` keeps every
reduce-admitted pair on the push branch, and that the rule stays idempotent
across an AQE re-run.
One thing worth adding to the description: this also restores
**transitivity** of `isCompatibleWith`. The old `isCompatible`-based answer was
non-transitive (`bucket(12) ~ bucket(4)` and `bucket(4) ~ bucket(8)`, but
`bucket(12) ≁ bucket(8)`), while `ValidateRequirements` compares every child
against `specs.head` only. `isSameFunction` is an equivalence, so that latent
hole closes too.
Below are the things I'd like to see addressed, most important first.
---
### 1. The `ValidateRequirements.validate` assertion is vacuous
`KeyGroupedPartitioningSuite.scala:1357`
```scala
assert(ValidateRequirements.validate(plan), "the plan that leaves must hold
up")
```
`spark.sql.adaptive.enabled` defaults to `true` and neither the test nor
`KeyGroupedPartitioningSuite.sparkConf` disables it, so `plan` is an
`AdaptiveSparkPlanExec`, which `extends LeafExecNode`.
`ValidateRequirements.validate` is `plan.children.forall(validate) &&
validateInternal(plan)`, and on a childless node `validateInternal`
short-circuits to `satisfied = true` with the `children.length > 1`
co-partitioning branch never entered. Revert `allowReduce = false` and re-add
only this assertion — it still passes.
The sibling assertions survive only because `collectShuffles` /
`collectGroupPartitions` go through `AdaptiveSparkPlanHelper.collect`, whose
`allChildren` explicitly descends into `AdaptiveSparkPlanExec.executedPlan`;
`ValidateRequirements` does not.
This suite already documents the trap and works around it — see
`KeyGroupedPartitioningSuite.scala:5528-5534` ("a query stage is a leaf, so
validating an AQE plan checks nothing"), which validates
`collect(stripAQEPlan(...)) { case smj: SortMergeJoinExec => smj }.head`
instead. Since the PR description cites this assertion as evidence, I'd fix it
the same way.
### 2. Two planning-time exceptions become newly reachable
`partitioning.scala:2001`
Pairs that used to take the `compatibleAsIs` shortcut now enter the push
branch. Two hard failures follow for queries that previously planned:
- `KeyedPartitioning.reduceKeys` takes the reduced types from the
connector's `Reducer.resultType()`, so a connector violating the `r(f1(x)) =
f2(x)` type contract now hits
`storagePartitionJoinIncompatibleReducedTypesError`.
- For the identity-vs-transform arm, `IdentityReducer.reduce` calls
`bound.eval(...)`, and `V2ExpressionUtils.loadV2FunctionOpt` accepts any
`BoundFunction` with no `ScalarFunction` check — so a transform bound to a
non-`ScalarFunction` gives `resolvedFunction = None` and raises
`cannotEvaluateExpressionError`.
Both are better than the wrong rows they replace, but they are user-visible
and the "Does this PR introduce any user-facing change?" section doesn't
mention either.
### 3. The unknown-keys branch is loosened, not tightened
`partitioning.scala:2062`
This is the one place the diff relaxes the predicate. The deleted arm was
`case (l: TransformExpression, r: TransformExpression) => l.isSameFunction(r)`;
`isExpressionCompatible` now routes a marked + reduced pair to
`hasSameReducedKeys`. Left `bucket(12, a).reducedTogetherWith(bucket8)` vs
right `bucket(8, b).reducedTogetherWith(bucket12)` was refused before and is
accepted now — and two marked sides route their undeclared rows at `hash(key) %
numPartitions` in their *original* key spaces.
I traced all four marker producers and your unreachability claim does hold
today. But it rests on three cooperating sites in three files
(`canCreatePartitioning`'s `expressionsDescribeKeys`,
`PartitionGrouping.isIdentity` being false whenever a reducer slot exists, and
`ShuffledJoin.clearUnknownPartitionKeys`) with no assertion anywhere. Either
keep the explicit refusal or add `assert(!hasReducedKeys(l) &&
!hasReducedKeys(r))`.
### 4. Rewritten comment 3.3 is factually false
`partitioning.scala:1994`
> 3.3 each pair of partition expressions at the same index must share the
same transform function.
`isCompatibleWith` does not require that. The reduced-keys arm answers
`hasSameReducedKeys`, which is true for differing `functionId`s:
`bucket8.reducedTogetherWith(bucket4)` has `numBucketsOpt = Some(8)` and
`bucket4.reducedTogetherWith(bucket8)` has `Some(4)`, so `isSameFunction` is
false while `hasSameReducedKeys` is true — and
`ShuffleSpecSuite.scala:757-759`, added by this PR, asserts that pair *is*
compatible. The first arm also admits two plain attributes, which share no
transform at all. The replaced wording ("compatible transform functions") was
loose but true.
The risk is concrete: someone reading 3.3 as a guarantee and tightening the
reduced-keys arm to `isSameFunction` sends every chained SPJ back to a shuffle
— exactly what `ShuffleSpecSuite:756-759` exists to protect.
### 5. The `allowReduce` scaladoc overstates what triggers the bug
`partitioning.scala:2028`
> it takes a connector reducer that reorders its key space to turn this into
wrong rows
For the identity-vs-transform arm the same paragraph names, there is no
connector `Reducer` at all — `reducersBothWays` synthesises
`IdentityReducer(t.withReference(a))` from the *other side's ordinary
transform*. Your own `FlipLowBitFunction` is a plain `ScalarFunction[Long]`,
not a `ReducibleFunction`, and it produces the wrong rows. All it takes is a
transform that preserves its argument type and is not order-preserving.
Same overstatement in the fixture comment at
`transformFunctions.scala:290-294`: "result type is its input type, and not the
identity on its inputs" does not exclude the existing `TruncateFunction`
(`StringType -> StringType`, not the identity). What actually makes
`flip_low_bit` the minimal fixture is that it is a **permutation** of the key
space; `truncate` is many-to-one, so it cannot mispair rows a raw-column join
would match. Worth stating that way — it's the sharper reason.
### 6. `allowReduce` defaults to the unsafe value
`partitioning.scala:2030`
There are exactly two production call sites: `EnsureRequirements.scala:703`
(wants `true`, and is the caller that runs the reduce) and
`partitioning.scala:2001` (wants `false`). Every other consumer asks the
as-they-stand question. A future site written as
`spec.areKeysCompatible(other)` — the natural spelling — silently gets the
reduce-allowed answer and reintroduces SPARK-59688 with no compile error and no
failing test.
Dropping the default and writing `allowReduce = true` at the one
`EnsureRequirements` call site costs a word; only `ShuffleSpecSuite`'s existing
assertions need the extra argument. Note those assertions ride the default too,
so flipping it later would silently reverse their meaning rather than fail to
compile.
---
### Smaller items
7. **`partitioning.scala:2100`** — the delegation silently widens the
marked-keys gate from `(_: AttributeReference, _: AttributeReference)` to `(_:
LeafExpression, _: LeafExpression)`. Unreachable today, but this is the gate
protecting the `partitionKeys` subset comparison, and reverting this hunk alone
leaves every test in the diff green — `ShuffleSpecSuite:662` is the only
coverage and all its inputs are `AttributeReference`/`TransformExpression`. At
least say in the comment that it's a widening rather than a restatement.
8. **Cogroup path has no test.** `checkKeyGroupCompatible` returns `None`
for every non-SMJ/SHJ parent, so `FlatMapCoGroupsInPandasExec` and friends
reach `pickCoPartitionTarget` as their only path. With
`v2BucketingShuffleEnabled=true`, that path had the same bug, and the failure
mode is worse than for joins because no equi-predicate filters the mispaired
groups — the join returns zero rows, a cogroup hands the user's function rows
that don't belong together. This PR fixes it silently;
`EnsureRequirementsSuite`'s keyed-cogroup tests (1502, 2042) use no
compatible-but-different transform pair.
9. **`partitioning.scala:2049`** — the unreachability argument names
`createPartitioning` as where marked layouts come from, but
`PartitioningCollection.fromPartitionings` also produces them: it ORs the
marker across members (line 1483) and stamps the OR'd flags onto the canonical
layout, and its `require` guards `describesSameKeys` but not the marker. The
operative guard is `ShuffledJoin.clearUnknownPartitionKeys`. The conclusion
holds; a reader checking a future change would check the wrong site.
10. **`partitioning.scala:2060`** — when `allowReduce = false` the
unknown-keys `forall` is textually identical to the outer one at 2039-2040, so
it is pure duplicate work; and since `isExpressionCompatible(l, r, false) ⇒
isExpressionCompatible(l, r, true)` on every arm, the outer pass is subsumed in
the `allowReduce = true` case too. Hoisting `val unknownKeys =
partitioning.mayContainUnknownPartitionKeys ||
other.partitioning.mayContainUnknownPartitionKeys` and calling the outer loop
with `allowReduce && !unknownKeys` collapses both into one pass, and states the
intent (the inner loop exists only to *downgrade* a lenient caller) which isn't
visible today.
11. **`partitioning.scala:2098`** — `val canReduce = allowReduce &&
canReduceKeys` is eager, so the `(Leaf, Leaf)` arm now evaluates `SQLConf.get`
plus three conf reads where it previously read none. That's the plain
`identity(id)` vs `identity(id)` shape, hit once per partition expression per
candidate pair over the `leftCandidates × rightCandidates` cross product. `lazy
val`, or a local `def`, restores the old cost.
12. **`InMemoryBaseTable.scala:347`** — `value ^ 1L` re-implements
`FlipLowBitFunction.produceResult` with a comment as the only sync, the fourth
such pair in this file. They already disagree on NULL: the `case (value: Long,
LongType)` type test doesn't match null so the write path throws, while
`ApplyFunctionExpression.eval` has no null guard and `getLong(0)` on a null
slot returns 0, mapping NULL to key `1L`. Avoidable — moving
`FlipLowBitFunction` into the existing
`sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/functions/`
package lets this arm call
`FlipLowBitFunction.produceResult(InternalRow(value))` (`InternalRow` is
already imported), and lets `ShuffleSpecSuite` use the real function too.
13. **`transformFunctions.scala:281`** — `UnboundFlipLowBitFunction` is
`UnboundSignedZerosFunction` with two strings changed. `SimpleFunction`
(`@since 4.2.0`) exists for exactly this, with precedent at
`DataSourceV2FunctionSuite.scala:753`; `object FlipLowBitFunction extends
SimpleFunction with ScalarFunction[Long]` drops the wrapper and still works
with `withFunction`. The copied type check is dead anyway — `loadV2FunctionOpt`
swallows the `UnsupportedOperationException` into `None`.
14. **`ShuffleSpecSuite.scala:715`** — the inline `flipFn` is behaviourally
identical to the `FakeBucket` this PR just hoisted, for everything this suite
observes (`resultType` and "is a `TransformExpression`"); the function is never
evaluated here, so the `a ^ 1` semantics the comment explains aren't exercised.
The same test builds a `FakeBucket` 29 lines later. Item 12 would let you use
the real function instead.
15. **`KeyGroupedPartitioningSuite.scala:1332`** — the title has no JIRA ID
(172 of 188 tests in this file are `SPARK-xxxxx: ...`; the 16 that aren't are
original feature tests, not bug fixes), so `git grep SPARK-59688` won't find
the regression test. It also lands between the 4th and 5th of five consecutive
`SPARK-59045:` tests; moving it after the last one (before the `SPARK-59121`
pair) keeps both the topical adjacency and the ticket grouping.
--
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]