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

[peter-toth]

Backport of #58527 to branch-4.3

Clean cherry-pick of 4605662b6e4, no tailoring. Rebased onto dacfdba1818, @ulysses-you's follow-up that fixes the subset-join-key test on this branch, as he asked on #58279.

Green on this branch: ShuffleSpecSuite 18 tests, and EnsureRequirementsSuite + KeyGroupedPartitioningSuite 195 tests, all passing. dev/lint-scala is clean.

---

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

[peter-toth]

Thank you @ulysses-you and @uros-b.