dongjoon-hyun commented on PR #58262:
URL: https://github.com/apache/spark/pull/58262#issuecomment-5440568415
I did a deep review pass over this change (8 review angles, each candidate
finding then adversarially verified against the code). Posting the findings
that survived verification, most severe first. The two type-divergence items
were established by full static traces but not executed end-to-end.
### Correctness
**1. Planning-time `ClassCastException`: `numPartitionsAfter` reads
reducer-rewritten keys at a stale declared type** (`EnsureRequirements.scala`,
`numPartitionsAfter`)
`GroupPartitionsExec.outputPartitioning` keeps the child's original
expressions over the reduced keys, and `KeyedShuffleSpec.reducers` synthesizes
type-changing reducers — `identity(ts)` reduced by `years` leaves
Integer-valued keys under a TimestampType-declared member, exactly the state
the comment above the memo describes. With tables partitioned `(identity(ts),
bucket(4, id))` and `(years(ts), bucket(4, id))`, an SPJ on both keys under
`pushPartValues` + `allowCompatibleTransforms` +
`allowKeysSubsetOfPartitionKeys`, and a window `PARTITION BY ts` above the
join: the new single-child path computes `clusterKeyPositions = {0}` (a strict
subset), the merges-nothing check calls `projectKeys`, and `key.row.get(0,
TimestampType)` on an Integer-backed row throws `ClassCastException` at
planning time. On master this path never projected keys, so the same query
planned (with the wrong result this PR fixes, but without crashing).
**2. The executed node can group with a different member than the one the
planner validated** (`GroupPartitionsExec.groupedPartitionsTuple`)
`splitKeyedPartitionings` ranks and count-validates a specific collection
member — memoized on `(positions, expressionDataTypes)` precisely because
members may disagree on declared types — but hands the node only
`joinKeyPositions`. `groupedPartitionsTuple` then `collectFirst`s the *first*
KP member and projects with *its* `expressionDataTypes`;
`PartitioningCollection.checkKeyedPartitioningInvariant` enforces the shared
`partitionKeys` reference and arity, not per-position types. With a
type-divergent collection (same shape as in item 1), the executed node can
produce a partition count different from the one the `requiredNumPartitions`
filter just validated, coalesce on differently-read key values, or hit the same
CCE when `outputPartitioning` is first computed. The memo key acknowledges the
divergence at planning; nothing reconciles it at execution — passing the
validated member (or its data types) to the node would.
**3. (pre-existing) The multi-child fallback applies the best spec's
`joinKeyPositions` to every compatible child**
This PR routes all co-partitioned subset projection to the multi-child block
(the `isCoPartitioned` gate), and that block's fallback calls
`withJoinKeyPositions(child, ...)` with the *best* spec's positions for every
compatible child, while `KeyedShuffleSpec.isCompatibleWith` matches
`joinKeyPositions` as `_`. For a non-SMJ/SHJ operator
(`checkKeyGroupCompatible` returns `None`), e.g. a cogroup on `i` with children
partitioned `(n, i)` and `(i, m)` over the same `i`-domain: both specs project
to grouped keys of `i` and are compatible, but applying the left's positions
`[1]` to the right selects `m` — misaligned partitions, silent wrong results.
The block is byte-identical on master, so this is not introduced here; noting
it because the new cogroup test only covers identical `(n, i)`/`(n, i)`
layouts, so the hole in the mechanism this PR designates as the projection
owner stays uncovered.
**4. (robustness) The `clusterKeyPositions` assert can fire for a
reference-free-expression KP**
With `allowKeysSubsetOfPartitionKeys` off, `groupedSatisfies`' final branch
`attributes.forall(...)` is vacuously true when every partition expression has
empty references; `clusterKeyPositions` then derives an empty position set over
non-empty expressions and `assert(positions.nonEmpty ||
kp.expressions.isEmpty)` throws at planning where the old code planned a plain
coalesce. Unreachable today because DSv2 scans gate on `supportsExpressions`,
but nothing at `KeyedPartitioning` construction enforces that for other
producers.
### Maintainability / efficiency
5. The "which partition-expression positions cover a cluster key" derivation
now lives in three places — `clusterKeyPositions`, `createShuffleSpec` via
`KeyedShuffleSpec.keyPositions`, and `groupedSatisfies`' subset branch — with
already-divergent semantics at the expression level, and `clusterKeyPositions`'
soundness depends on the `references.size == 1` invariant inside
`groupedSatisfies` with no back-pointer at that site. A shared helper on
`KeyedPartitioning` would keep the single-child and shuffle paths from drifting.
6. `splitKeyedPartitionings` returns two mutually exclusive `Option`s whose
exclusivity only the scaladoc enforces, and the caller pays with guarded
`.get`s (`satisfying.orElse(needsGrouping.map(_._1)).get`,
`needsGrouping.get._2`). A small private sealed ADT (`SatisfiedAsIs(kp)` /
`NeedsNode(kp, positions)`) would make the invariant type-enforced and drop the
`orElse`/`map`/`get` chain.
7. The keyed analysis — per-member `satisfies`/`groupedSatisfies`,
`clusterKeyPositions`, and the `O(numPartitions)` projections in
`numPartitionsAfter` — now runs before the caller's
`other.exists(_.satisfies(distribution))` short-circuit; master deferred all
keyed satisfaction work behind it. A collection mixing a satisfying
`HashPartitioning` with large KPs pays the full reconciliation per operator and
discards it.
8. `projectedNumPartitions` is scoped to one invocation, so the same
unchanged KP flowing through stacked operators (the no-node "merges nothing"
shape; windows with the same `PARTITION BY` but different `ORDER BY` don't
collapse) repeats the per-key projection at every operator. The identity arm
also recomputes the distinct count `KeyedPartitioning.apply` already derived
for `isGrouped` and discarded — a `@transient lazy val` on `KeyedPartitioning`
would compute it once per plan.
9. `k.satisfies(distribution)` already evaluates `groupedSatisfies`
internally for a grouped member (`satisfies0 = nonGroupedSatisfies ||
(isGrouped && groupedSatisfies)`), and the guard `satisfies ||
k.groupedSatisfies(distribution)` evaluates it a second time when false; each
pass rebuilds the `AttributeSet` and re-runs the `semanticEquals` scans.
Computing it once into a local would halve the per-member work.
For completeness, the review also probed and could not fault: the
`OrderedDistribution` arm's dropped positions (provably always `None` there),
the `sliding(2)` fix, the `requiredNumPartitions` filter-then-rank order, the
containment prune's tie semantics, and the memo's cross-member sharing — those
all check out as described in the PR.
🤖 Posted from a [Claude Code](https://claude.com/claude-code) review session
(Claude Fable 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]
