[SPARK-59080][SQL] Pick one ShuffleSpecCollection member for the SPJ pushdown and the re-shuffle

[peter-toth]

What changes were proposed in this pull request?

EnsureRequirements stops asking a ShuffleSpecCollection for a single answer. It resolves the one member the matched children agreed on, preferring the finest when several qualify, and uses that member to build a re-shuffled child's partitioning.

• flattenSpec replaces the head read. It recurses, because ShuffledJoin.outputPartitioning builds PartitioningCollection.fromPartitionings(Seq(left, right)) for an inner join, so a chain of same-key joins nests collections.
• The chosen member has to be compatible with every matched child. When no member is, there is no shared layout, so every child takes the ordinary shuffle. Reaching that needs three or more clustered children, and no operator has three today.
• The joinKeyPositions pushed into a compatible child now come from that child's own matching member, because they index into that child's partition expressions.
• ShuffleSpecCollection.createPartitioning is untouched. Its require stays as a guard, and a new unit test pins it.

Why are the changes needed?

Under spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled, KeyedPartitioning.createShuffleSpec projects each member of a PartitioningCollection onto its own join-key subset and drops the duplicate keys that projection creates. The members of the resulting ShuffleSpecCollection can therefore end up with different numPartitions. EnsureRequirements then asks the collection for a shuffle template, and ShuffleSpecCollection.createPartitioning requires all members to agree:

java.lang.IllegalArgumentException: requirement failed: expected all specs in the collection to have the same number of partitions

so planning fails outright. With items partitioned by [identity(id), identity(arrive_time)], one row per split, purchases unpartitioned, and v2BucketingShuffleEnabled=true, partiallyClusteredDistribution=false, allowKeysSubsetOfPartitionKeys=true:

SELECT /*+ MERGE(i, p) */ id, t1, t2, i.price AS purchase_price, p.price AS sale_price
FROM (SELECT id, arrive_time AS t1, arrive_time AS t2, price FROM testcat.ns.items) i
JOIN testcat.ns.purchases p ON i.id = p.item_id AND i.t1 = p.time

Selecting arrive_time twice under two aliases makes the alias cross-product produce members that cover different numbers of join keys, which is where the counts diverge.

The collection cannot answer that question locally. isCompatibleWith succeeds when any member matches, so the collection alone never said which member the two sides agreed on, and createPartitioning fell back to specs.head, whichever the alias cross-product enumerated first. Narrowing the collection to its finest members would satisfy the require, but it would still be a guess: the right member is the one the other side matched, and that is only visible in EnsureRequirements.

Does this PR introduce any user-facing change?

Yes. The query above failed to plan and now runs, producing one shuffle and the right rows.

The joinKeyPositions half is user-facing too. I originally wrote here that it was latent, on the grounds that a cogroup's grouping key is synthesized so neither side stays keyed. @sunchao pointed out that this is only true of the Scala CoGroupExec, whose key comes from an AppendColumns that no KeyedPartitioning satisfies. A Pandas or Arrow cogroup groups on real columns, so two keyed children do reach the per-child branch, and this suite already had a FlatMapCoGroupsInPandasExec test with two of them.

So: with two keyed children whose partition expressions are laid out differently, the second side is handed the first side's positions and ends up grouped on its other partition column. The plan test below reproduces it and fails without the fix, reporting List(Some(List(1)), Some(List(1))) where List(Some(List(1)), Some(List(0))) is right. I have not built an end-to-end query for it.

The only part that stays latent is the three-or-more clustered children case, which no operator has.

How was this patch tested?

Four new tests. Each was measured against the same commit with only the EnsureRequirements change reverted.

test	on base
KeyGroupedPartitioningSuite: both sides of the join land on the same collection member	fails with the require above
EnsureRequirementsSuite: the re-shuffled side lands on the member the keyed side was matched on	fails with the require above
EnsureRequirementsSuite: pushed-down positions index into the child's own partition expressions (a Pandas cogroup over two keyed children)	fails
ShuffleSpecSuite: a collection whose members cover different key subsets disagrees	passes, by design

The last one pins the guard rather than the fix. It asserts that the members disagree on purpose, that every one of them stays available for isCompatibleWith, and that asking the collection for a single partitioning throws. That turns the require from untested prose into a pinned contract, which matters now that the method has no production caller.

Green: ShuffleSpecSuite, EnsureRequirementsSuite, KeyGroupedPartitioningSuite, 202 tests in all. dev/lint-scala is clean.

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

Generated-by: Claude Code

[dongjoon-hyun]

+1, LGTM. Thank you, @peter-toth !

[sunchao]

One nonblocking P3 finding; no P1/P2 defects found. Reviewed 6136919314809cd61473fb045d0521c202b41fff with five independent reviewers.

Finding

[P3] Correct the claim that keyed cogroups cannot reach this branch. The new comment overlooks Pandas/Arrow cogroups, which preserve named grouping expressions. The existing SPARK-58968 test in this suite already exercises two keyed cogroup children. Qualify the comment and the PR description’s claim that this correction is latent. The implementation itself improves this reachable path.

Validation

Coverage	Result
Affected suites in CI	202 passed
Local focused tests	69 passed
Additional nested-layout and fallback probes	6 passed
Reverted implementation	Reproduced all three claimed regression failures
Full-patch and last-commit whitespace checks	Passed

Local validation used verified CI binaries and source overlays, not a clean local build. The PR head remained unchanged. Nothing was posted to GitHub.

[LuciferYang]

createKeyedShuffleSpec still resolves a PartitioningCollection with collectFirst, taking the first satisfying member, while this PR makes the per-child branch resolve the finest member both sides agreed on. The default requireAllClusterKeysForCoPartition=true keeps partial-coverage members out of the candidates, so this only shows up with requireAllClusterKeysForCoPartition=false plus allowKeysSubsetOfPartitionKeys=true: when both sides report alias cross-products with a coarse member first, SPJ pairs the coarse members and joins on them, leaving the finer pairing unused; with coarse first on one side and fine first on the other, the picks fail to pair and the per-child branch takes over, grouping each side on its finest matching member, so what is skipped is the SPJ partition-value pushdown, not a shuffle. Results stay correct. If you align it, the choice cannot be an independent per-side finest (a side with only coarse members would then fail to pair); it has to pair like bestMemberOpt does, by finest member compatible with some member of the other side. Either fixing it here or a follow-up works for me.

[peter-toth]

Good catch, and I took it into #58531 rather than a separate ticket. createKeyedShuffleSpecs now returns every member's spec and checkKeyGroupCompatible picks the pair, with your constraint respected: not an independent per-side finest, but the pair that agrees on the keys and offers the most parallelism.

Two things I measured that are worth passing back.

Your reachability framing is slightly off, in a way that does not change your conclusion. requireAllClusterKeysForCoPartition = true does not always keep a collection down to one candidate: I instrumented createKeyedShuffleSpec over KeyGroupedPartitioningSuite and EnsureRequirementsSuite and got 31 collection reaches, of which 2 had two qualifying members under the default. But in every one of those the qualifying members produced identical specs, so the pick was not observable. The observable case does need the relaxed conf, as you said — just not for the reason that the members are filtered out.

The shape is narrower than it looks, too. Two members can only differ in coverage if the clustering is wider than the partitioning arity, since the collection requires its members to have matching arity. With that, the repro is: each side offering a member for a different clustering-key subset, in the opposite order, so taking each side's first member pairs two that do not agree. On the fix that pairs correctly and no shuffle is needed at all; without it checkKeyGroupCompatible declines and each side is merely grouped on its own keys.

I also checked the sibling collectFirst at the top of the partially-clustered branch, which recovers the original KeyedPartitioning while the positions come from the picked member. That one is harmless and I am not touching it: projectKeys reads only partitionKeys and the types derived from them, and the collection invariant makes partitionKeys a shared reference, so any member gives the same answer.

[LuciferYang]

bestSpecOpt is still picked with ShuffleSpecCollection.numPartitions, which comes from specs.head. A collection like [coarse head with 2 partitions, fine member with 200] competes as 2 in the maxBy, loses to a sibling child's plain 100-partition spec, and the side holding the fine member gets re-shuffled onto 100 partitions when it could have stayed unshuffled as the best. The selection line predates this PR, and in this shape it behaves exactly as before (the losing side took the same full re-shuffle), so this is purely pre-existing. Not blocking; worth a follow-up comparing best candidates at member level at the call site (changing numPartitions itself would affect other readers).

[peter-toth]

Agreed, and this is exactly what I had queued as the follow-up: #58531 (SPARK-59256). It takes the max over the flattened members at the call site, for the reason you give — changing numPartitions itself would affect other readers.

Worth naming the other reader, since it turned out to want a different aggregation: SinglePartitionShuffleSpec.isCompatibleWith reads other.numPartitions == 1, reached from ValidateRequirements, and there the right answer is exists, not max. Two consumers wanting two different aggregations is what convinced me the collection should not answer numPartitions at all rather than answer it better, so that PR removes it along with createPartitioning.

[peter-toth]

Thanks @sunchao, and you are right — I checked and it is worse than a wording problem.

My "the cogroup key is synthesized" claim came from a measurement I did for a different follow-up, and that measurement only covered the Scala CoGroupExec, whose key comes from an AppendColumns no KeyedPartitioning satisfies. A Pandas or Arrow cogroup groups on real columns, so two keyed children do reach the per-child branch, exactly as you say, and this suite already had a FlatMapCoGroupsInPandasExec test with two of them.

So the joinKeyPositions half is not latent. I rebuilt the plan test on FlatMapCoGroupsInPandasExec instead of the synthetic parent I had used: both sides declare the same two key columns in the opposite order, so the cogroup key sits at position 1 on the left and 0 on the right, and both project onto the same key set. Without the production change the right side is handed the left side's positions and ends up grouped on its other partition column — List(Some(List(1)), Some(List(1))) where List(Some(List(1)), Some(List(0))) is right. I have not built an end-to-end query for it, so I am not claiming a wrong-results reproduction, only the plan.

Pushed as f6a1442, and the description no longer calls that half latent. The only part that stays latent is the three-or-more clustered children case.