[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. Four 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. SinglePartitionShuffleSpec.isCompatibleWith reads other.numPartitions, the same head;
4. createKeyedShuffleSpec collapses a collection to one member with collectFirst, per side, before either side has seen the other.

This PR is 2, 3 and 4, and it makes the type say why a collection cannot answer alone. 2 and 4 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

flatten replaces a private helper #58527 added to EnsureRequirements and returns Seq[LeafShuffleSpec]. Sealing is not what makes it total — it is a virtual method with an implementation per kind. What sealing buys is that "every spec is one layout or a choice of layouts" becomes a guarantee, which is what makes that return type honest, and that SinglePartitionShuffleSpec's two-case match compiles as exhaustive. Nothing outside partitioning.scala extends ShuffleSpec.

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, and it now answers forall over the members. Deliberately not the exists the other specs use for a collection: they ask whether some member matches them and then plan on that member, while this asks a property of the child itself. The child is whichever member the other side picks, so a single partition has to hold for every one of them. The head read could claim a co-partitioning that does not hold, because a member's count is the count after the projection createShuffleSpec promises, not the child's own — a three-partition child whose head member projects to one used to answer "co-partitioned with a single partition".

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.

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. The zero-node path is not what such a preference would protect: it needs left.outputPartitioning.exists(_ == leftPartitioning), which only an identity projection satisfies, and an identity projection reports the side's own physical count, which is maximal. maxByOption keeps the earliest element on a tie (Maximized folds with a strict gt), and the earliest pair is the old per-side pick, so the pairing never loses the as-is path that pick used to find.

What the ranking does trade, and this one is measured rather than hypothetical: between two pairs that both need grouping nodes, it takes 4 partitions with 2 empty ones on one side over 2 exactly-matched partitions. More parallelism for some padding. If reviewers prefer the other side of that trade, the change is one find before the maxByOption.

When no pair agrees, each side's first member is reported — what the per-side pick took — and the checks below fail on it exactly as before, so the method returns None as it did.

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?

Yes, and correcting that is the main thing this revision does. I first wrote "no" here on the basis that both gating configs default to false. spark.sql.sources.v2.bucketing.pushPartValues.enabled defaults to true, and it is all the push branch in checkKeyGroupCompatible needs, so the pairing changes plans out of the box.

Measured on the transform-difference shape with no config set at all: taking each side's first member makes the join decline, bestSpecOpt is then empty because a keyed spec cannot be a shuffle reference without v2BucketingShuffleEnabled (which is off by default), and both children are shuffled onto the default 200 partitions. With the pairing the join runs with no shuffle and two grouping nodes. It is an improvement, but it is a plan change for anyone with V2 bucketing on, not a no-op.

The other two changes are no-ops unless allowKeysSubsetOfPartitionKeys is on: PartitioningCollection requires its members to agree on numPartitions, and every spec reports its own partitioning's count, so max equals head and the forall equals the head read.

For the pairing to matter at all, a collection's members have to differ in a way areKeysCompatible can see, and there are two ways. spark.sql.requireAllClusterKeysForCoPartition off lets members cover different clustering keys. A transform difference does it with that requirement at its default, because a collection forces its members to agree on the key rows, and through those on the transforms' result types, but not on the transforms themselves.

One path changes reachability rather than behaviour, and it is worth naming: reducersBothWays now runs on a pair the old code could not form, so a connector whose Reducer.resultType() violates the r(f1(x)) = f2(x) contract can raise storagePartitionJoinIncompatibleReducedTypesError where the join used to fall back silently. That is what the error exists to catch, but it is a new failure mode.

How was this patch tested?

Eight new tests. A plain fail-on-base measurement is not available: 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	maxBy(_.flatten.head.numPartitions)	fails
EnsureRequirementsSuite: the join is planned on a member pair, not on each side's first member	(leftCandidates.head, rightCandidates.head)	fails
EnsureRequirementsSuite: the finest agreeing pair wins, not the first one	agreeingPairs.headOption	fails, and also under the per-side pick
ShuffleSpecSuite: a single-partition side needs every member of a collection to be one	isCompatibleWith(specs.head)	fails
EnsureRequirementsSuite: the pairing reaches a default configuration through a transform difference	(leftCandidates.head, rightCandidates.head)	fails — and it sets no config at all
EnsureRequirementsSuite: no agreeing pair leaves the join alone however many members each side has	(leftCandidates.head, rightCandidates.head)	passes, by design — a negative pin, asserting on the shuffle rather than on the absence of grouping nodes, which a wrong fallback would not change
ShuffleSpecSuite: an empty collection is rejected where it used to answer for its head	n/a — replaces the assertion that went with the removed numPartitions	passes, by design
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 four discriminating ones 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.

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 with it. A dangling pre-split scaladoc for ShuffleSpec, stranded above ShufflePartitionIdPassThrough, is deleted rather than left to contradict the new trait doc.

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 holds — the members share their key rows and the types those rows were built with, so every one reduces to the same values, and only the key attribute is re-targeted.

Smaller things in the same direction: CoalescedHashShuffleSpec.from is typed LeafShuffleSpec, since it is structurally the one spec the coalesced spec was built from. Refining the declared return types of HashPartitioning, NullAwareHashPartitioning, KeyedPartitioning and PartitioningCollection's createShuffleSpec to their own spec types is what lets that be a type rather than a cast, and it deletes three asInstanceOf on the way, one of them in production code.

Planning cost, for the record: the per-side derivation is now linear in the number of satisfying members instead of stopping at the first, and the pairing is their cross product. Bounded in practice, since members of one collection reference different attributes and usually only one satisfies a given join's clustering.

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

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

Generated-by: Claude Code