[SPARK-59256][SQL] Choose a ShuffleSpecCollection member by pairing, not by enumeration order

[peter-toth]

What changes were proposed in this pull request?

Stacked on #58527 (SPARK-59080), which is the first two commits here. Please review only the last commit, and do not merge this before #58527. Draft until then.

A PartitioningCollection offers several layouts, and which one is right depends on what the other side matched. Three places in EnsureRequirements decided it by enumeration order instead, and none of them could see both sides:

1. the per-child branch that shuffles a child read specs.head through ShuffleSpecCollection.createPartitioning — fixed by [SPARK-59080][SQL] Pick one ShuffleSpecCollection member for the SPJ pushdown and the re-shuffle #58527;
2. the child ranking, finalCandidateSpecs.values.maxBy(_.numPartitions), reads the collection's numPartitions, which is specs.head's;
3. createKeyedShuffleSpec collapses a collection to one member with collectFirst, per side, before either side has seen the other.

This PR is 2 and 3, and it makes the type say why a collection cannot answer alone. Both were raised by @LuciferYang in review of #58527.

The type split. A new LeafShuffleSpec sub-trait carries the two single-member methods:

ShuffleSpec (sealed)             isCompatibleWith, canCreatePartitioning, flatten
  <- LeafShuffleSpec             + numPartitions, createPartitioning
       SinglePartitionShuffleSpec, RangeShuffleSpec, HashShuffleSpec,
       NullAwareHashShuffleSpec, CoalescedHashShuffleSpec, KeyedShuffleSpec,
       ShufflePartitionIdPassThroughSpec
  <- ShuffleSpecCollection

Sealing is what makes the split mean something: it is why flatten can return Seq[LeafShuffleSpec] with no fallback case, and it is checked — nothing outside partitioning.scala extends ShuffleSpec. flatten replaces a private helper #58527 added to EnsureRequirements, and lives on the hierarchy because that is where the totality comes from.

The two answers each caller now gives itself. The collection's createPartitioning had no production caller after #58527, and only a runtime require stood between a future caller and a wrong partitioning. numPartitions had two consumers wanting two different aggregations:

• EnsureRequirements ranks the children by parallelism to pick the reference layout. It wants the max, and it now takes it over the flattened members.
• SinglePartitionShuffleSpec.isCompatibleWith asks whether the other side is a single partition. It wants exists, and it now unwraps a collection the way every other spec already does.

The pairing. createKeyedShuffleSpecs returns every member's spec, and checkKeyGroupCompatible picks the pair that agrees on the keys and offers the most parallelism. The derivation of leftSpec / rightSpec is the only thing that changes; everything after it works on the chosen pair as before. When no pair agrees, any pair fails the same checks and the method returns None, as it did before.

The pick cannot be an independent per-side finest, as @LuciferYang noted: a side whose only members are coarse would then fail to pair. It is also one rule and not two — a pair that is co-partitioned as it stands is not preferred over a finer one that needs a grouping node. I could not measure a case where the two rules differ, so the extra step would have been an unmeasured refinement.

Why are the changes needed?

isCompatibleWith succeeds when any member matches, which is the whole point of the type. The two single-member questions are not wrong in the same way. createPartitioning has no local answer at all, since the right member is the one the other side matched. numPartitions has an answer, but two different ones depending on who asks, which is the same thing as not having one.

For the pairing the cost is concrete: two sides pick members that do not agree, checkKeyGroupCompatible declines, and the join loses the storage-partitioned pushdown even though a pairing existed.

Does this PR introduce any user-facing change?

No, under a default configuration. The three behaviour changes are gated differently, though, and it is worth not conflating them.

The ranking and the SinglePartitionShuffleSpec change are no-ops unless spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled is on. PartitioningCollection requires its members to agree on numPartitions, and every spec reports its own partitioning's count, so max equals head and exists(_.numPartitions == 1) equals head.numPartitions == 1. Only the projection that config turns on can make members disagree.

The pairing reads no count, so that argument does not cover it. It changes the plan when a collection has two members that both qualify and sit at different clustering key positions, which needs spark.sql.requireAllClusterKeysForCoPartition off — that check is what otherwise refuses a member not covering every key. Not the default either, so still no default-configuration change.

One thing worth flagging rather than leaving to be found. The only production consumer of the changed SinglePartitionShuffleSpec answer is ValidateRequirements, through specs.tail.forall(_.isCompatibleWith(specs.head)); EnsureRequirements never reaches it, because SinglePartitionShuffleSpec.canCreatePartitioning is false and the collection's is a forall. And there the answer is about a layout nobody builds: a projected KeyedShuffleSpec's numPartitions is the count after the GroupPartitionsExec that EnsureRequirements would insert, while ValidateRequirements inserts nothing. So "the other side is a single partition" can be true of a child that physically has more. The head read had that same problem on a different arbitrary member, so this is a pre-existing conflation rather than a regression, but exists does widen it from one member to any. I could not build a query that reaches it, and I did not add a test, because pinning an answer I am not sure is right would be worse than leaving it uncovered.

How was this patch tested?

Four new tests. A plain fail-on-base measurement is not available here: the tests name LeafShuffleSpec, so they cannot compile against the parent commit. Instead I reinstated each old decision under the new types, one at a time, and ran the tests against that.

test	old decision reinstated	result
EnsureRequirementsSuite: a collection is ranked on its best member, not on whichever came first	ranking reads flatten(spec).head	fails
EnsureRequirementsSuite: the join is planned on a member pair, not on each side's first member	(leftCandidates.head, rightCandidates.head)	fails
ShuffleSpecSuite: a single-partition side matches a collection through any member	isCompatibleWith(specs.head)	fails
ShuffleSpecSuite: flattening reaches the members of a nested collection	n/a — pins the helper's contract, which the parent's private version also had	passes, by design

The first three assert their member counts as an ordered Seq rather than a Set, so each pins its own premise: that the member a head read would take is the wrong one.

I also wrote two further tests, for the pair ranking and for the pairing's config gating, and deleted them again: both passed with the old decision reinstated, so they measured nothing. A PartitioningCollection forces matching arity on its members, so "one member covers fewer clustering keys than another" needs a partition expression outside the clustering, which neither fixture built. Rather than ship tests that assert nothing, I am naming the gap: the ranking among agreeing pairs is not covered, only the fact that pairing happens at all.

Deleted: createPartitioning: other specs had a case asserting that a collection delegates to specs.head. The method is gone and the type now rejects the call. The same test's expected className in the UNSUPPORTED_CALL error moves from ShuffleSpec to LeafShuffleSpec, since the defaulted method moved. A dangling pre-split scaladoc for ShuffleSpec, stranded above ShufflePartitionIdPassThrough, is deleted rather than left to contradict the new trait doc.

CoalescedHashShuffleSpec.from is typed LeafShuffleSpec too — it is structurally the one spec the coalesced spec was built from, which is the refactor's own argument applied to the one field that had escaped it. Refining the declared return types of HashPartitioning.createShuffleSpec and NullAwareHashPartitioning.createShuffleSpec to their own spec types means that needs no cast.

Also corrected: GroupPartitionsExec carried the only explanation of why re-deriving a different collection member is safe, and it explained it in terms of createKeyedShuffleSpec's collectFirst, which no longer exists. It now gives the argument that actually holds — the members share their partitionKeys reference and arity, and reducing reads the key values at their own types, so every member reduces to the same values.

Green: ShuffleSpecSuite, DistributionSuite, EnsureRequirementsSuite, ValidateRequirementsSuite, KeyGroupedPartitioningSuite, PlannerSuite, GroupPartitionsExecSuite, 323 tests in all. dev/lint-scala is clean.

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

Generated-by: Claude Code