[SPARK-58968][SQL] Fix data correctness issue when SPJ allowKeysSubsetOfPartitionKeys

[ulysses-you]

What changes were proposed in this pull request?

For a non-join operator, when a grouped KeyedPartitioning satisfies a ClusteredDistribution whose clustering keys are a strict subset of the partition keys (via v2BucketingAllowKeysSubsetOfPartitionKeys), EnsureRequirements now inserts a GroupPartitionsExec that projects the partition keys to the operation keys, so partitions sharing the same operation key are coalesced.

Why are the changes needed?

When v2BucketingAllowKeysSubsetOfPartitionKeys is enabled and a storage-partitioned source is partitioned by more keys than the operation requires (e.g. a table partitioned by (id, name) with a window PARTITION BY id), the grouped KeyedPartitioning satisfies the ClusteredDistribution through the subset-key relaxation, but the partitions are still grouped by the full partition keys rather than by the operation keys. Without coalescing them, the downstream operator ranks/aggregates each partition independently, producing incorrect results.

For example, a top-k window over PARTITION BY id on a (id, name)-partitioned table would surface id=1 twice (once per (1,'aa')/(1,'bb') partition) instead of once.

Does this PR introduce any user-facing change?

Yes, it fixes a data correctness issue: storage-partitioned windows (and other non-join operators) with PARTITION BY a subset of the partition keys now produce correct results when v2BucketingAllowKeysSubsetOfPartitionKeys is enabled.

How was this patch tested?

Added regression tests in KeyGroupedPartitioningSuite covering:

• window top-k over PARTITION BY a subset of the partition keys
• window top-k over a duplicated PARTITION BY key
• window top-k over union output partitioning

Verified org.apache.spark.sql.connector.KeyGroupedPartitioningSuite and org.apache.spark.sql.execution.exchange.EnsureRequirementsSuite pass.

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 think about this issue tomorrow.

[gengliangwang]

0 blocking, 3 non-blocking, 0 nits.
The fix is correctly scoped and well gated behind its config, with one design question about how the new guard picks the operators it applies to.

Design / architecture (1)

• sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:67: isJoin keys on the ShuffledJoin marker trait rather than on whether the operator enters the two-child co-partitioning block, so the guarded and handled sets do not match in either direction. -- see inline

Suggestions (2)

• sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:4210: Add an aggregate (GROUP BY a subset) test and a two-clustered-child (cogroup) test; the three new tests are all top-k windows. -- see inline
• sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:124: Only joinKeyPositions is read here, but createShuffleSpec also builds a projected KeyedPartitioning over every partition key that this call site discards. -- see inline

Verification

CI is fully green on this head (31 check runs, no failures or pending), so my pass focused on the parts a green build does not settle.

Four things I checked and found sound. The unchecked asInstanceOf[KeyedShuffleSpec] cannot fail, because KeyedPartitioning.createShuffleSpec constructs KeyedShuffleSpec(this, distribution) and returns either it or a copy of it on both branches. The strict-subset comparison against expressions.indices is exact rather than incidentally correct: joinKeyPositions is built as keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2), always ascending and always a subsequence, so no permutation of clustering keys can make a full key set look like a subset. Stateful streaming is structurally excluded, since StatefulOpClusteredDistribution is not a ClusteredDistribution. And the sibling nonGroupedSatisfiesAsIs branch cannot hide the same unfixed bug, because nonGroupedSatisfies delegates to the base Partitioning.satisfies0, which answers only UnspecifiedDistribution and AllTuples and returns false for every ClusteredDistribution.

One thing I checked that is safe but not for the stated reason. SortMergeAsOfJoinExec is a third ShuffledJoin, so isJoin excludes it, yet checkKeyGroupCompatible returns None for it - the guard's justification does not cover it. It is nonetheless correct today: with v2BucketingShuffleEnabled at its default false, KeyedShuffleSpec.canCreatePartitioning is false, bestSpecOpt is empty, and each clustered child falls through to a plain ShuffleExchangeExec, which establishes the distribution outright. An as-of join gets correct results, just not storage-partitioned execution - the same as before this PR.

What I could not settle is the CoGroupExec interaction in the finding below; it needs a run with two non-default configs, which is why it is posted as a question rather than an assertion.

PR metadata suggestions

• Document: the change also affects non-join operators that have two clustered children (CoGroupExec and the Python cogroups), not just single-child operators like the window in the example - that is where the new inline grouping meets the existing two-child co-partitioning block.
• Name aggregates explicitly alongside the window example: GROUP BY a subset of the partition keys hits the identical branch and duplicates groups without the fix, and it is the shape most users will recognize.

[gengliangwang]

This guard states a different condition from the one the comment appeals to, and the two sets come apart in both directions.

The comment says joins are excluded because checkKeyGroupCompatible handles their projection. But that helper matches only SortMergeJoinExec and ShuffledHashJoinExec, while the block that calls it is entered for any operator with two clustered children (parent.isDefined && children.length == 2 && childrenIndexes.length == 2).

One direction is currently harmless: SortMergeAsOfJoinExec is also a ShuffledJoin, so it is excluded here yet unhandled there. It still gets correct results, but by falling through to a plain shuffle rather than by anything this comment describes.

The other direction concerns me, and I could not finish confirming it - hence a question. CoGroupExec requires ClusteredDistribution on both children (objects.scala:638) and is not a ShuffledJoin, so this branch does run for it. children is then reassigned with the wrapped child, and specs a few lines down are computed from children(i).outputPartitioning - the projected partitioning. checkKeyGroupCompatible returns None, so areChildrenCompatible is false and each clustered child reaches withJoinKeyPositions(child, joinKeyPositions), which rewrites the node in place via g.copy(joinKeyPositions = Some(positions)). Those positions index the projected expression list, but GroupPartitionsExec applies them to its child's unprojected partitioning. For tables partitioned by (name, id) cogrouped on id: this branch computes [1] and inserts a node projecting to id; the block recomputes [0] against the now one-element list and overwrites, so the node projects position 0 of (name, id) and coalesces by name.

I did not verify the observable result. Reaching the overwrite also needs v2BucketingShuffleEnabled on, since it gates KeyedShuffleSpec.canCreatePartitioning - with the default false, bestSpecOpt is empty and the fallback shuffles and unwraps the node instead. Could you confirm whether a cogroup over two storage-partitioned tables with both configs on and a non-leading grouping key produces wrong groups?

Either way I'd gate on the structural fact rather than the trait: run this branch only when the operator will not enter that block, i.e. childrenIndexes.length <= 1. That is what the comment is really appealing to, it covers cogroups and as-of joins without enumerating join classes, and it keeps one owner of the projection per operator so a future ShuffledJoin or a new parent case in checkKeyGroupCompatible cannot move the boundary silently.

[gengliangwang]

Consider adding an aggregate variant next to these. HashAggregateExec requires ClusteredDistribution(groupingExpressions), so SELECT id, sum(price) FROM items GROUP BY id on an (id, name)-partitioned table reaches this exact branch, and without the fix it emits one row per (id, name) partition instead of per id - a duplicated group, which is arguably a more familiar symptom than a ranking artifact. All three new tests go through ROW_NUMBER() OVER (PARTITION BY ...), so a reader could reasonably think the fix is window-specific.

A cogroup test would be worth more still, because that case is one the change newly affects rather than fixes: df1.groupByKey(...).cogroup(df2.groupByKey(...)) over two storage-partitioned tables is a two-clustered-child operator that is not a ShuffledJoin, so it takes the new branch and then also goes through the co-partitioning block. That is the interaction I asked about on the isJoin line.

[gengliangwang]

Only spec.joinKeyPositions is used, but with the config enabled createShuffleSpec also runs projectKeys(joinKeyPositions)._2 across every partition key and then .distinct on the result, to build a projectedPartitioning that this call site drops - and GroupPartitionsExec recomputes the same projection later from child.outputPartitioning. That is two O(number of partitions) passes per qualifying operator at planning time, on a path that exists precisely for tables partitioned finely enough to need coalescing.

KeyedShuffleSpec(kp, c).keyPositions gives the same information without the discarded projection, and reads more directly as "which partition expressions does this operation actually key on" than routing through a shuffle-spec factory. The cast also goes away with it.

[peter-toth]

@ulysses-you I ended up trying a different shape for this one - classifying a KeyedPartitioning by what still has to happen to it rather than by isGrouped, so a single place inserts GroupPartitionsExec. Still a draft, and not a review of this PR: #58262. Linking it only so the two approaches can be compared.