[SPARK-58968][SQL] Fix SPJ allowKeysSubsetOfPartitionKeys correctness for non-join operators

[peter-toth]

What changes were proposed in this pull request?

This builds on two changes that are now on master.

#58351 (SPARK-59057) renamed KeyedPartitioning.isNarrowed to isCollapsed and split groupedSatisfies into keysSatisfy and mayGroupToSatisfy. Two things here follow from that. The classification asks whether a member can satisfy the distribution once a GroupPartitionsExec is allowed, which is keysSatisfy for a grouped member and mayGroupToSatisfy for a non-grouped one, since only the second may coalesce duplicate keys and only it needs the collapse permission. That composition is a new KeyedPartitioning.keysMaySatisfy, which keeps keysSatisfy private. And the tests that build a KeyedPartitioning state isCollapsed explicitly, because that parameter lost its default.

#58420 (SPARK-59120) made every reader of a KeyedPartitioning's partition keys take its types from keyDataTypes, the schema the keys were written under, rather than from the partition expressions. This PR projects those keys on more paths, so it needs that fix underneath it. Without it, projecting the keys of a partitioning whose keys a reducer rewrote throws ClassCastException at planning.

This is an alternative to #58245, which fixes the same JIRA by adding the projection to one of the two branches below.

EnsureRequirements split a child's KeyedPartitionings by isGrouped and then had two branches that each had to insert a GroupPartitionsExec. This PR classifies by what still has to happen to the data instead.

• splitKeyedPartitionings now takes the required distribution and answers two questions, in this order. Whether a non-KeyedPartitioning member already satisfies it, and if not, how a KeyedPartitioning member can. As it is, or after a GroupPartitionsExec projecting to the partition expression positions returned with it, with None positions when the node only has to coalesce duplicate partition keys. At most one partitioning comes back, because the caller acts on a single one of them. The order is deliberate, since the keyed half is the one that projects partition keys, so it only runs when no plain member is enough.
• A new clusterKeyPositions helper derives those positions from the required clustering, and a new KeyedPartitioning.numPartitionsProjectedOn answers how many partitions a projection onto them would leave.
• The four-way match collapses to three cases, because the two arms that each had to insert a GroupPartitionsExec become one. It now matches on the distribution and the resolution together, since two of the three arms are decided by the resolution rather than by the distribution. The OrderedDistribution arm keeps its own insertion, as on master.

clusterKeyPositions keeps a partition expression when it is one of the operation keys. keysSatisfy recognises that at the reference level, where a bucket(4, a) transform covers the cluster key a, and also at the expression level, where the cluster key is the partition expression itself. Both are honoured here, which is why the positions are derived from the clustering rather than taken from KeyedShuffleSpec.keyPositions. That answers only the first, which is the right question for a storage-partitioned join but would drop a partition expression that is itself an operation key. The expression-level test is what keeps a position whose expression is clustered on while its references are not. No production path builds such a clustering today. V2ExpressionUtils.toCatalystTransformOpt maps IdentityTransform to the resolved attribute itself, so the reference-level test matches the same position anyway, and DistributionAndOrderingUtils.prepareQuery maps resolveTransformExpression over a write's clustering, so a TransformExpression does not survive into one. The test builds the shape by hand, with a table partitioned by (id, years(ts)) and clustered on those same two expressions. The check is here because keysSatisfy already accepts that shape, so deriving the positions any other way would make the two disagree.

A member whose expressions cover no operation key at all is skipped rather than projected to no position, which would put every partition into one. Only a partitioning whose expressions have no references can get there. A DSv2 scan cannot report one, because supportsExpressions refuses it, but nothing rejects it at KeyedPartitioning construction.

A partitioning that satisfies the distribution is kept as it is when nothing is left for a node to do, which holds in two ways. Either the projection drops no position, so satisfies is not the over-claim this fix is about. Or a position is dropped but the projection merges nothing, so every operation key already lives on a single partition. Keeping the partitioning is then better than projecting, because KeyedPartitioning([id, name]) and KeyedPartitioning([id]) describe the same number of partitions and only the first lets a downstream operator co-partition on name too. This is where the collapse gives the single insertion point a new way to be wrong, namely inserting a node that merges nothing, which cannot happen on master, where a grouped partitioning got no node at all.

That question is asked of every member of the child's partitioning, not just of the one a node would be built from, because a child can report a whole PartitioningCollection and its members can disagree about which positions are operation keys. An inner join is where they do. Its outputPartitioning is the two sides' partitionings, and unlike AliasAwareOutputExpression it does not enumerate the mixed combinations, so a window keyed on one side's first key column and the other side's remaining ones sees one member covering position 0 and one covering the rest. The first needs no node, the second would coalesce for nothing.

Among the members that do need a node, the one whose projection leaves the most partitions is used, and that answer is exact rather than a heuristic. Only the position sets not contained in another one have to be projected. Projecting to fewer positions can merge partitions but never split them, so a contained set can never leave more partitions than the set containing it, and can only tie. Dropping it therefore costs no parallelism, and on a tie it settles the choice toward the wider set, which still names the keys the narrower one would have dropped. In the ordinary case one set contains all the others and a single projection settles it, and only sets that genuinely disagree, neither containing the other, each cost one. That matters because the projection is the expensive step here, since KeyedPartitioning.projectKeys allocates a row per input partition and InternalRowComparableWrapper.hashCode is uncached. So the count is memoized per position set as well, a set covering every position needs no projection at all, being the identity on the key values, and a single surviving candidate with no required count needs none either, which is the shape the default config produces. Keeping one member per position set costs nothing either, because PartitioningCollection guarantees its members share the partitionKeys reference and their arity, so position i addresses the same key column in all of them.

Distribution.requiredNumPartitions needs care, because a GroupPartitionsExec derives its number of partitions from the partition keys it is handed rather than from the operator. The requirement is therefore checked against the count such a node would produce, not against the count the partitioning happens to have now, and it filters the candidates rather than vetoing the winner. One that would land on the required number must not lose the ranking to one that cannot honour it and send the whole child to a shuffle instead. As-is satisfaction goes through satisfies, which enforces the count on its own.

For an operator that co-partitions more than one child no projection is done here, because that belongs to the multi-child block below, which a storage-partitioned join reaches through checkKeyGroupCompatible and anything else through withJoinKeyPositions. Projecting inline as well would leave that block deriving positions from an already projected partitioning and applying them to the unprojected keys.

KeyedPartitioning.satisfies is not touched, so nothing outside EnsureRequirements changes behaviour. It does still answer true for a partitioning that needs a projection first, which means ValidateRequirements cannot catch a missing GroupPartitionsExec. Giving that check the strict test directly, without changing what satisfies answers, is a follow-up we are working on.

The OrderedDistribution arm also loses a MatchError. It tested that the partition keys are sorted with partitionKeys.sliding(2), which yields one short window for a single-key partitioning, and case Seq(k1, k2) cannot match it, so planning threw. A v2 table whose rows all share one partition value reports such a partitioning, since DataSourceV2ScanExecBase has no single-partition short-circuit, and with v2BucketingAllowSorting on, a global sort on the partition key reaches this branch. The test is now a zip of the keys with their successors, which is vacuously true for one key. Pre-existing, but the branch is rewritten here anyway.

Two scaladocs are corrected as well. KeyedPartitioning taught isGrouped as the axis this PR replaces. GroupPartitionsExec.joinKeyPositions described its projection as being "for join compatibility", and it now carries the projection for a single-child operator too.

Why are the changes needed?

With v2BucketingAllowKeysSubsetOfPartitionKeys enabled, KeyedPartitioning.keysSatisfy only requires that some operation key overlaps the partition attributes. A partitioning grouped on (id, name) therefore reports that it satisfies ClusteredDistribution([id]) while rows sharing an id still sit on separate partitions. A storage-partitioned join projects the keys down to the operation keys in checkKeyGroupCompatible, and for a non-join operator nothing did.

isGrouped is the wrong thing to classify on, because it only says the full partition keys are unique and says nothing about whether the projected keys are. A GroupPartitionsExec is needed either to coalesce duplicate partition keys, or to project down to the operation keys, or both, and splitting by provenance put those two reasons in different branches. Both branches returned wrong results.

1. A grouped partitioning over a superset of the operation keys got no node at all, so a top-k window over PARTITION BY id on an (id, name)-partitioned table surfaced id=1 twice, once per (1,'aa') and (1,'bb') partition.
2. A non-grouped partitioning over a superset got a node without a projection, so it coalesced by the full partition keys and left the operation key split. Any source reporting more than one split per partition value hits this. SUM(price) OVER (PARTITION BY id) over the same table with two splits for (1,'aa') returned 25.0 and 20.0 instead of 45.0.

Both reach back to 4.2.0, where GroupPartitionsExec and this classification were introduced.

A reduced storage-partitioned join reaches the same wrong result through a third shape. Joining an identity(ts)-partitioned table to a years(ts)-partitioned one under v2BucketingAllowCompatibleTransforms leaves both sides grouped on (year, bucket), so two rows sharing a ts in different buckets sit on separate partitions. SUM(v) OVER (PARTITION BY ts) then returns 10 and 20 on master where the answer is 30 and 30, and the projection this PR inserts is what merges the two partitions. Before #58420 that same query threw ClassCastException at planning once the keys were read, so the wrong answer only became observable when that fix landed.

Does this PR introduce any user-facing change?

Yes, it fixes a data correctness issue. With v2BucketingAllowKeysSubsetOfPartitionKeys enabled, a single-child operator whose keys are a subset of the partition keys now produces correct results, namely a window PARTITION BY and the single-pass aggregate shapes (FlatMapGroupsInBatchExec, ArrowAggregatePythonExec, MapGroupsExec). A two-phase SQL aggregate was already correct, because its partial HashAggregate is a PartitioningPreservingUnaryExecNode, so it narrows KP([id, name]) to KP([id]) before the final aggregate sees it. A cogroup is unaffected either way, because the projection there is left to the multi-child block, exactly as before.

Two further changes are not gated on that config.

The requiredNumPartitions rule applies whether allowKeysSubsetOfPartitionKeys is on or off, since a scan still needs spark.sql.sources.v2.bucketing.enabled to report a KeyedPartitioning at all, and master has no count check on the grouping path. A non-grouped KeyedPartitioning with 3 partitions and 2 distinct keys under ClusteredDistribution([k], requiredNumPartitions = Some(3)) got a GroupPartitionsExec with 2 partitions on master and now gets a shuffle with 3. I could not find a query where such a distribution meets a KeyedPartitioning today, so this is robustness rather than a reachable wrong result. Only StatefulOperatorPartitioning and AQEUtils ever set the requirement. The AQEUtils one fires only over a HashPartitioning child. StatefulOperatorPartitioning sets it through a plain ClusteredDistribution when spark.sql.streaming.statefulOperator.useStrictDistribution is off, and a streaming scan never reports a KeyedPartitioning, because MicroBatchScanExec, ContinuousScanExec and RealTimeStreamScanExec all leave keyGroupedPartitioning at None. The initial-state child of flatMapGroupsWithState and transformWithState is a batch plan though, so it can be a scan that does. That operator co-partitions its two children, so the shape is the multi-child block's to decide either way, and what changes here is only that a candidate which cannot honour the count no longer wins the resolution. The requirement is part of the ClusteredDistribution contract though, the two could meet more widely as this area grows, and honouring it costs nothing where nobody asks for a count, so it seemed the wrong thing to keep ignoring.

The MatchError fix is gated on v2BucketingAllowSorting instead, which is also off by default.

explain gains one label where the new projection happens. A GroupPartitionsExec inserted for a single-child operator now carries joinKeyPositions, so the node prints JoinKeyPositions: [...] where it printed nothing before. The name is historical, and the scaladoc now says the projection is not join-specific. This needs allowKeysSubsetOfPartitionKeys on, so no golden file moves.

How was this patch tested?

Added regression tests in KeyGroupedPartitioningSuite:

• window top-k over PARTITION BY a subset of the partition keys, for both PARTITION BY id and the duplicated PARTITION BY id, id
• window top-k over union output partitioning
• a plain window over a subset of the partition keys on a non-grouped KeyedPartitioning, asserting the inserted node projects to the operation key rather than only coalescing
• no GroupPartitionsExec and no shuffle when projecting to the operation keys merges nothing
• a window over an inner join's two-member PartitioningCollection, keyed on the left side's first key column and the right side's remaining two, asserting no node is inserted because the left member needs none, even though the right one covers more operation keys
• a window over a join that reduced one side's keys onto the other side's key space, asserting the two rows sharing a ts end up on one partition

and in EnsureRequirementsSuite:

• a FlatMapCoGroupsInPandasExec over (n, i)-partitioned children grouped on i, asserting both sides are grouped on i and not on n
• a grouped KeyedPartitioning whose count differs from requiredNumPartitions, asserting the count is still honoured with a shuffle
• an (n, i)-partitioned KeyedPartitioning whose count matches requiredNumPartitions but which needs a projection, asserting it falls back to a shuffle rather than to a node that would break the count it just passed
• a projecting KeyedPartitioning whose post-projection count matches requiredNumPartitions, asserting it still groups on the operation key with no shuffle
• a non-grouped KeyedPartitioning whose post-grouping count matches requiredNumPartitions, asserting it still groups without a shuffle
• the same count rule with allowKeysSubsetOfPartitionKeys left at its default, asserting the shuffle
• an (id, years(ts))-partitioned KeyedPartitioning clustered on those same two expressions, with and without requireAllClusterKeys, asserting no node is inserted when a cluster key is the partition expression itself
• a (bucket(4, a), b)-partitioned KeyedPartitioning clustered on a alone, asserting the node projects to the bucket position. This is the one test where the reference-level match decides a kept position, and it is the shape the single-reference soundness argument is about. master inserts no node at all here, so the partitions sharing a bucket stay apart.
• a single-partition KeyedPartitioning under OrderedDistribution, asserting planning no longer throws a MatchError
• a KeyedPartitioning with no partition expressions, asserting it is still kept as it is, since without that guard it reaches the shuffle branch, where UnspecifiedDistribution.createPartitioning throws
• a KeyedPartitioning whose expressions reference no column, asserting the child is shuffled rather than projected to no position at all
• a PartitioningCollection whose two members cover a different number of operation keys, asserting the wider one supplies the projection
• two members whose position sets are nested and whose projections leave the same number of partitions, asserting the containing set still wins, so the projection keeps naming the key the other would have dropped
• two members whose position sets are nested, where only the narrower one's projection lands on requiredNumPartitions, asserting it is used rather than pruned and lost to the wider one
• two members whose position sets are not nested and where the narrower one leaves more partitions, asserting the narrower projection wins over the wider coverage
• two members covering one position each whose projections leave different numbers of partitions, asserting the one leaving the most supplies the projection
• two members covering one position each whose projections leave the same number of partitions, asserting the one the child reports first wins, in both collection orders

and in ProjectedOrderingAndPartitioningSuite a grouped and collapsed KeyedPartitioning, asserting keysMaySatisfy accepts it where mayGroupToSatisfy refuses it, because a grouped partitioning has no duplicate keys left for a node to coalesce.

Seventeen of the twenty-four fail without the production change in this commit, measured on master. Those are the four window tests, the reference-free expressions, the transform position, the MatchError, four of the six requiredNumPartitions tests, all five multi-member ones and the keysMaySatisfy one, which does not even compile there. The other seven pass already and guard this PR's own insertion decision. Removing the condition each one covers makes it fail, with one exception worth naming. the candidate covering the most operation keys wins survives the removal of either the containment prune or the ranking, because either mechanism alone picks the same winner, so it earns its place by documenting the shape rather than by pinning a single condition. The join test is one of the seven. The base leaves that plan alone too, and it is here because an earlier revision of this PR inserted a node there. Removing the isCoPartitioned guard also fails two pre-existing SPJ tests.

DistributionSuite and ShuffleSpecSuite pass with 29 tests, and KeyGroupedPartitioningSuite, EnsureRequirementsSuite, ProjectedOrderingAndPartitioningSuite, GroupPartitionsExecSuite and PlannerSuite pass with 295.

The window PARTITION BY tests come from #58245.

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

Generated-by: Claude Code

Co-authored-by: Xiduo You ulyssesyou@apache.org

[ulysses-you]

I'm fine with this fix, the main change is same that adding an extra GroupPartitionExec with projected key position for non-join operators. This fix makes change inlines the method splitKeyedPartitionings and return the needsGrouping which including the key positions.

[dongjoon-hyun]

I ran a deep review of this PR (line-by-line scan, removed-behavior audit, cross-file tracing, plus reuse/simplification/efficiency angles, with each candidate adversarially verified against the checked-out tree). Overall the change looks solid: I verified that the old four-way classification maps faithfully onto the new two-bucket one (nonGroupedSatisfies, AllTuples, UnspecifiedDistribution, OrderedDistribution, and requiredNumPartitions paths all produce identical outcomes), that the inline and multi-child projection paths are mutually exclusive (no double projection), that re-applying the rule over an inserted GroupPartitionsExec is idempotent under AQE, and that using needsGrouping.head is safe given PartitioningCollection's shared-partitionKeys invariant. The findings below are ordered by severity; the first three are correctness/design, the rest are cleanup.

1. Pre-existing crash carried into the rewritten function: sliding(2) MatchError for a single-partition KeyedPartitioning under OrderedDistribution (confirmed)

spark/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala

Lines 100 to 102 in 8b74639

	if (satisfyingKeyedPartitioning.partitionKeys.sliding(2).forall {
	caseSeq(k1, k2) => keyOrdering.lteq(k1, k2)
	}) {

sliding(2) on a 1-element partitionKeys yields a single window of size 1, which case Seq(k1, k2) cannot match, so planning throws scala.MatchError. Reachable: a v2 table partitioned by identity(id) whose rows all share one partition value reports a one-key KeyedPartitioning; with spark.sql.sources.v2.bucketing.sorting.enabled=true, SELECT id FROM t ORDER BY id (global SortExec requires OrderedDistribution, admitted via areAllClusterKeysMatched) reaches this branch. Not a regression of this PR — the line is carried over verbatim — but since the branch is being rewritten anyway, case Seq(k1, k2, _*) (or a size guard) plus a small test would close it here. No existing sorting test covers a single-partition table.

2. The co-partitioned PartitioningCollection gap is documented but remains a live wrong-results hole (plausible, pre-existing)

The new comment at L74-L76 accurately states that a co-partitioned child reporting a PartitioningCollection gets a projection from neither the inline path (suppressed by isCoPartitioned) nor the multi-child block (withJoinKeyPositions never matches a ShuffleSpecCollection best spec). The consequence is the same wrong-results class this PR fixes, just behind a rarer plan shape: e.g. FlatMapCoGroupsInPandasExec on key i over children reporting collections of KPs grouped on (n, i) — compatibility is judged on the projected specs while the children run unprojected. Since it's the identical bug class, it may deserve a JIRA now rather than an open-ended follow-up note.

3. KeyedPartitioning.satisfies still over-claims, so ValidateRequirements certifies the exact plans this PR deems wrong (plausible, deferred by design)

The compensation (positions.isEmpty) lives only inside EnsureRequirements; ValidateRequirements.validateInternal still goes through satisfies, and it is the safety gate in AdaptiveSparkPlanExec.optimizeQueryStage and OptimizeSkewedJoin. The pre-fix wrong plan (grouped KP([id, name]) under ClusteredDistribution([id]), no GroupPartitionsExec) passes validation today, which also means no validator-based regression test can catch a reintroduction of this bug class. The PR description already flags this as a follow-up — agreed, but I'd track it with a JIRA for the same reason as (2).

4. The projected distinct count is computed twice per candidate partitioning (confirmed)

spark/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala

Lines 908 to 913 in 8b74639

	privatedefprojectionCoalesces(kp: KeyedPartitioning, positions: Seq[Int]):Boolean= {
	// The number of partitions a non-projecting GroupPartitionsExec would produce.
	valgroupedNumPartitions=
	if (kp.isGrouped) kp.numPartitions else kp.partitionKeys.distinct.size
	kp.projectKeys(positions)._2.distinct.size < groupedNumPartitions
	}

projectionCoalesces computes kp.projectKeys(positions)._2.distinct.size, and groupedNumPartitions (L954-L955) recomputes the identical expression for the requiredNumPartitions check; the no-projection fallback there also misses the isGrouped shortcut projectionCoalesces has. InternalRowComparableWrapper.hashCode is uncached (interpreted Murmur3 per call), so with many partition keys this is two full projection + distinct passes where one suffices — and the two formulas can drift, desynchronizing the (load-bearing, tested) requiredNumPartitions gate from the count the inserted node actually produces. Having projectionPositions return the positions together with the resulting partition count would collapse both call sites.

5. Per-collection-member repetition of O(n) key passes (confirmed)

split runs projectionPositions/projectionCoalesces for every PartitioningCollection member, although the collection invariant guarantees all KP members share the same partitionKeys reference — an m-member collection does m sets of identical passes, and the cost also lands on the common no-op path (a satisfying partitioning forces positions merely to evaluate positions.isEmpty, where the pre-PR code traversed keys zero times). This runs per child per rule invocation, again per AQE stage. Memoizing the unprojected distinct count per partitionKeys reference (and the projected count per positions value) would remove the repetition.

6. Admission and projection are decided by different code — drift risk (design)

splitKeyedPartitionings admits a KP into needsGrouping via KeyedPartitioning.groupedSatisfies, but what to project is computed by the rule-local projectionPositions — now a third independent implementation of cluster-key coverage alongside groupedSatisfies and KeyedShuffleSpec.keyPositions. The two must agree per position, or the inserted GroupPartitionsExec produces a partitioning that doesn't satisfy the distribution it was inserted for (the PR's own years(ts) unit test demonstrates the failure mode for a naive derivation). If groupedSatisfies is ever widened (e.g. GetStructField references, compatible-transform matching) without mirroring projectionPositions, the result is the L898 assert firing at planning time or a silent mis-projection. Housing the position computation next to groupedSatisfies on KeyedPartitioning — one matcher feeding both answers — would make the drift impossible, and would also give the currently-defensive expression-level disjunct a non-speculative home. Fine as a follow-up given the fix is config-gated.

7. The duplicated-PARTITION BY test duplicates the whole fixture (minor)

The "window top-k over duplicated PARTITION BY key" test is a full copy of the preceding subset test with only PARTITION BY id → PARTITION BY id, id, and its comment concedes it cannot fail independently. Looping the window spec (Seq("id", "id, id")) inside one test keeps the pinned behavior at half the fixture and runtime cost.

Refuted candidates, for the record: the assert(positions.nonEmpty) at L898 is provably unreachable for any KP satisfying supportsExpressions (every production producer preserves it); the OrderedDistribution comparator cannot bind a transform output to the wrong domain (areAllClusterKeysMatched requires 1:1 semanticEquals between partition expressions and SortOrder children); and GroupPartitionsExec.outputPartitioning's collection rebuild cannot trip PartitioningCollection's equal-count require in stock Spark (no producer emits a mixed KP + non-KP collection).

[peter-toth]

Thanks @dongjoon-hyun for the review, I've iterated on the code multiple times today, but this is still a draft. Will try to wrap it up tomorrow.

[dongjoon-hyun]

Got it~ I'll revisit when the PR becomes out of Draft status, @peter-toth .

[peter-toth]

This is more complex than I initially thought. Let me iterate on it a bit more.

[peter-toth]

@ulysses-you thanks for closing #58245 in favour of this one, and for the three window tests, which I brought over verbatim.

Since the two PRs look quite different, here is why I ended up somewhere else rather than iterating on yours.

The root cause is one level up from the branch you patched.splitKeyedPartitionings split the child's KeyedPartitionings by isGrouped, and isGrouped is not the question that matters here. It only says whether the full partition keys are unique. It says nothing about whether the projected keys are, and under allowKeysSubsetOfPartitionKeys those are different things. A GroupPartitionsExec is needed either to coalesce duplicate partition keys, or to project down to the operation keys, or both, and splitting by isGrouped puts those two reasons in different branches. So patching one branch leaves the other one wrong.

And the sibling branch was wrong in the same way. For a ClusteredDistribution a non-grouped KeyedPartitioning always lands in case _ => GroupPartitionsExec(child), with no positions, because nonGroupedSatisfies is Partitioning.satisfies0 and answers false there. I measured it on your head: an (id, name)-partitioned table with two splits for (1,'aa') and SUM(price) OVER (PARTITION BY id) returned 25.0 and 20.0 instead of 45.0. Worth knowing why your tests could not catch it: WindowGroupLimit Final also requires a ClusteredDistribution, so the projecting node lands above it and the wrong grouping underneath is harmless, because a group limit can only keep too many rows. It takes a plain window with no group limit above it.

isJoin also turned out to be the wrong predicate, and it regressed a cogroup. The multi-child block owns the projection for any operator with more than one ClusteredDistribution child, FlatMapCoGroupsInPandasExec included, not only for a ShuffledJoin. With KP([n, i]) and a cogroup on i, master grouped on i (2 groups) and the PR grouped on n (3 groups). The predicate that works is requiredChildDistributions.count(_.isInstanceOf[ClusteredDistribution]) > 1, which is what this PR uses.

On the positions themselves. Taking them from createShuffleSpec answers the reference-level question, which is exactly right for a storage-partitioned join: a bucket(4, a) transform covers the cluster key a. But a cluster key can also be the partition expression itself, and then the reference-level lookup drops that position and coalesces for nothing. That is why this PR derives the positions from the required clustering instead, honouring both forms.

One more difference worth naming, and it is plan quality rather than correctness. Your version decided from a single member, the one groupedSatisfies picked out of the child's partitioning. This one asks every member, because a PartitioningCollection's members can disagree about which of their positions are operation keys, and an inner join is where they do: ShuffledJoin.outputPartitioning unions the two sides without enumerating the mixed combinations, so a window keyed on one side's first key column and the other side's remaining ones sees one member covering position 0 and one covering the rest. Taking the wrong one is still correct - any admitted member's projection satisfies the distribution - but it can insert a node where another member needs none, which costs a CoalescedRDD layer and some parallelism. That case is measured in the suite.

The extra look is cheaper than it sounds. A position set contained in another one is dropped without projecting it, since projecting to fewer positions can merge partitions but never split them. In the ordinary case one set contains the rest and it is a single projection, same as yours; it is only more when the sets genuinely disagree.

[peter-toth]

@ulysses-you, @dongjoon-hyun updated and out of draft, so this is ready for another look.

@dongjoon-hyun thanks for the review, all seven are addressed. The MatchError you found is fixed with a test that throws without it. The projection is now computed once per distinct position set rather than once per PartitioningCollection member, and the duplicated PARTITION BY test shares a fixture with the subset test above it. Your point about admission and projection being two implementations of the same question is right, and it is the next thing we want to take on in this area - it changes what satisfies answers for every consumer, so it did not belong in a bugfix.

Since the earlier reviews the branch has also been rebased onto master twice, past SPARK-59025 and SPARK-59027, and it picked up a few things from my own re-reads: the ranking of candidate projections is now exact rather than a coverage heuristic, the "needs no node" question is asked of every member of the child's partitioning rather than only of the winner, and requiredNumPartitions filters the candidates rather than vetoing the winner. Nineteen tests, thirteen of which fail on master; the description has the breakdown.

[peter-toth]

@ulysses-you@dongjoon-hyun so you both know what I am planning in this area, here is what I have collected while working on this. Each will get its own JIRA as I get to it.

1. Reduced SPJ partition keys: the reported expressions still do not describe them.[SPARK-59045][SQL] Fix SPJ ClassCastException when reducer changes partition key data type #58335 fixes the identity-versus-transform shape. Two remain. When both sides are transforms and only one reduces, an exact expression does exist - the other side's transform retargeted at this side's child - but the un-reduced one is reported. When both sides reduce, the keys land in a space no transform describes, and the ClassCastException survives; I measured that on [SPARK-59045][SQL] Fix SPJ ClassCastException when reducer changes partition key data type #58335's head with the existing SPARK-56164 test plus allowKeysSubsetOfPartitionKeys. The stale expression is also a silent wrong-results bug at a chained join, and that one reaches released branches. This is the follow-up I offered on [SPARK-59045][SQL] Fix SPJ ClassCastException when reducer changes partition key data type #58335, and I am happy either way: I can take it, or carry it after whatever you do there.

   Update: shape 2 has since been handled in [SPARK-59045][SQL] Fix SPJ ClassCastException when reducer changes partition key data type #58335 itself, so only the both-sides-reduce shape is left, and I am taking that one.

2. The multi-child block pushes the wrong join key positions.withJoinKeyPositions is handed the best spec's positions for every child rather than each child's own specs(idx), and KeyedShuffleSpec.isCompatibleWith compares the projected partitionings without comparing the positions. So two sides can be compatible while their positions differ, and one side ends up grouped on the wrong column. Pre-existing, silent wrong results. SPARK-59025 makes it reachable in one more shape, because the head of an unwrapped ShuffleSpecCollection now supplies the positions for a child whose own spec may carry different ones.

   Update: I went looking for a query that reaches this and could not build one, so it is latent rather than live. A cogroup is the only non-join operator with two clustered children, and its grouping key comes from AppendColumns, so no KeyedPartitioning satisfies it and both sides shuffle before the block decides anything. With a join, checkKeyGroupCompatible handles the two-keyed-sides case itself and pushes each side's own positions - measured with tables on (dept, id) and (id, dept) joined on id: [1] and [0], no shuffle, right answers. When that path bails, at least one side has no keyed spec, so the block never sees two keyed children whose positions differ. The wrong-column grouping is still real in the code path, and the one-line fix stands, but I am treating it as hardening and taking item 3 first.

3. ShuffleSpecCollection members can disagree on numPartitions. Under allowKeysSubsetOfPartitionKeys, KeyedPartitioning.createShuffleSpec projects each member of a PartitioningCollection to its own join-key subset, so the resulting specs can differ in numPartitions. ShuffleSpecCollection.numPartitions then reports the head's count while createPartitioning requires them all equal and throws. Reproduces both before and after SPARK-59025, so it is independent of that change.

4. A coalesce keeps a stale non-KeyedPartitioning sibling.GroupPartitionsExec.outputPartitioning rewrites only the KeyedPartitioning members of its child's partitioning and passes any other member through with its old numPartitions, so rebuilding the enclosing PartitioningCollection trips its equal-count requirement. No producer builds such a collection from a query today, so this is reachable only from a hand-built partitioning. The repair belongs in the node: after coalescing, an inherited HashPartitioning is not mis-counted, it is false, and should be dropped.

5. Make the validation check strict.satisfies answers true for a partitioning that still needs a projection, so ValidateRequirements cannot catch a missing GroupPartitionsExec. Giving that check the strict test, and housing admission and projection next to each other so the two cannot drift apart, is the shape I want. This is what @dongjoon-hyun asked about in his review here.

[ulysses-you]

Empirical verification performed

Base = upstream/master HEAD e2cc1d64ab36 (identical to this PR's base). Applied the full patch transiently to compile and run, then reverted:

Run	Tree	Result
Full patch applied	compile clean (30s, warm cache)	EnsureRequirementsSuite -z SPARK-58968: 14/14 pass; KeyGroupedPartitioningSuite -z SPARK-58968: 5/5 pass
Negative control (test files applied, main code reverted)	same suites	exactly 6 pass / 13 fail on master, matching the description's enumeration precisely (the 6 passing are exactly the insertion-decision guards claimed, incl. the join-collection-member-needs-none test)

Line-length and non-ASCII checks on the changed files pass (splitKeyedPartitionings has only one caller, private to EnsureRequirements, nothing serializes the new planner-only structures).

I also traced every correctness claim through the peer code rather than trusting comments. The load-bearing facts verified against master sources:

• Candidates-dedup / exec-side consistency.candidates is keyed only on BitSet ("first member wins"), while runtime projection uses GroupPartitionsExec.groupedPartitionsTuple's collectFirst { case k: KeyedPartitioning => k } (GroupPartitionsExec.scala:136). These agree only becausePartitioningCollection.checkKeyedPartitioningInvariant (partitioning.scala:848) requires all members to share the partitionKeys reference and arity, and fromPartitionings interns by value-equality (wrapper equality includes dataTypes, InternalRowComparableWrapper.scala:65). Given that invariant, projected key values read at positions i..j are byte-identical whichever member applies them, so dedup-by-position-set, the (BitSet, Seq[DataType]) memo key on numPartitionsAfter, and the "most partitions" ranking are all internally exact — count equality ⟺ no merges, since a projection can coalesce but never split.
• Over-coalescing is safe, under-coalescing is the bug. For single-child operators (window top-k, batch aggregates), merging two tasks together never splits an operation key into separate groups because window/aggregation grouping happens per-row within the task; only leaving rows sharing an op key in different tasks is wrong. So choosing the candidate leaving the most partitions, tie-breaking toward the containing set, is not merely plan-quality — dropping it would be incorrect. This is a real invariant but it is never stated anywhere.
• Assert reachability.clusterKeyPositions's assert(positions.nonEmpty ...) was probed against every satisfaction path: nonGroupedSatisfies is true only via base satisfies0 (Unspecified/Broadcast → returns all indices anyway); groupedSatisfies under requireAllClusterKeys goes pairwise positionally, under subset-config requires overlap, default branch requires every attribute ∈ clustering — each positionally guarantees ≥1 covered index. Unreachable through current branches.
• AQE re-entry fixed point. After insertion, re-running splitKeyedPartitionings on the projected grouped KP satisfies again with positions.size == len → satisfiedAsIs → stable; no double insertion under AQE re-optimization.
• @transientjoinKeyPositions. New single-child usage leans on it for correctness, unlike pre-existing SPJ-only uses — but grouping happens entirely driver-side at RDD construction (doExecute → CoalescedRDD), so executor-side nulling is irrelevant. No new hazard.
• Copartitioned equivalence. With isCoPartitioned=trueclusterKeyPositions returns all indices, so Option.when(size < len) yields None = master's bare GroupPartitionsExec(child); childrenIndexes relocation does not change preferSinglePartition semantics (still evaluated on mapped children after the map). checkKeyGroupCompatible/withJoinKeyPositions operate independently of this rewrite.
• MatchError fix. Confirmed real: sliding(2) on a 1-element seq yields one size-1 window and case Seq(k1, k2) threw; the added case _ => true is correct (single partition trivially sorted).

Findings (non-blocking)

N1. First-member-wins correctness rests on an invariant enforced three classes away

EnsureRequirements.scala (splitKeyedPartitionings) / GroupPartitionsExec.scala:63

• Defect statement: the planner picks positions per BitSet assuming "the same set projects to the same keys whichever member applies it", but the executor re-derives the member independently via collectFirst; these coincide only if PartitioningCollection keeps members interned on identical partitionKeys — a requirement living in partitioning.scala:848 with a value-equality escape hatch in fromPartitionings.
• Failure scenario if relaxed: collection [P1(non-satisfying), P2(satisfying)], P2 selected at {1,2} → exec projects P1's expressions at {1,2}; projected values stay correct (shared keys) but output metadata names P1's columns as group keys, poisoning downstream spec creation/ordering claims. Today impossible; one future relaxation of checkKeyedPartitioningInvariant makes this silent op-key mislabeling.
• Peer citation: the invariant itself (partitioning.scala:857-862) predates this PR; this PR multiplies the number of consumers relying on it without adding a probe where the coupling crosses files.
• Verdict: CONFIRMED-sound-today (traced end-to-end); report as fragility — suggest one sentence in clusterKeyPositions' scaladoc pointing at PartitioningCollection.fromPartitionings as the actual guarantee, or a cheap debug require(k.partitionKeys eq firstMember.partitionKeys).

Review verdict overall: no blocking correctness issues found; both behavior changes outside the gated config are either strictly safer than master or unreachable today, and every axis probed (codegen/interpreted, WSCG, AQE re-entry idempotency, ordering metadata, collation-aware key hashing) held up under tracing plus the test runs above.

[ulysses-you]

Non-blocking (N1 of my review below): first-member-wins here only coincides with what actually happens at runtime.

[ulysses-you]

Nit/non-blocking: "first member wins" is sound only because runtime projection re-derives the member independently via GroupPartitionsExec.groupedPartitionsTuple's collectFirst { case k: KeyedPartitioning => k } (GroupPartitionsExec.scala:135), which is guaranteed to agree with whichever member was recorded here only by PartitioningCollection.checkKeyedPartitioningInvariant (partitioning.scala:848, value-equality interning in fromPartitionings) -- an invariant enforced three classes away from this call site.

If that invariant is ever relaxed, e.g. collection [P1(non-satisfying), P2(satisfying)] with P2 selected at {1,2}: the node projects P2's positions, but the executor projects P1's expressions at them; projected key values stay correct (shared keys) while output metadata names P1's columns as group keys -- silent op-key mislabeling poisoning downstream spec creation/ordering claims.

Two cheap hardenings: point this comment at fromPartitionings as the actual guarantee (the coupling crosses files today), or add a debug check like require(recorded.partitionKeys eq firstCollected.partitionKeys).

[peter-toth]

Agreed, and thanks for tracing it across the files. I have added a sentence at the candidates declaration naming checkKeyedPartitioningInvariant and the interning in fromPartitionings as the actual guarantee, and saying that relaxing it means changing GroupPartitionsExec's collectFirst at the same time rather than this side alone.

Worth adding that the assumption is already only partly guaranteed: the partitionKeys reference and the arity are enforced, the per-position expressionDataTypes are not. Two members of one collection really can declare different types over the same keys - pushPartValues plus allowCompatibleTransforms, an identity(ts)-partitioned table joined to a years(ts)-partitioned one - which is why the projected-count memo in this method is keyed on (BitSet, Seq[DataType]) rather than on the position set alone.

I would rather not add a require on the reference identity here: PartitioningCollection already asserts it on construction, so a second check would go stale the day that one moves.

[peter-toth]

One thing in the verification section I would push back on, because it is the kind of claim that hardens into a constraint if it goes unanswered.

Over-coalescing is safe, under-coalescing is the bug. [...] So choosing the candidate leaving the most partitions, tie-breaking toward the containing set, is not merely plan-quality — dropping it would be incorrect.

The premise is right and worth having written down. The conclusion does not follow, and it points the other way: if coalescing more is the safe direction, then taking the candidate that leaves the most partitions is the less conservative choice, not the one whose removal would break correctness.

Every admitted candidate is correct on its own, whichever the ranking picks. groupedSatisfies' subset branch also requires expressions.forall(_.references.size == 1), so each position it keeps is a function of a single cluster key. Rows that agree on all the cluster keys therefore agree on that one, and land in one partition. Covering a single cluster key is sufficient - which is exactly why the config is allowed to admit a subset in the first place.

Measured rather than argued: in the no GroupPartitionsExec when a join collection member needs none test, an earlier revision of this PR picked the other candidate and inserted a node. checkAnswer passed - the rows were right - and only the plan was worse, 3 partitions where 4 were available. That is the whole cost of picking wrong here.

So I would keep the description's "plan quality rather than correctness". If the ranking were recorded as load-bearing for correctness, the next person to look at this would not dare simplify it, and there is a simplification worth having later - the exact ranking only matters when a collection's members disagree on which positions are operation keys, which needs a join whose outputPartitioning unions two sides without the mixed combinations.

[dongjoon-hyun]

Is this PR ready, @peter-toth ?

[peter-toth]

Is this PR ready, @peter-toth ?

Yes it is.

[dongjoon-hyun]

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 collectFirsts the first KP member and projects with itsexpressionDataTypes; 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 Options whose exclusivity only the scaladoc enforces, and the caller pays with guarded .gets (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.

[peter-toth]

Thanks @dongjoon-hyun, this was a good catch on 1. All four correctness items are addressed in a new commit, and 6 and 7 with them.

1. Confirmed, and it is a regression this PR introduces. Your repro plans and returns the right rows though, so it is worth writing down what it takes and where the throw actually comes from, because the mechanism is not the one in the finding.

key.row.get(0, TimestampType) does not throw: BaseGenericInternalRow.get ignores the requested type and hands back the Integer. The projected keys are then wrapped as TimestampType over Integer values, and distinct hashes them - which tolerates the mismatch too. The throw is in InternalRowComparableWrapper.equals, i.e. in the ordering built for the declared types, so it needs two projected keys that collide. In your example the two rows are 2020 and 2021, they project to different years, equals is never called, and the query plans.

Two rows in the same year and different buckets do it: (1, '2020-01-01') and (2, '2020-01-01') in both tables, reduced keys (50, 1) and (50, 2), both projecting to 50. On this PR's pushed head that is ClassCastException: class java.lang.Integer cannot be cast to class java.lang.Long at planning. And the same data shows why the projection has to happen: on master the query plans with no node and returns s = 10 and s = 20 for two rows that share a ts, where the answer is 30 and 30. So it is the wrong-result shape of this JIRA and the crash in one query, which is now a test.

The fix is to read the keys at the types they were built with rather than at the types their expressions declare. KeyedPartitioning.keyDataTypes returns partitionKeys.head.dataTypes - the schema each key row was written under, and the one it is hashed and compared under - and projectKeys uses it. The expressions describe the partitioning; a reducer made that description stale, but the rows never lied, so they are the sound source for reading a row. This is not the general fix for the stale expressions, which is #58335 plus the both-sides-reduce case left after it; it is the part this PR needs, because this PR is what starts projecting those keys.

2. Closed by the same change.GroupPartitionsExec now takes its base types from keyDataTypes too, so both sides read the keys the same way. The count the node produces no longer depends on which member collectFirst picks: the positions are the same, the keys are shared by the PartitioningCollection invariant, and the types now come from those keys rather than from a member's expressionDataTypes. That also let the numPartitionsAfter memo drop the data types from its key.

3. Agreed, and it is on my list as item 2 of the follow-up comment above - pre-existing, measured identical on master, and I will file it with the repro. You are right that the new cogroup test does not cover it: it uses identical layouts on both sides, so the positions agree and the hole stays invisible. A test for it belongs with the fix, since it fails on master today.

4. Fixed by skipping such a member instead of asserting. Relaxing the assert would not do: the empty position set flows on as the projection, and a node projecting to no position at all puts every partition into one - an earlier round of this PR measured exactly that (3 partitions to 1, and the result did not satisfy the distribution it was inserted for). So a member covering no position is not a candidate, and the child is shuffled, which is correct for a partitioning grouped on a reference-free expression. EnsureRequirementsSuite has a test that threw the AssertionError before.

5. Agreed, and it is the same thing you asked for in your first review (your #6 there, item 5 of my follow-up list): one matcher next to groupedSatisfies feeding both admission and projection, so the two cannot drift. That is the follow-up I want to take next in this area; it is more than a bug fix should carry.

6. Done, with Either rather than a new type.splitKeyedPartitionings returns Option[Either[KeyedPartitioning, (KeyedPartitioning, Option[Seq[Int]])]] behind a private type alias: Left satisfies as it is, Right needs a node with these positions, None needs a shuffle. The orElse(...).get and needsGrouping.get._2 are gone. One wrinkle worth knowing if you touch it: Left and Right have to be written scala.Left / scala.Right in this file, because catalyst.expressions has its own.

7. Done. The walk now only sorts the members into keyed and non-keyed, the non-keyed question is asked immediately after it, and the keyed analysis runs in a second method that is only entered when no plain member satisfies. Nothing touches a partition key before that answer.

8. Agreed on the per-plan cache, and it belongs on KeyedPartitioning. The distinct count that KeyedPartitioning.apply computes for isGrouped and discards is the clearest case. I left it out here because it is a catalyst change in a fix that otherwise only changes EnsureRequirements and the node, and because the memo does share across the members of one operator, which is the part that repeats within a call. The stacked-operators case is real and I will carry it into the follow-up above.

9. The second evaluation is the price of asking satisfies rather than re-deriving it. Only a member that fails satisfies pays it. Folding the two would mean computing groupedSatisfies once and combining it locally with the requiredNumPartitions gate - which is what an earlier revision of this PR did, and it dropped the gate for a grouped member: a blind review round measured an Exchange with 5 partitions on master against 3 partitions and no shuffle on the branch. Memoizing groupedSatisfies per distribution would fix it without that risk, and again on the partitioning rather than here.

KeyGroupedPartitioningSuite, EnsureRequirementsSuite, ValidateRequirementsSuite, PlannerSuite, ProjectedOrderingAndPartitioningSuite, UnionSuite, DataSourceV2Suite, AdaptiveQueryExecSuite, DistributionSuite and ShuffleSpecSuite are green - 575 tests. Both new tests fail on master, and on this PR's previous head they fail with the ClassCastException and the AssertionError respectively.

[ulysses-you]

I reviewed head 9fab82b (checked out in a worktree, both commits), read the full surrounding files at head and merge-base, traced every distribution type and plan shape through the old vs. new classification, and ran the suites locally (results below).

[ulysses-you]

the key-type axis is fixed in projectKeys but not in the ordering consumers — toGrouped, keyRowOrdering, and the OrderedDistribution arm still read at expressionDataTypes

• Where:partitioning.scala:547-556 (keyRowOrdering/keyOrdering/toGrouped), reached
  via KeyedPartitioning.createShuffleSpec subset branch at partitioning.scala:620-629; also
  EnsureRequirements.scala:93-94.
• Defect:toGrouped sorts with groupedKeyRowOrdering(expressionDataTypes) — the declared
  expression types — while this PR establishes that keys must be read at keyDataTypes. For a
  reducer-rewritten partitioning (Integer year values under a TimestampType-declared ts
  expression), sorting compares GenericInternalRow.getLong against a boxed Integer
  (rows.scala: getLong(ordinal) = getAs(ordinal)) ->
  ClassCastException: java.lang.Integer cannot be cast to java.lang.Long — the exact error this
  PR fixes elsewhere.
• Concrete failure:v2BucketingAllowCompatibleTransforms +
  v2BucketingAllowKeysSubsetOfPartitionKeys on; identity(ts)/years(ts) SPJ join (the PR's
  own test shape) -> any PartitioningPreservingUnaryExecNode above it (final HashAggregate)
  preserving the stale-typed KP -> a join above that calls createShuffleSpec -> subset branch
  runs projectKeys (now fixed) then .toGrouped -> CCE at planning with >=2 distinct projected
  keys. Same mismatch in the OrderedDistribution arm (EnsureRequirements.scala:93):
  RowOrdering.create(o.ordering, attrs) compares reduced keys with an ordering bound to the
  stale expression types, under v2BucketingAllowSorting.
• Peer code / invariant: the PR's own keyDataTypes scaladoc (partitioning.scala:550-564)
  states "the types each key was built with, and the ones it is hashed and compared under" — the
  hash/compare path (InternalRowComparableWrapper.equals/hashCode) uses the wrapper's own types
  and is safe; toGrouped's sort does not go through the wrapper's ordering, so it violates the
  stated invariant.
• Severity honesty:not a regression — on master these plans CCE'd even earlier (in
  projectKeys). The PR widens the set of plannable queries past the first crash and lands on
  the next instance of the same bug. Recommend either deriving keyRowOrdering from
  keyDataTypes (one line, consistent with the new scaladoc), or explicitly listing the
  remaining consumers in the follow-up JIRA alongside the ValidateRequirements gap.
• Verdict: CONFIRMED (line-by-line trace; cast mechanics verified in source). Runnable repro,
  works on master without the PR, and would still fail after it — appendable to
  EnsureRequirementsSuite:

test("SPARK-58968: createShuffleSpec must sort keys at the types they were built with") {
valexprTs=AttributeReference("ts", TimestampType)()
valexprId=AttributeReference("id", IntegerType)()
valfactory=InternalRowComparableWrapper
.getInternalRowComparableWrapperFactory(Seq(IntegerType, IntegerType))
// Reduced keys (year, bucket) under stale expressions (ts: TimestampType), as// GroupPartitionsExec.outputPartitioning reports after applying reducers.valkeys=Seq(InternalRow(2020, 0), InternalRow(2021, 1)).map(factory)
valkp=newKeyedPartitioning(Seq(exprTs, exprId), keys, isGrouped =true)
withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->"true") {
vale= intercept[ClassCastException](kp.createShuffleSpec(ClusteredDistribution(Seq(exprTs))))
assert(e.getMessage.contains("java.lang.Integer cannot be cast to java.lang.Long"))
}
}

Run:
build/sbt 'sql/testOnly org.apache.spark.sql.execution.exchange.EnsureRequirementsSuite -- -z "createShuffleSpec must sort keys"'

[peter-toth]

Confirmed, and fixed by deriving keyRowOrdering from keyDataTypes. Thanks - the trace is right, and it made me walk every reader rather than just this one.

Four places read key rows at the expressions' types and now read them at keyDataTypes: keyRowOrdering (so toGrouped, and PushDownUtils' key sort through keyOrdering), reduceKeys, the base types the reduce path passes on both sides in EnsureRequirements, and PushDownUtils' wrapper factory for the keys a scan reports after runtime filtering - that last one for consistency only, since it sees a scan's own partitioning, where the two sources coincide. The keyRowOrdering one also closes a latent divergence: GroupPartitionsExec.groupAndSortByKeys already sorted the reduced keys at the reduced types, so for exactly these partitionings the two sides of the contract in groupedKeyRowOrdering's scaladoc did not agree.

One place keeps expressionDataTypes, with a comment now saying why: ShuffleExchangeExec wraps lookup keys it evaluates from the expressions per row, and the stored keys it matches them against have to be declared the same way, so the two move together. For a reducer-rewritten partitioning neither choice works - the stored keys are in the reduced key space and the evaluated ones are not - so such a partitioning must not be shuffled onto at all. That gate is KeyedShuffleSpec.canCreatePartitioning, and closing it belongs with the stale-expression follow-up.

The OrderedDistribution arm you also point at is not a type swap. RowOrdering.create(o.ordering, attrs) binds the distribution's sort orders to the partition attributes, so with reduced keys the ordering is over the wrong space rather than merely at the wrong type; rebuilding it over the key types would have to carry the distribution's directions and null orderings, and refusing the partitioning is probably the better answer. A query does reach it, by the way - v2BucketingAllowSorting with a global sort on the partition key over a reduced join - and it throws the same ClassCastException on master, at the same place. On the follow-up list, named explicitly as you suggest.

Test added in EnsureRequirementsSuite, createShuffleSpec sorts the projected keys at the types they were built with - your repro inverted. It throws your ClassCastException before the fix and asserts the sorted projected keys after.

[ulysses-you]

Redundant local val

[peter-toth]

Fixed, thanks.

[dongjoon-hyun]

I reviewed this with a multi-angle pass (correctness, removed-behavior, cross-file tracing, efficiency/cleanup), with each candidate finding adversarially verified. The core splitKeyedPartitionings rework and its ranking/memoization logic held up under every scenario checked, and the test coverage is thorough. Leaving 8 inline comments: 2 correctness notes (one is a genuine planning-to-execution failure-mode shift worth a one-line gate, the other is the already-acknowledged deferred OrderedDistribution issue) and 6 minor efficiency/cleanup suggestions.

[dongjoon-hyun]

[correctness] This comment asserts that a reducer-rewritten partitioning "cannot be shuffled onto at all", but nothing enforces it: KeyedShuffleSpec.canCreatePartitioning checks only isGrouped and the expression shapes, never keyDataTypes == expressionDataTypes.

Moreover, this PR removes the accidental planning-time fail-fast that used to stop this path: pre-PR, createShuffleSpec's subset path read the keys at expressionDataTypes and threw CCE at planning for a reduced partitioning; post-PR it succeeds at keyDataTypes, so under allowCompatibleTransforms + v2BucketingShuffleEnabled + allowKeysSubsetOfPartitionKeys, an SPJ output (TimestampType-declared expressions over reduced IntegerType year keys) can become bestSpec, the other side gets shuffled onto it, and the valueMap below wraps the Integer keys at expressionDataTypes — CCE at execution, or silent misrouting for same-width reductions. The failure moves from planning time to execution time.

A keyDataTypes == expressionDataTypes clause in canCreatePartitioning would make this comment true and restore the planning-time fallback to a shuffle.

[peter-toth]

Taken, and it shipped separately as #58420 (SPARK-59120), now on master, 4.2, 4.3 and 4.x. KeyedShuffleSpec.canCreatePartitioning ends with partitioning.expressionsDescribeKeyShape, which is the clause you asked for. It compares shapes rather than plain types, because createPartitioning puts the other child's expressions over these keys, so a struct field name can legitimately differ with no reducer involved.

The whole keyDataTypes half left this PR with it, so the comment you flagged is not in this diff any more.

What the clause still does not catch is recorded next to it. Matching shapes are only a proxy: bucket(12) and bucket(8) reducing onto bucket(4) keep the type, pass the gate, and misroute rows anyway. SPARK-59121 replaces the proxy with the real test.

[dongjoon-hyun]

[correctness / acknowledged-deferred] This arm still binds the ordering to the expressions' declared attribute types and evaluates it over partitionKeys rows (keyOrdering.lteq below, and the sortBy in the else branch) — the one key-reading site in this rule not migrated to keyDataTypes. With v2BucketingAllowSorting + allowCompatibleTransforms, a global ORDER BY over a reduced join reaches this and throws CCE at planning (verified identical on master, and the PR description already defers it).

The deferral rationale is sound — reduced keys live in a different key space, so a type fix alone would give silently wrong ordering. One thought: an interim planning-time gate refusing keyDataTypes != expressionDataTypes here would trade part of the crash surface for a shuffle until the follow-up lands, though it would not cover same-typed reductions.

[peter-toth]

Still deferred, and now tracked rather than only mentioned in a description. The arm is one of three readers that reduced keys break in a way no choice of types can fix, and the keyDataTypes scaladoc names all three: this arm, the UnionExec key merge, and KeyedShuffleSpec.reducers. SPARK-59121 covers them together.

I did not take the interim gate. It would refuse the differently-typed reductions and let the same-typed ones through, so SPARK-59121 has to solve that class anyway and the gate would come straight back out. The crash is identical on master and needs v2BucketingAllowSorting, which is off by default.

[dongjoon-hyun]

[efficiency, minor] The memo keeps only the count and discards the projected keys, so whenever a count was computed (satisfying-narrowing member, requiredNumPartitions filter, or the maxBy ranking), the inserted GroupPartitionsExec re-runs the identical projectKeys + grouping in the same planning pass as soon as its outputPartitioning is consulted (and tryEnableSortedMerge's copy recomputes once more). The default-config single-candidate shape is unaffected thanks to the ranked.size == 1 fast path.

No drop-in fix — expectedPartitionKeys has different semantics and a cached-keys parameter would fight the documented independent re-derivation — so this is just a noted trade-off / possible follow-up.

[peter-toth]

Agreed, and left as it is, for the reason you gave. The comment at the memo now says what it keys on and why, so the trade is on the record rather than implied.

The default shape does not pay it. A single surviving candidate takes the maximal.size == 1 fast path and no count is computed at all.

[dongjoon-hyun]

[efficiency, minor] This eagerly pays an O(#partitions) projection for an early satisfying-but-narrowing member even when a later member of the same collection turns out to satisfy with full positions and zero key work (the earlier projection is then discarded). Narrow shape (subset config + asymmetric coverage across collection members) and planning-time only — worth a two-pass restructure only if it stays simple against the candidate-ordering invariants documented above.

[peter-toth]

Fixed in 3509a88.

satisfiedAsIs is two passes now instead of one loop over a var. The first find asks every admitted member whether its positions cover all of its expressions, which touches no partition key. Only the orElse pass computes a count. So a narrowing member no longer pays a projection that a later full-coverage member discards, and the scaladoc at the two finds says that is why the order is what it is.

[dongjoon-hyun]

[efficiency, nit]satisfies0 is nonGroupedSatisfies || (isGrouped && groupedSatisfies) and Partitioning.satisfies has no caching, so for a grouped member that fails satisfies, groupedSatisfies runs twice (SQLConf lookup, AttributeSet build, semanticEquals scans). Evaluating nonGroupedSatisfies/groupedSatisfies once as locals and deriving satisfies from them (plus the count gate) would avoid it. Constant-factor only.

[peter-toth]

Half of it, by construction. satisfies is asked first and || short-circuits, so every member that is admitted matches once. A member that fails satisfies still matches twice, which is the case you named.

I left that. Composing the two out of locals means putting satisfies0's structure in the caller, and nonGroupedSatisfies and keysSatisfy are both private now, which is what keysMaySatisfy exists to keep. Widening them to save a constant factor at planning time looks like the wrong trade. The comment above admitted records why the strict question is asked first.

[dongjoon-hyun]

[simplification, nit] The || k.expressions.isEmpty disjunct keeps alive a zero-expression KeyedPartitioning that no in-tree producer can construct (the scan, AliasAwareOutputExpression, and GroupPartitionsExec all guarantee at least one expression), and the branch is untested — even the reference-free test uses Seq(Literal(1)) and exercises the skip path. Dropping it would let a zero-expression member shuffle uniformly like the reference-free case and simplify the trickiest guard in this function.

[peter-toth]

Kept, and there is a test for it now: a KeyedPartitioning with no partition expressions is kept as it is.

Dropping it does not make such a member shuffle uniformly like the reference-free one. It reaches the shuffle branch under UnspecifiedDistribution, whose createPartitioning throws outright, so the guard is what stops a planning failure rather than a bad plan. The test's comment says that, and it also says what you did, that no in-tree producer builds one.

[dongjoon-hyun]

[reuse, nit] This hand-rolled recursive traversal duplicates the existing PartitioningCollection.flatten; flatten(partitioning) followed by a partition on KeyedPartitioning yields the same two sequences without local recursion or mutable buffers. The recursion predates this PR, but since the function is rewritten wholesale anyway, the cleanup is in scope.

[peter-toth]

Fixed in 3509a88.

splitKeyedPartitionings starts with PartitioningCollection.flatten(partitioning) and then splits by type. No local recursion, no mutable buffers.

[dongjoon-hyun]

[altitude, follow-up] This is now a third derivation of "operation-key positions", with deliberately different matching rules from KeyedShuffleSpec.keyPositions and createShuffleSpec's joinKeyPositions (the expression-level semanticEquals branch, tolerance of reference-free expressions) — and the co-partitioned path still derives positions the keyPositions-only way, so the expression-level case the new test pins is honoured on the single-child path only. Latent today (analyzed queries don't put a TransformExpression in a ClusteredDistribution), but consolidating this next to keyPositions on KeyedPartitioning/KeyedShuffleSpec in the follow-up would make the divergence visible in one place.

[peter-toth]

Agreed, and it is the follow-up's second piece, one derivation of the operation-key positions. The count half already moved onto the type here, as KeyedPartitioning.numPartitionsProjectedOn, so what is left there is the positions.

The co-partitioned path keeping its own derivation is deliberate for now, and the description says so. Projecting inline there would leave the multi-child block deriving positions from an already projected partitioning and applying them to the unprojected keys.

[peter-toth]

Moving this back to draft while it waits on two other PRs. It is stacked on #58351, and the key-type half has been split out into #58420. I will rebase onto plain master and take this out of draft once both are in.

[peter-toth]

@dongjoon-hyun@ulysses-you this is out of draft and ready for review.

It is rebased onto master, which now carries both PRs it was stacked on, #58351 (SPARK-59057) and #58420 (SPARK-59120). The keyDataTypes half left this PR and went out as #58420, so what is left here is only the classification change.

Since your round, two of your findings are fixed in the code, one of them shipped as #58420, and the rest are answered in their threads. The description is rewritten around the current shape.

One test was added, a transform position is kept when its reference is a cluster key. It is the only one where the reference-level match in clusterKeyPositions decides a kept position, which is the shape the single-reference soundness argument is about. master inserts no node at all for it, so the partitions sharing a bucket stay apart.

[ulysses-you]

thank you @peter-toth !

What I verified (the invariant audit)

Equivalence alignment. Both the new count and the actual coalescing share one representation. numPartitionsProjectedOn counts projectKeys(positions)._2.distinct.size; GroupPartitionsExec.grouping groups via reducedKeys.groupMap(_._1) on the sameInternalRowComparableWrapper. equals = RowOrdering.compare==0, routing float/double through SQLOrderingUtil.compareDoubles (if (x==y) 0 else Double.compare) — so -0.0/0.0 compare equal and all NaN bit-patterns collapse. Predicted count and executed grouping cannot diverge. Checked the -0.0/NaN/collation cells; the wrapper hash is isCollationAware=true, predates this PR, and is used consistently — no new mismatch introduced.

Soundness of keeping a position.clusterKeyPositions keeps a position only if the expression or one of its references is a cluster key. The reference-level arm is safe only because admission through the subset branch of keysSatisfy requires expressions.forall(_.references.size == 1) (partitioning.scala:287) — confirmed present, so a kept expression is a function of a single cluster key and coalescing on it cannot separate rows sharing that key. The multi-reference counter-example (keep a+b when only a is clustered) is unreachable through keysMaySatisfy.

Single-child vs co-partitioned split.isCoPartitioned = childrenIndexes.length > 1 is computed from requiredChildDistributions only, so it is identical to master's value despite being hoisted. When true, clusterKeyPositions returns all positions, so the inline path degrades to Left(keep) or Right(_, None)(pure coalesce) and never projects — projection stays owned by the multi-child block (checkKeyGroupCompatible → applyGroupPartitions for joins; withJoinKeyPositions for cogroup). Traced the cogroup case: no double-application of positions. GroupPartitionsExec re-derives its member via collectFirst, but members share the partitionKeys reference and arity (PartitioningCollection.checkKeyedPartitioningInvariant), so projecting by the chosen member's positions lands on the same key columns whichever member executes — data co-location holds; only reported expressions may differ (metadata, not data).

satisfies over-claim is neutralized. A grouped superset satisfies via keysSatisfy but positions.size < expressions.length and the projection merges something, so satisfiedAsIs (both finds) misses and it falls through to a projecting GroupPartitionsExec — the actual fix. Conversely, when the projection merges nothing (numPartitionsAfter == numPartitions), keeping the wider partitioning is genuinely correct and strictly better for downstream co-partition. requiredNumPartitions: as-is satisfaction enforces the count through satisfies (partitioning.scala:229), and grouping satisfaction filters eligible by the post-projection count rather than vetoing the winner — a candidate that can honour the count can't lose to one that can't.

Cross-path / ordering axes. Planner-level; no codegen/interpreted asymmetry (GPE is exchange-like, doExecute only). AQE runs the same rule in queryStagePreparationRules, and the top-k window test asserts the wrong result under default (AQE-on) config. GPE's outputOrdering is untouched and conservative; ordering restored by the subsequent sort pass. OrderedDistribution selection is equivalent to master, and the sliding(2)→zip(drop(1)) change is a genuine MatchError fix (single-key partitioning). StatefulOpClusteredDistribution is a Distribution not a ClusteredDistribution, never matches the subset branch, always shuffles — consistent with its contract. UnspecifiedDistribution/AllTuples/BroadcastDistribution still resolve to keep/shuffle/broadcast.

Verdict: every enumerated axis held under tracing. I did not build/run the suites (a from-scratch sql/core compile is disproportionate where tracing found no candidate bug to reproduce); the author's report that 17/24 new tests fail without the fix matches my read that these pin real behavior rather than tautologies.

[dongjoon-hyun]

+1, LGTM. Thank you, @peter-toth and @ulysses-you .

[peter-toth]

Merge Summary:

• merged into master 88d7ee1
• merged into branch-4.x c2b4cd4
• merged into branch-4.3 93ab1bb

Posted by merge_spark_pr.py

[peter-toth]

Thank you all for the review!