[SPARK-58996][SQL] Fix SPJ partially clustered data correctness when EnsureRequirements re-runs

[ulysses-you]

What changes were proposed in this pull request?

EnsureRequirements is not idempotent for a storage-partitioned join that uses a partially
clustered distribution: re-running it on a plan it already produced stacks a second
GroupPartitionsExec on top of the first, and the result duplicates rows.

The re-run is not an AQE quirk: AdaptiveSparkPlanExec builds one EnsureRequirements
instance, and ConvertSortMergeJoinToShuffledHashJoin and OptimizeSkewedJoin both hand the
whole tree back to that same instance after rewriting some other join, all within one
queryStagePreparationRules pass. Idempotency is therefore a requirement of the current rule
composition. On that second pass a join child is SortExec(GroupPartitionsExec(...)) rather
than a bare scan. A partially clustered KeyedPartitioning reports isGrouped = false by
design, so the distribution step treats it as satisfied "only after grouping" and adds a plain
GroupPartitionsExec on top; applyGroupPartitions then writes the join's
expectedPartitionKeys and distributePartitions into that fresh outer node.

The alignment is therefore re-derived from an already-aligned layout. The inner node replicates
an input partition across the expected partitions, and the outer node concatenates those
replicas back into a single partition before replicating again, so every row of the replicated
side is emitted twice:

outer GroupPartitions groupedPartitions = [([1],[0,1]), ([1],[0,1]), ([2],[2])]
inner GroupPartitions groupedPartitions = [([1],[0]), ([1],[0]), ([2],[1])]

The fix makes the rule reuse what an earlier pass inserted instead of deriving from it:

• innermostGroupPartition is the single descent through the nodes this rule itself inserted
  directly above the child: a GroupPartitionsExec and a local SortExec. A
  GroupPartitionsExec hidden behind any other node belongs to a different operator and is
  never reused. Both the rewrite and the reads below go through it.
• rewriteGroupPartitions rewrites the innermost GroupPartitionsExec and drops what sits
  above it, reproducing the plan a single pass would have produced. Only
  applyGroupPartitions calls it, reached from checkKeyGroupCompatible, which runs for joins
  alone.
• unwrapGroupPartitions peels down to the pre-alignment plan along the same descent, so the
  statistics-based replicate-side choice and the original partition keys below are read from
  that plan on every pass. Peeling one level reads them from the sort this rule added, which
  carries no logicalLink and reports the aligned layout, and deterministically flips the
  replicate-side choice. The positions projecting the original keys come from the innermost
  grouping as well, for the same reason applyGroupPartitions keeps them, and the
  partition-count fallback compares the pre-alignment split counts for the same reason.
• applyGroupPartitions keeps the joinKeyPositions a reused node already holds: they were
  computed against the node child's raw partition keys, while the incoming ones were computed
  against the node's own, already projected report, and applying them would project a second
  time.
• withJoinKeyPositions, reached from the children loop for every multi-child operator and not
  just joins, reuses only a topmost GroupPartitionsExec and does not descend.
• The ShuffleExchangeExec site strips every grouping this rule inserted instead of one level:
  a replicating grouping repeats every row, so none of them may feed a shuffle.

Why are the changes needed?

The join returns duplicated rows. Reproduction, against a DSv2 catalog that reports KeyGroupedPartitioning with one split per row (the in-memory test catalog does this with numRowsPerSplit = 1; on a real connector the same shape is a partition value with more than one data file whose files are not combined into a single task):

spark.sql.sources.v2.bucketing.enabled=true
spark.sql.sources.v2.bucketing.pushPartValues.enabled=true
spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled=true
spark.sql.adaptive.maxShuffledHashJoinLocalMapThreshold=100m

CREATETABLEsp1 (id BIGINT, data STRING) PARTITIONED BY (id);
INSERT INTO sp1 VALUES (1, 'aa'), (1, 'ab'), (2, 'bb'); -- two splits for id = 1CREATETABLEsp2 (id BIGINT, data STRING) PARTITIONED BY (id);
INSERT INTO sp2 VALUES (1, 'p'), (2, 'q');
CREATETABLEnp1 (id BIGINT, data STRING);
INSERT INTO np1 VALUES (7, 'x');
CREATETABLEnp2 (id BIGINT, data STRING);
INSERT INTO np2 VALUES (7, 'y');
SELECT/*+ MERGE(a, b) */a.idAS k FROM sp1 a JOIN sp2 b ONa.id=b.idUNION ALLSELECTc.idAS k FROM np1 c JOIN np2 d ONc.id=d.id;

Returns (1, 1, 1, 1, 2, 7); the correct answer is (1, 1, 2, 7).

One detail is load-bearing. The np1/np2 branch exists only to create a materialized shuffle stage, which is what makes ConvertSortMergeJoinToShuffledHashJoin fire and re-run EnsureRequirements over the whole plan -- the storage-partitioned join has no shuffle of its own, so it cannot trigger the re-run by itself.

A single pass is self-consistent: within one ensureDistributionAndOrdering call the distribution step creates the GroupPartitionsExec and applyGroupPartitions rewrites that same node. The stacking only appears once the rule is applied to its own output.

Does this PR introduce any user-facing change?

Yes. It fixes wrong results (duplicated rows) for a storage-partitioned join under a partially clustered distribution when AQE re-runs EnsureRequirements. For the query above, the result changes from (1, 1, 1, 1, 2, 7) to the correct (1, 1, 2, 7).

It also changes the plan in one corner: the partition-count fallback that chooses the replicate side when there are no plan statistics now compares the pre-alignment split counts instead of the aligned report's distinct keys. On a first pass of a query where one side holds more than one split per key, the replicated side can differ from earlier releases (the side with fewer splits is replicated, which is the cheaper choice).

How was this patch tested?

New tests:

• KeyGroupedPartitioningSuite, "partially clustered join keeps its row count when EnsureRequirements re-runs" -- the end-to-end query above, whose row count was wrong before the fix. It also asserts the storage-partitioned side stays shuffle-free and no GroupPartitionsExec is stacked over another.
• KeyGroupedPartitioningSuite, "partially clustered join keeps its replicate-side choice when EnsureRequirements re-runs" -- the smaller side is replicated by plan statistics on the first pass. The data keeps rows on both sides of the multi-split key, so the re-run's flipped choice returns wrong results already on master, and the smaller side's five splits for id = 1 pin the guard: the flipped distribute side overflows the expected count of one (padTo never truncates) and the join sides end up with an unequal number of partitions.
• KeyGroupedPartitioningSuite, "partially clustered subset-key join keeps its join key positions when EnsureRequirements re-runs" -- the left table is partitioned by (extra, id) and the join uses only id, the second partition key. On master the stacked re-grouping duplicates rows; without the position fix the re-run projects the pre-alignment keys with aligned-space positions and throws the PartitioningCollection invariant.
• EnsureRequirementsSuite, "only a local sort is looked through when reusing GroupPartitionsExec" -- covers the bare, local-sort, stacked and global-sort cases, so a GroupPartitionsExec serving a global sort's OrderedDistribution is never reused.
• EnsureRequirementsSuite, "withJoinKeyPositions reuses only a topmost GroupPartitionsExec" -- non-join multi-child operators reach this path, so it must not descend to nodes that may belong to another operator.
• EnsureRequirementsSuite, "reusing a GroupPartitionsExec keeps its tags" -- tags are instance state a copy does not carry, so every rewrite copies them back.
• EnsureRequirementsSuite, "the shuffle site strips every grouping the rule inserted" -- pins the ShuffleExchangeExec site's descent directly, as no query reaches its stacked shape today.
• EnsureRequirementsSuite, "the replicate-side fallback counts pre-alignment partitions" -- forces the statistics-less fallback and pins that it counts the pre-alignment splits rather than the aligned reports.
• EnsureRequirementsSuite, "a single-child operator over a partially clustered layout still gets grouped" -- pins the intentional non-idempotence of the children loop's wrap for single-child operators.

Existing suites run locally: KeyGroupedPartitioningSuite, GroupPartitionsExecSuite, ProjectedOrderingAndPartitioningSuite, EnsureRequirementsSuite, PlannerSuite, AdaptiveQueryExecSuite, DataFrameJoinSuite, DisableUnnecessaryBucketedScanWithoutHiveSupportSuite(AE).

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code

[ulysses-you]

cc @peter-toth@cloud-fan thank you

[peter-toth]

Let me check this tomorrow.

[uros-b]

The fix looks correct and complete to me, adding @peter-toth to reviewers

[dongjoon-hyun]

Reviewed with a focus on the re-run (idempotency) semantics this PR introduces. The fix is correct for the reproduced scenario, but the second pass still re-derives its inputs from the already-aligned layout, and the new write-to-the-inner-node behavior surfaces that in a few configurations. Inline comments below; two findings that fall outside the diff hunks:

1. EnsureRequirements.scala:277 (case gpe: GroupPartitionsExec => ShuffleExchangeExec(newPartitioning, gpe.child)) still strips only one level. On a re-run the child at this site is GroupPartitionsExec(SortExec(local, GPE_replicated(scan))), so this strips the harmless outer node and plants the shuffle over the replicating inner GPE. A partially clustered GPE with distributePartitions = false emits Seq.fill(numSplits)((key, splits)) — its raw output stream contains numSplits copies of every row for a key, which is only safe when consumed against the distributing join side. Shuffling that stream duplicates rows. I could not construct a query where pass 2 fails checkKeyGroupCompatible after pass 1 succeeded, so this may be unreachable today, but nothing asserts "a replicating GPE is never shuffled", and this PR makes SortExec(GPE) an acknowledged shape while leaving this sibling site with the old one-level assumption. Worth routing through the same innermost-descent helper.

2. Design note. Four sites now each encode their own answer to "where is the GPE that owns the alignment": line 119 wraps unconditionally, lines 277 and 730 strip one level, line 763 descends two shapes — and the inline findings below are exactly the sites that disagree. Any partitioning-preserving node not in the line-763 whitelist (or a third-party queryStagePrepRules rewrite) reintroduces the SPARK-58996 duplication with no diagnostic. A deeper fix — letting an aligned GPE's reported partitioning carry enough information that the distribution step adds nothing on a re-run, or asserting the difference between "no GPE below" and "GPE behind an unrecognized node" before wrapping — would make the rule idempotent by construction rather than per-shape.

[dongjoon-hyun]

Re-run still flips the replicate-side decision, and the flipped flags now land on the inner node over the raw scan — unequal join partition counts.

unwrapGroupPartitions (line 730) was not updated alongside this helper: on the re-run it returns the ER-inserted SortExec, whose logicalLink is None (ER-built sorts carry no LOGICAL_PLAN_TAG, and AQE's inherited-tag propagation stops at the join node). The stats branch at lines 614-629 is therefore skipped, and the fallback leftPartKeys.size < rightPartKeys.size compares two already-aligned key lists that are equal by construction (alignToExpectedKeys emits exactly numSplits entries per key in both modes) — so pass 2 deterministically picks replicateLeftSide = false, flipping pass 1 whenever stats chose true.

Concretely: left key k = 3 raw splits (small bytes), right = 1 split (large). Pass 1: replicateLeftSide = true, expected counts {k: 1}, both sides emit 1 partition. Pass 2: flip; numExpectedPartitions is read off the aligned layout ({k: 1}), and this helper writes distributePartitions = true into the inner GPE over the raw scan — splits.map(Seq(_)).padTo(1, ...) never truncates, so the left emits 3 partitions while the replicating right emits 1 → Can't zip RDDs with unequal numbers of partitions from SortMergeJoinExec (or duplicated/misaligned rows in multi-key variants).

Pre-PR the flipped values landed on a fresh outer node over the aligned child (1 vs 1, correct), so this shape is a regression. The new test avoids it only because sp1 > sp2 makes pass 1 already choose replicateLeftSide = false. Suggestion: give unwrapGroupPartitions the same descent as this helper so pass 2 reads stats and the original partitioning from the pre-alignment plan (which also fixes the stale numExpectedPartitions source at lines 664-675).

[ulysses-you]

Confirmed -- a regression of this PR's descent. unwrapGroupPartitions now shares the rewrite descent (innermostGroupPartition), so the statistics and the original partition keys read from the pre-alignment plan on every pass. New regression test has the replicated side hold more splits for id = 1 than the other side, so a flipped decision overflows padTo on the distribute side and the join sides end up with an unequal number of partitions. 5af8667

[dongjoon-hyun]

joinKeyPositions index-space mismatch when reusing the inner node.

With spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled=true, pass 2's KeyedShuffleSpec computes joinKeyPositions against the GPE's reported partitioning, whose expressions are already projected (GroupPartitionsExec.outputPartitioning line 76: joinKeyPositions.fold(k.expressions)(_.map(k.expressions))). But the node this now rewrites applies the positions to its child's rawKeyedPartitioning (groupedPartitionsTuple, GroupPartitionsExec.scala:128-138).

Example: table partitioned by (a, b), join on b. Pass 1: inner GPE.joinKeyPositions = Some(Seq(1)), reported expressions [b]. Pass 2 computes Some(Seq(0)) in projected space and this code writes it through the local sort onto the inner node — which now projects position 0 of the raw [a, b], i.e. groups by a while expectedPartitionKeys hold b-values. alignToExpectedKeys's keyMap.getOrElse(key, Seq.empty) misses on every key: empty partitions, silently dropped rows. Same exposure via withJoinKeyPositions (line 258 path).

Pre-PR the fresh outer wrapper's child exposed [b] at index 0, so Seq(0) was self-consistent (the bare-GPE case had the hazard already; the descent extends it to the SortExec and stacked shapes). Suggestion: compose the positions with the reused node's existing joinKeyPositions (or skip the overwrite when they are defined) so the stored positions stay in the raw child's index space. No existing test combines allowJoinKeysSubsetOfPartitionKeys with a second EnsureRequirements pass, so nothing in the suites would catch this.

[ulysses-you]

Confirmed. applyGroupPartitions keeps the positions a reused node already holds: they were computed against the child's raw partition keys and stay authoritative for this join, while the incoming ones are in the node's already projected report. The new regression test partitions the left table by (extra, id) and joins on the second partition key -- joining on the first column would not discriminate, since the positions agree on both passes there. 5af8667

[dongjoon-hyun]

The safety argument here is not accurate.rewriteGroupPartitions is also reached from withJoinKeyPositions (line 806), which is called at line 258 from the generic children.zip(requiredChildDistributions) loop — gated only on childrenIndexes.length > 1, not on the parent being a join. CoGroupExec and FlatMapCoGroupsInBatchExec require ClusteredDistribution on two children and qualify.

What actually keeps a non-grouped GPE away from that path today is a config coincidence: KeyedShuffleSpec.canCreatePartitioning requires !v2BucketingPartiallyClusteredDistributionEnabled (partitioning.scala:1374-1376). But the OrderedDistribution branch (lines 100-106) produces a non-grouped GPE under v2BucketingAllowSorting alone, and RemoveRedundantSorts can delete the global sort shielding it — leaving a shape where line 258 under a cogroup could drop the outer grouping and leave the required distribution unsatisfied.

Suggest rewording this paragraph to state the real invariant (and/or guarding the drop, e.g. only drop when the inner node's partitioning still satisfies the required distribution), so the safety condition is enforced rather than documented.

[ulysses-you]

Confirmed. The descent and drop are now confined to applyGroupPartitions, which is reached from checkKeyGroupCompatible for joins alone; withJoinKeyPositions reuses only a topmost node again, since every multi-child operator reaches it. The scaladoc states that invariant, and a unit test pins the top-most behavior. 5af8667

[dongjoon-hyun]

Nit: the local-sort look-throughs elsewhere use the named form — AQEUtils.getRequiredDistribution, ConvertSortMergeJoinToShuffledHashJoin.stripSort, and simpleCosting all write case s: SortExec if !s.global. The positional SortExec(_, false, _, _) (mirrored in the new test) hard-codes the arity and field order: a new SortExec parameter breaks this match at compile time in two places, and a reordering of global/child would break it silently.

cases: SortExecif!s.global =>
rewriteGroupPartitions(s.child)(f).map(c => s.copy(child = c))

(with copyTagsFrom(s) if you switch from withNewChildren to copy, since withNewChildren propagates tags).

[ulysses-you]

Done -- case s: SortExec if !s.global. Kept withNewChildren, which already propagates the sort's tags. 5af8667

[dongjoon-hyun]

Both callers pass the identical three-line lambda (bind g.copy(...), copyTagsFrom(g), return). Tag propagation is an invariant of "rewrite this node", not of either caller — consider folding copyTagsFrom(g) into the helper's GroupPartitionsExec branch so callers pass a bare copy:

rewriteGroupPartitions(plan)(_.copy(joinKeyPositions =Some(positions)))
.getOrElse(GroupPartitionsExec(plan, joinKeyPositions =Some(positions)))

A future third caller then cannot forget the tag copy (a bug class with no compile error and no test).

[ulysses-you]

Done -- copyTagsFrom moved into the helper's GroupPartitionsExec branch, callers pass a bare copy. A unit test now pins tag retention on the bare, local-sort and top-most rewrite paths. 5af8667

[dongjoon-hyun]

Two suggestions to keep this regression test from going vacuously green:

1. The ALTER TABLE split trick is unnecessary. The suite's own createTable helper (line 261) passes numRowsPerSplit = 1, and InMemoryBaseTable.withData splits per row (rows.size >= numRowsPerSplit), so a single INSERT INTO ... VALUES (1, 'aa'), (1, 'ab'), (2, 'bb') yields two splits for id = 1 — the idiom every other partially clustered test here uses (e.g. line 812). This is the only raw CREATE TABLE testcat in the suite, and the current construction couples the test to an incidental schema-evolution branch of the fixture: if that behavior changes, the test keeps passing while no longer covering the fix.

2. checkAnswer alone cannot fail on unfixed code if the setup drifts.Seq(1, 1, 2, 7) is the correct answer for any plan, including a plain shuffle join. The test only exercises the fix while AQE stays on by default, ConvertSortMergeJoinToShuffledHashJoin fires on the np branch, and the SPJ branch stays shuffle-free — none of which is asserted. Adding assert(collectShuffles(df.queryExecution.executedPlan).isEmpty) for the SPJ side plus an assertion that no GroupPartitionsExec has a GroupPartitionsExec descendant (the stacking this PR eliminates) would make it fail loudly instead.

[ulysses-you]

Done both. The test is rebuilt on the suite's createTable idiom (numRowsPerSplit = 1), and it asserts the storage-partitioned side stays shuffle-free and no GroupPartitionsExec is stacked over another. 5af8667

[dongjoon-hyun]

Note that this test never reaches the changed code and passes byte-identically on the base commit: with one child, the childrenIndexes.length > 1 block is skipped, so neither applyGroupPartitions nor withJoinKeyPositions (and hence rewriteGroupPartitions) runs. The grouping it observes comes from the untouched wrap at line 118-119.

It is still a meaningful guard — if someone later "generalizes" the reuse into that wrap (rewriteGroupPartitions(child)(identity).getOrElse(GroupPartitionsExec(child))), the returned child would stay non-grouped and this assertion fails, which is exactly what the scaladoc argues must not happen. But as named, a reader will assume it covers the drop logic in rewriteGroupPartitions. Suggest a comment (or rename) making explicit that it pins the intentional non-idempotence of the line-119 wrap for single-child operators, not the new helper.

[ulysses-you]

Done -- the test's comment now states that it pins the children loop's wrap for single-child operators and never runs rewriteGroupPartitions (with one child the multi-child block is skipped). 5af8667

[peter-toth]

There are quite a few fixes/changes in progress for SPJ. Can we wait a bit with this one and sort out the simpler ones first? Also, my #58262 will change the shape a bit.

[peter-toth]

#58262 has been merged, but I'm still thinkig about an alternative. Let me come back to you tomorrow.

[peter-toth]

Thanks @ulysses-you, and sorry for the wait. I said I would come back with an alternative, so this is
what I have: I could not find one that works. Below is the mechanism as I traced it, the three shapes I
tried, and how I would close @dongjoon-hyun's three correctness threads if this direction stays.

What it does.EnsureRequirements is re-run on plans it already produced, so a join child arrives
as SortExec(GroupPartitionsExec(...)) rather than as the scan the first pass aligned. The
distribution step then adds a plain GroupPartitionsExec on top, because a partially clustered
KeyedPartitioning reports isGrouped = false by design, and applyGroupPartitions writes the freshly
merged alignment into that outer node. Pass 1's alignment survives underneath, so the inner node
distributes a key's splits and the outer one concatenates them back and replicates again.
rewriteGroupPartitions descends to the innermost node, rewrites that, and drops what sits above it.

One reframing, because it changes how the fix should be read. This is not an AQE quirk and it does
not need a reOptimize. AdaptiveSparkPlanExec builds one EnsureRequirements instance, and
ConvertSortMergeJoinToShuffledHashJoin and OptimizeSkewedJoin both hand the whole tree back to that
same instance after rewriting some other join. So the double application happens inside a single
queryStagePreparationRules pass, and idempotency is a requirement of the current rule composition
rather than a nice property. Worth saying in the description, since it is what makes this a bug and not
a hardening.

Confirmed by instrumentation, printing inner-first from inside the helper on your own test:

GPE expected=Some(2) distribute=true <- inner, pass 1, already aligned
local SortExec
GPE expected=None distribute=false <- outer, fresh, gets dropped

Still live after #58262. I cherry-picked your commit onto today's master: clean pick, identical
diffstat, 176 tests green, and partially clustered join keeps its row count when EnsureRequirements re-runs still fails with Results do not match without the production change. #58262 did not mask it.

What I tried instead

Three shapes: two measured dead, one only partway. I am recording them so nobody proposes them again,
including me.

1. Skip the wrap at the decision site. In the children loop, do not wrap a child that already
   carries a GroupPartitionsExec through a local sort. Your repro passes and 145 SPJ tests pass, but
   a single-child operator over a partially clustered layout still gets grouped fails with
   newChild.outputPartitioning.satisfies(distribution) was false. Your own guard test catches it.
   The reason is structural: at that point the rule cannot tell a join child about to be re-aligned from
   an aggregate or a window that genuinely needs its non-grouped input grouped, because
   childrenIndexes and checkKeyGroupCompatible run afterwards.

2. Normalize the join children in checkKeyGroupCompatible and plan from the raw children, so a
   second pass sees the same input as the first. Your repro passes under it too, and then 44 of 126
   KeyGroupedPartitioningSuite tests fail, including ordinary SPJ: partitioned join: exact distribution, all eight SPARK-42038 cases, SPARK-47094 compatible buckets, SPARK-56046
   reducers. createKeyedShuffleSpec gates on satisfies, which gates on isGrouped, and the node the
   first pass inserted is exactly what makes a non-grouped child satisfy. The information the method
   needs, that the child can satisfy once a node is added, is what the wrapping encodes.

3. Ask keysMaySatisfy instead of satisfies, on top of that normalization: give
   checkKeyGroupCompatible its own spec builder, and compare on toGrouped's layout for a non-grouped
   source. The toGrouped half is required, not cosmetic, because createShuffleSpec builds the
   projected and grouped view only when allowKeysSubsetOfPartitionKeys is on, so otherwise the key
   comparison and the key merging would run on unaligned keys. That takes the 44 failures down to 7,
   which says the satisfies gate was the blocker.

Why the repro passing under 1 and 2 is not evidence: in that shape the scan's own KeyedPartitioning is
already grouped, because pushing partition values gives one partition per key, and the non-groupedness
comes from alignToExpectedKeys on the node's output.

Shape 3 is not a working alternative either, and I want to be exact about that: one of the 7 is your
own partially clustered join keeps its row count when EnsureRequirements re-runs, failing with
Results do not match, so it reintroduces wrong rows on the very shape this PR fixes. Another puts a
GroupPartitionsExec where a fully clustered join needs none. Both say the "child partitionings not
modified" fast path and the merged-key and expected-count computation have to be reworked together
with the gate, not after it. That is a redesign of the SPJ decision, with its own ticket, and whatever
happens here should not wait for it.

For the record, the direction I would take that in: the children loop should not put a placeholder
GroupPartitionsExec above a co-partitioned child at all. Decide that child's node once, in the place
that already knows the join's merged keys, and let the spec creation ask whether a child can satisfy
after grouping so that it can reason about the raw child. A second pass then has nothing to re-derive,
because the node is built once from an input that does not change between passes, and both
rewriteGroupPartitions' drop and unwrapGroupPartitions fall out rather than being fixed.

If this direction stays, how I would close @dongjoon-hyun's three correctness threads

Suggestions, not requests, and I have not measured any of the three. What I do think is that none of
them needs index arithmetic.

• The unwrapGroupPartitions regression
  is real, and the clean fix is one descent helper serving both the stats read and the rewrite, rather
  than two depths of the same question. Note the plan-reading sites do not need satisfies at all, so
  normalizing for them is safe even though it is not safe for the spec derivation.
• The joinKeyPositions index space
  is the one that pinches, and I could not construct it. The positions are written by pass 1's own
  applyGroupPartitions, so on pass 2 the node's reported partitioning really is projected. I would not
  translate positions between the two spaces. The invariant is that the innermost node's positions were
  computed in the raw space and are the authoritative ones, so a later pass keeps them instead of
  recomputing: applyGroupPartitions sets positions only when the node has none. That composes with the
  drop this PR already does.
• The withJoinKeyPositions safety argument
  I found independently before reading the thread. Passing the caller's knowledge that the parent is a
  join is enough, and it is information the caller already has.

Findings

Nothing blocking. Two Non-blocking, inline.

1. rewriteGroupPartitions' scaladoc justifies the local-sort restriction only through the
   OrderedDistribution argument, and does not state the general rule.
2. The rule's set of "this is my own output" cases is incomplete and inconsistent across four sites,
   and this PR adds a fifth. Not this PR's to fix, worth naming.

[peter-toth]

[Non-blocking, docs] This justifies the restriction through the OrderedDistribution case, which
is the sharpest consequence, but it does not state the rule the descent actually follows: descend only
through a node this rule itself inserted directly above this child. Without that, a later reader has no
reason not to generalise the descent to a ProjectExec, and it looks like an omission rather than a
decision.

It is worth stating because generalising it would be wrong. I instrumented the descent over
KeyGroupedPartitioningSuite: a non-SortExec node hides a GroupPartitionsExec from it 6 times, and
refusing to descend is right every time. They are Project > SortMergeJoin > Sort > GroupPartitions and
Project > Filter > Window > WindowGroupLimit, where the hidden node belongs to a different operator,
so reusing it would move another operator's alignment.

[ulysses-you]

Done -- the scaladoc states the rule the descent follows (only nodes this rule itself inserted directly above the child, with the shapes you measured) and now lives on the shared descent helper innermostGroupPartition, per your suggestion of one descent serving both the statistics read and the rewrite. 5af8667

[peter-toth]

[Non-blocking, altitude] Not this PR's to fix, but worth having on the record, because the reason
this fix is needed at all is that the rule has no single answer for "this node is my own output". There
are four shapes today and this makes five:

site	what it does with the rule's own output
the ShuffleExchangeExec case	s.copy(outputPartitioning = ...), reuses in place
the GroupPartitionsExec case beside it	ShuffleExchangeExec(newPartitioning, gpe.child), unwraps it
unwrapGroupPartitions	peels one level, only to reach logicalLink
the children loop plus applyGroupPartitions	did not account for it at all, which is this bug
rewriteGroupPartitions, new here	descends through a local sort and drops what is above

Consolidating them is the follow-up I described in the body: stop inserting a placeholder above a
co-partitioned child, and decide that child's node once where the merged keys are known. That retires
both rewriteGroupPartitions' drop and unwrapGroupPartitions. I am not asking for it here.

[ulysses-you]

Noted -- the five-shape table and the decide-the-node-once direction go into a separate follow-up ticket.

[ulysses-you]

Thanks @peter-toth@dongjoon-hyun -- everything addressed in 5af8667 (rebased onto current master):

Correctness findings (all three confirmed against the code, all regressions of this PR):

• unwrapGroupPartitions now shares the rewrite descent, so the statistics-based replicate-side choice and the original partition keys read from the pre-alignment plan on every pass.
• applyGroupPartitions keeps the joinKeyPositions a reused node already holds; the incoming ones are in the node's already projected report.
• The descent and drop are confined to the join path (applyGroupPartitions); withJoinKeyPositions is back to reusing only a topmost node.

Per @peter-toth's suggestion the descent is unified in one helper, innermostGroupPartition, shared by the statistics read, the rewrite and the ShuffleExchangeExec site (which now strips every grouping the rule inserted).

Shape 1 was considered gated on join parents, but without the outer wrap pass 2 fails createKeyedShuffleSpec's satisfies gate and silently degrades to a shuffle join; the no-placeholder shape belongs to the decide-the-node-once redesign, which goes into the follow-up with the five-shape consolidation. The description now carries the re-run reframing.

Tests: 3 end-to-end (row count, replicate-side choice, subset-key positions) and 5 unit tests. Each mechanism is pinned by a shape verified to fail with its fix reverted in isolation. Local suites green: KeyGroupedPartitioningSuite, EnsureRequirementsSuite, GroupPartitionsExecSuite.

[peter-toth]

Re-checked through 5af8667 — findings 1, 2 resolved (the descent scaladoc now states the rule and sits on the shared innermostGroupPartition; the five-shape consolidation goes to the follow-up), nothing regressed. KeyGroupedPartitioningSuite, EnsureRequirementsSuite and GroupPartitionsExecSuite are 203 green on this head.

One Blocking finding, and it is the half of r3864640486 and r3864640493 that the fix does not reach. unwrapGroupPartitions now reads the pre-alignment plan, but the positions those keys are projected with still come from the aligned node's spec. At round 1 I said I could not construct that index space. I can now, and on this head it throws.

Blocking

• 3.Pre-alignment keys projected with aligned-space positions (new):partiallyClusteredSpec.joinKeyPositions is computed against the reused node's already projected report and applied to the raw partitioning unwrapGroupPartitions just returned. Under allowKeysSubsetOfPartitionKeys a re-run throws All KeyedPartitionings in a PartitioningCollection must have equal partitionKeys; master returns duplicated rows for the same query. Measured, with a fix and a test. inline: EnsureRequirements.scala:670

Non-blocking

• 4.The replicate-side test does not fail on master (new): it passes on 08757655cc2 run on its own, so it pins the guard rather than the bug, and its comment's "before the fix" does not describe master. Two lines of data make it fail there too, 27 rows against a correct 7. inline: KeyGroupedPartitioningSuite.scala:3525
• 5.A wrong expected count is not caught at the node (new):alignToExpectedKeys derives isGrouped from numSplits while padTo never truncates, so a count below the real split count emits duplicate keys under an isGrouped = true report. That is why finding 3 surfaces three steps later at a PartitioningCollection invariant instead of at the node. Deriving isGrouped from the emitted partitions, or asserting splits.size <= numSplits, would localise it — sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:131-146. Follow-up, not this PR.
• 6.Affects version and backports (new): the JIRA affects 5.0.0 only, but the one-level unwrapGroupPartitions and the top-node applyGroupPartitions are identical in v4.2.0-rc6, branch-4.2, branch-4.3 and branch-4.x, so this is a wrong-results bug in a released line. branch-4.3 and branch-4.x carry ConvertSortMergeJoinToShuffledHashJoin(ensureRequirements), so the new tests reproduce there as they are. branch-4.2 has only OptimizeSkewedJoin(ensureRequirements), so a 4.2 backport needs a skew-based repro.

Minor

• 7.Inconsistent shape notation in the descent scaladoc (new): the first shape ends at the hidden GroupPartitions, the second stops at WindowGroupLimit, so a reader cannot tell where the node sits in the second — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:743-744.
• 8.applyGroupPartitions' scaladoc still describes the one-level behaviour (new): "either the given plan node ... or we can create a new one" predates the descent — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:806-811.

[peter-toth]

Finding 3. The pre-alignment plan is read here, but the positions its keys are projected with at line 677 still come from the aligned node's spec, so the two are in different index spaces.

partiallyClusteredSpec is leftSpec/rightSpec, built at line 167 from children(i).outputPartitioning. On a re-run that is the reused GroupPartitionsExec's report, whose expressions are already projected (GroupPartitionsExec.scala:84). So joinKeyPositions comes back as Some(Seq(0)) over one expression, while originalKeyedPartitioning is the raw [extra, id] that unwrapGroupPartitions just returned. projectKeys(Seq(0)) then reads extra, numExpectedPartitions is keyed by extra values, no merged id key matches, and every count stays at 1.

That is the same invariant r3864640493 is about and that applyGroupPartitions now honours: the innermost node's positions were computed in the raw space and stay authoritative. Both readers of the pre-alignment plan need them, not only the writer. r3864640486 named this site too — "which also fixes the stale numExpectedPartitions source at lines 664-675".

Measured.multi_part partitioned by (extra, id) with two extra values sharing id = 1, joined on id, pushPartValues + partiallyClusteredDistribution + allowKeysSubsetOfPartitionKeys on, the same np1/np2 branch the other tests use:

• this head throws java.lang.IllegalArgumentException: requirement failed: All KeyedPartitionings in a PartitioningCollection must have equal partitionKeys, from PartitioningCollection.fromPartitionings through SortMergeJoinExec.outputPartitioning inside the second pass. The distribute side reports [1, 1, 2], because padTo(1, ...) never truncates its two splits, against the replicate side's [1, 2].
• master returns 6 rows, [1,10] [1,10] [1,11] [1,11] [2,20] [7,0], of which 4 are correct.

The fix below is measured: the query returns the correct 4 rows, and KeyGroupedPartitioningSuite + EnsureRequirementsSuite stay green (188).

val (partiallyClusteredChild, partiallyClusteredPositions) =if (replicateLeftSide) {
(unwrappedRight,
innermostGroupPartition(right).flatMap(_._1.joinKeyPositions)
.orElse(rightSpec.joinKeyPositions))
} else {
(unwrappedLeft,
innermostGroupPartition(left).flatMap(_._1.joinKeyPositions)
.orElse(leftSpec.joinKeyPositions))
}

and at line 677:

valprojectedOriginalPartitionKeys= partiallyClusteredPositions
.fold(originalKeyedPartitioning.partitionKeys)(
originalKeyedPartitioning.projectKeys(_)._2)

partiallyClusteredSpec has no other use, so it goes away.

The subset-key test covers this shape once its left table has a second split per join key, which also makes that test fail on master rather than only with the orElse reverted:

 sql("INSERT INTO testcat.ns.multi_part VALUES (1, 10, 'x'), (1, 11, 'x2'), (2, 20, 'y')")

with checkAnswer(df, Seq(Row(1L, 10L), Row(1L, 11L), Row(2L, 20L), Row(7L, 0L))).

[ulysses-you]

Fixed as suggested: partiallyClusteredPositions comes from the innermost grouping, falling back to the spec's when there is none, and partiallyClusteredSpec is gone. Revert-verified -- without it the updated subset test throws exactly the measured PartitioningCollection invariant. The second split for id = 1 makes the test fail on master as well, where the stacked re-grouping duplicates rows. 83ed835

[peter-toth]

Finding 4. Measured: this test passes on the merge base 08757655cc2 when run on its own, so it does not reproduce a master defect. It does pin the guard — revert unwrapGroupPartitions to the one-level peel on this head and it fails — which is the method you describe in the conversation. The comment above it claims something stronger: "before the fix the re-run flipped the choice and emitted an unequal number of partitions per join side". On master this query returns the correct answer.

The reason is the data. The flip does happen on master, but sp_small has no id = 2 rows, so the id = 2 partitions the other side over-replicates never join anything and the extra copies stay invisible.

Putting rows on both sides of the multi-split key makes it a master-level test. Measured: 27 rows on 08757655cc2 where 7 are correct, and it passes on this head.

 createTable("sp_small", spColumns, Array(identity("id")))
sql("INSERT INTO testcat.ns.sp_small VALUES (1, 'a'), (2, 'b')")
createTable("sp_large", spColumns, Array(identity("id")))
sql("INSERT INTO testcat.ns.sp_large VALUES "+"(1, 'p'), (2, 'q1'), (2, 'q2'), (2, 'q3'), (2, 'q4'), (2, 'q5')")

with checkAnswer(df, Row(1L) +: Seq.fill(5)(Row(2L)) :+ Row(7L)). Left stays the smaller side, so pass 1 still replicates it by statistics. On master pass 2 flips, the right side's five id = 2 partitions each hold all five of its splits, and the join returns 25 rows for id = 2 where 5 are correct.

[ulysses-you]

Adopted -- the test now fails on master as well. One adjustment: with the single-split shape on its own, reverting the guard stayed self-consistent in my arms (the flipped side's one split per key fills the aligned counts), so the data keeps the five splits for id = 1 on the smaller side too: with them the flipped distribute side overflows the expected count of one, and reverting the guard fails the PartitioningCollection invariant (measured). 83ed835

[ulysses-you]

thank you @peter-toth , everything from the re-review is addressed in 83ed835:

Finding 3 (Blocking) -- fixed as suggested: the positions projecting the pre-alignment keys come from the innermost grouping, falling back to the spec's when there is none; partiallyClusteredSpec is gone. Revert-verified: without it the updated subset test throws the measured PartitioningCollection invariant, and the second split for id = 1 makes the test fail on master as well.

Finding 4 -- the replicate-side test is now master-level and still pins the guard; data details in the inline reply.

Finding 5 -- agreed, follow-up: deriving isGrouped from the emitted partitions, or asserting splits.size <= numSplits, would localise a wrong expected count at the node instead of three steps downstream.

Finding 6 -- updating the JIRA affects versions to 4.2.0 / 4.3.0; backports follow, with a skew-based repro for branch-4.2 as noted.

Finding 7 / 8 -- the shape notation now ends both chains at the hidden GroupPartitionsExec, and the applyGroupPartitions scaladoc describes the descent.

Local suites green: KeyGroupedPartitioningSuite, EnsureRequirementsSuite, GroupPartitionsExecSuite (203).

[peter-toth]

Re-checked through 83ed835. Findings 3, 4, 6, 7 and 8 are resolved, nothing regressed.

Measured on this head:

• KeyGroupedPartitioningSuite, EnsureRequirementsSuite and GroupPartitionsExecSuite are 203 green.
• All three end-to-end tests now fail on the merge base 08757655cc2 with Results do not match, so finding 4 is closed for tests 2 and 3 as well.
• Reverting only this commit's partiallyClusteredPositions on this head fails the subset-key test and nothing else, so that test pins finding 3's mechanism on its own.

Nothing blocking left. The three items below sit in the replicate-side block and none of them is a defect on this head. They are all about reading one fact from one place, which is the class findings 3 and 4 came out of.

Findings 2 and 5 are still waiting on their follow-up ticket, and I could not find one filed. Finding 5 is self-contained, take it if you like. I would rather keep the decide-the-node-once direction behind finding 2, since I already have the two probes and their numbers.

Non-blocking

• 9.One descent for the pre-alignment plan and its positions (new):unwrapGroupPartitions(right) at :618 and innermostGroupPartition(right) at :671 walk to the same node, and the pair is sound because they do. Taking both from one value makes that structural instead of a coincidence of two call sites. inline: EnsureRequirements.scala:668
• 10.The partition-count fallback still compares aligned key counts (late catch): the statistics read moved to the pre-alignment plan, leftPartKeys.size < rightPartKeys.size at :644 did not. It never fired in 187 tests, so this is about holding the invariant at both branches rather than a defect. inline: EnsureRequirements.scala:663
• 11.reducers is overwritten where joinKeyPositions is kept (late catch): on a re-run both come from the same already projected report. What keeps the overwrite safe is a short-circuit two methods up, and that is worth a line of comment here. inline: EnsureRequirements.scala:830

[peter-toth]

Finding 9.unwrapGroupPartitions(right) at :618 and innermostGroupPartition(right) here run the same descent and land on the same node. That is what makes the pair sound. The keys come from that node's child, the positions from that node itself. The code derives each one separately and pairs them by hand, so the pairing holds because two call sites agree rather than because it is one value.

That is the shape findings 3 and 4 came out of. The plan moved to the pre-alignment node, the positions stayed on the spec, and nothing failed to compile. One value removes the class. At :617-618:

valleftGrouping= innermostGroupPartition(left)
valrightGrouping= innermostGroupPartition(right)
valunwrappedLeft= leftGrouping.map(_._1.child).getOrElse(left)
valunwrappedRight= rightGrouping.map(_._1.child).getOrElse(right)

and here:

val (partiallyClusteredChild, partiallyClusteredPositions) =if (replicateLeftSide) {
(unwrappedRight,
rightGrouping.flatMap(_._1.joinKeyPositions).orElse(rightSpec.joinKeyPositions))
} else {
(unwrappedLeft,
leftGrouping.flatMap(_._1.joinKeyPositions).orElse(leftSpec.joinKeyPositions))
}

The descent then runs twice per join instead of three times. unwrapGroupPartitions is left with its one remaining caller, the ShuffleExchangeExec site at :288.

[ulysses-you]

Done -- one descent per side: leftGrouping/rightGrouping are computed once at the top of the block, and the unwrapped plans, the statistics, the original keys and the positions all derive from them. unwrapGroupPartitions keeps its one remaining caller at the shuffle site. 94afe6e

[peter-toth]

Finding 10. This is about :644, which sits just above the hunk, so this is the nearest line I can anchor on.

The statistics read at :623-645 now comes from the pre-alignment plan. The branch it falls back to does not. leftPartKeys and rightPartKeys at :551-552 are leftSpec/rightSpec's keys, and those specs are built at :167 from the aligned node's report. Both sides align to the same mergedPartitionKeys with the same per-key counts, so on a re-run the two lists have equal length and leftPartKeys.size < rightPartKeys.size is false whatever pass 1 chose. That is the flip r3864640486 describes. The fix reached the trigger and left the flip.

I could not reach it on this head and the argument here is from reading. After the fix unwrapGroupPartitions lands on the join child pass 1 planned, which carries a logicalLink, so the statistics branch wins unless a side reports sizeInBytes <= 1. I instrumented this fallback and ran KeyGroupedPartitioningSuite plus EnsureRequirementsSuite. It did not fire once in 187 tests.

So this is about holding the invariant at both branches, not about a defect today. The pre-alignment counts are one collectFirst away, and :684 already does that extraction for one side:

// The pre-alignment split counts, for the same reason the statistics above read the// pre-alignment plan. On a re-run both aligned reports hold the same number of keys.defrawNumKeys(plan: SparkPlan, aligned: Int):Int=
plan.outputPartitioning match {
casee: Expression=> e
.collectFirst { casek: KeyedPartitioning=> k.numPartitions }
.getOrElse(aligned)
case _ => aligned
}
rawNumKeys(unwrappedLeft, leftPartKeys.size) <
rawNumKeys(unwrappedRight, rightPartKeys.size)

[ulysses-you]

Done -- the fallback compares the pre-alignment split counts, read through a new PartitioningCollection.numKeyedPartitions (the representative keyed member, collections included). Pinned by a unit test that forces the fallback -- dummy plans carry no logicalLink -- and reads the choice back off the distributePartitions flags; reverting the fix fails its first arm ((true, false) against the expected (false, true)). 94afe6e

[peter-toth]

Finding 11.joinKeyPositions one line up keeps what the reused node holds, because the incoming value was computed against the node's already projected report. reducers is computed the same way. leftSpec.reducersBothWays(rightSpec) at :564 runs over the reported expressions, which on a re-run are the reduced ones this node produced. It is written straight through.

Nothing goes wrong on this head, and I traced why. A re-run with reducers never reaches here. Both sides report the same reduced keys and the same expressions after pass 1, so isCompatible at :525-528 is true. v2BucketingPartiallyClusteredDistributionEnabled is false, because a reducer needs canReduceKeys and that excludes it. So the whole if at :529 is skipped.

I measured it rather than only tracing it. Printing g.reducers and the incoming value here and running KeyGroupedPartitioningSuite plus EnsureRequirementsSuite, it fires 86 times, and every one is existing=false, incoming=true, expectedPartitionKeys set=false. The node is always the one the children loop just created, never an aligned one from an earlier pass.

So this is a request for the reason in the source, not a fix. Next to a defended orElse the plain assignment reads like an oversight. And if that short-circuit ever weakens, the reused node silently loses its reducer, its keys go back to the raw space, and every expected key misses.

Suggested change

	reducers = reducers,
	// Unlike `joinKeyPositions`, these need no `orElse`. A re-run with reducers never reaches
	// here. Both sides then report the same reduced keys, so `isCompatible` above is true and
	// the whole block is skipped.
	reducers = reducers,

[ulysses-you]

Done -- comment added verbatim. 94afe6e

[ulysses-you]

Thanks @peter-toth -- 9, 10 and 11 are addressed in 94afe6e:

• 9 -- one descent per side: leftGrouping/rightGrouping are computed once, and the unwrapped plans, the statistics, the original keys and the positions all derive from them. unwrapGroupPartitions keeps its one remaining caller at the shuffle site.
• 10 -- the fallback compares the pre-alignment split counts, read through a new PartitioningCollection.numKeyedPartitions (the representative keyed member, collections included). A unit test pins it by forcing the fallback -- dummy plans carry no logicalLink -- and reads the choice back off the distributePartitions flags; reverting the fix fails its first arm ((true, false) against the expected (false, true)).
• 11 -- the comment is in verbatim.

Follow-up for finding 5 filed as SPARK-59234; finding 2 is left with you as discussed. Local suites green (204).

[peter-toth]

Re-checked through 94afe6e — findings 9, 10 and 11 resolved (one descent per side, the fallback on the pre-alignment counts, the reducers comment in verbatim), and finding 5 is filed as SPARK-59234. Nothing regressed. Finding 2 stays with me.

Measured on this head against the merge base 0875765:

• the 9 new tests all pass on the head;
• on the base, the 3 end-to-end tests plus the shuffle site strips every grouping the rule inserted and the replicate-side fallback counts pre-alignment partitions all fail, so every test that pins a fix is either master-level or revert-verified;
• a single-child operator over a partially clustered layout still gets grouped passes on the base, which is what its own comment says it does;
• the fallback's answer on a first pass differs between the two heads, on 2 of the 3 shapes I tried. That is finding 12.

Non-blocking

• 12.The pre-alignment split counts also change a first pass (new): the fallback now reads the same number on both passes, which is what finding 10 asked for, and it gets there by moving the first pass rather than by leaving it alone. A first pass compared the aligned report's distinct keys, so the replicate side flips wherever a side holds more than one split per key. Measured on both arms; I would keep this head's number and name the change in the comment and in the user-facing-change section. inline: EnsureRequirements.scala:651
• 13.The repro's ALTER TABLE paragraph is stale (late catch): "Why are the changes needed?" still opens the second split for id = 1 with ALTER TABLE sp1 ADD COLUMN extra STRING, calls that detail load-bearing, and explains it through the in-memory catalog's write schema. Since R2 the test uses the suite's numRowsPerSplit = 1 instead, per r3864640525. The merge script takes the commit message from this description, so the stale mechanism lands in git history.

Minor

• 14.numKeyedPartitions is public for one sql/core caller (new):private[sql] covers EnsureRequirements and keeps representativeOf's private[physical] from reaching public through a one-line wrapper. inline: partitioning.scala:1141

[peter-toth]

Finding 12. This implements finding 10 and the invariant it asked for holds now: the fallback reads one number per side, from the pre-alignment plan, on every pass. It reaches that by changing what a first pass compares, and that half is not stated anywhere.

Base's number is leftSpec/rightSpec's key count, i.e. the aligned report's. Finding 10 was about what that is on a re-run: equal on both sides, so the comparison decides nothing. On a first pass it is a third thing. The children loop has already wrapped a non-grouped child, so the report holds one key per distinct value and base compared distinct keys. This head compares raw splits, and the two differ for exactly the shape partial clustering is about, a side with more than one split per key.

Measured on a bare first pass -- two DummySparkPlans under a SortMergeJoinExec, no logicalLink, so the fallback fires -- reading (left.distributePartitions, right.distributePartitions) back off the planned children:

left	right	0875765	this head
3 splits, 1 distinct key	2 splits, 2 keys	(false, true)	(true, false)
4 splits, 2 distinct keys	3 splits, 3 keys	(false, true)	(true, false)
2 splits, 2 keys	3 splits, 3 keys	(false, true)	(false, true)

I am not asking for the old number back. Splits are the better proxy, and your new unit test's first arm is the case that shows it: one split against three, and the base picks the three-split side to replicate. The two coherent options are this head and reverting finding 10, and this head is the right one.

What is missing is one sentence in each of two places. The comment says "on a re-run both aligned reports hold the same number of keys, so comparing them decides nothing", which reads as if a first pass is untouched, and "Does this PR introduce any user-facing change?" lists only the duplicated rows. This is a wrong-results fix headed for branch-4.2 and branch-4.3, so a plan change riding along in a maintenance backport is worth naming:

// As a simple heuristic, we pick the side with fewer number of partitions to// apply the grouping & replication of partitions. The counts read the// pre-alignment plans, for the same reason the statistics do: on a re-run both// aligned reports hold the same number of keys, so comparing them decides nothing.// This also changes a first pass, which compared the aligned report's distinct// keys rather than the splits behind them.

[ulysses-you]

Done -- the comment names the first-pass change verbatim, and the test gained two bare first-pass arms: three splits under one distinct key against two splits under two keys, where the pre-alignment count replicates the side with fewer splits (revert-verified: the old distinct-key count picks the other side), plus a control where both counts agree. 77c17e9

[peter-toth]

Finding 14. This is public, and its one caller is EnsureRequirements, so private[sql] covers it. Worth narrowing because the method is a one-line wrapper over representativeOf, which is private[physical] deliberately: as written, the count of a keyed member becomes reachable from anywhere on the classpath in one hop.

private[sql] defnumKeyedPartitions(partitioning: Partitioning):Option[Int] =

Separately, the scaladoc understates what holds. require at :1049 already forces every member to agree on numPartitions, keyed or not, so the answer is partitioning.numPartitions whenever there is a keyed member at all -- the representative only decides whether there is one. Saying that is shorter than the agreement argument and it is the stronger fact.

[ulysses-you]

Done -- narrowed to private[sql], and the scaladoc now states the stronger fact: the collection's require already unifies the count across all members, so the representative only decides whether a keyed member exists. 77c17e9

[ulysses-you]

Thanks @peter-toth -- 12, 13 and 14 are addressed in 77c17e9:

• 12 -- the comment names the first-pass change verbatim, the user-facing section calls out the fallback's plan change, and the fallback test gained two bare first-pass arms: one where the pre-alignment split count and the old distinct-key count pick different sides (revert-verified), and one control where they agree.
• 13 -- the repro no longer uses ALTER TABLE; it matches the test's numRowsPerSplit = 1 shape, with the real-connector reading kept in one sentence.
• 14 -- narrowed to private[sql]; the scaladoc states the stronger fact (the collection's require already unifies the count across all members, so the representative only decides whether a keyed member exists).

Local suites green (204).

[cloud-fan]

Review summary

The idempotence repair is coherent across the rewrite, pre-alignment reads, raw-key positions, tag preservation, shuffle stripping, and fallback partition counts, and the new tests exercise the important regression shapes. The remaining findings are four localized comment corrections; no runtime correctness, API, serialization, concurrency, or test-coverage defect remains at the pinned head.

Findings

4 total: 0 P0, 0 P1, 0 P2, 4 P3.

Nit (P3)

• Correct the grouping pass ownership — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:775 — see inline.
• Add the missing verb — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:617 — see inline.
• Fix the comparative phrase — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:645 — see inline.
• Remove the dangling modifier — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:763 — see inline.

[cloud-fan]

Nit (P3): This comment reverses the ownership of the two layers. On a rerun, this invocation's distribution step has just added the outer grouping; the inner GroupPartitionsExec is the one inherited from the earlier pass and owns the alignment we need to preserve. Please describe the outer node as current-pass and the inner node as earlier-pass.

[ulysses-you]

Done -- the comment now attributes the outer grouping to this invocation's distribution step and the one below to the earlier pass.
3d07619

[cloud-fan]

Nit (P3): This sentence is missing a finite verb. Please change read once to are read once so the subject (The pre-alignment plan ... and the grouping ...) has a predicate.

[ulysses-you]

Done -- are read once.
3d07619

[cloud-fan]

Nit (P3):fewer number of partitions is ungrammatical. Please use fewer partitions (or a smaller number of partitions).

[ulysses-you]

Done -- fewer partitions.
3d07619

[cloud-fan]

Nit (P3): The opening modifier dangles because the following subject is the non-SortExec shapes. Please use a construction such as Instrumentation of the descent over KeyGroupedPartitioningSuite found these non-SortExec shapes ....

[ulysses-you]

Done -- reworded as suggested: "Instrumentation of the descent over KeyGroupedPartitioningSuite found these non-SortExec shapes hiding a node: ...".
3d07619