peter-toth opened a new pull request, #58531:
URL: https://github.com/apache/spark/pull/58531
### What changes were proposed in this pull request?
**Stacked on #58527 (SPARK-59080), which is the first commit here. Please
review only the second commit, and do not merge this before #58527. Draft until
then.**
`ShuffleSpecCollection` answers two kinds of question with one type.
Matching (`isCompatibleWith`) is a collection question, and succeeding when
*any* member matches is why the type exists. The shuffle template
(`createPartitioning`, `numPartitions`) needs one member, and the collection
reads `specs.head`, whichever member the alias cross-product enumerated first.
A new `LeafShuffleSpec` sub-trait carries the two single-member methods, the
seven concrete specs extend it, and the collection extends `ShuffleSpec` alone:
```
ShuffleSpec isCompatibleWith, canCreatePartitioning
<- LeafShuffleSpec + numPartitions, createPartitioning
SinglePartitionShuffleSpec, RangeShuffleSpec, HashShuffleSpec,
NullAwareHashShuffleSpec, CoalescedHashShuffleSpec, KeyedShuffleSpec,
ShufflePartitionIdPassThroughSpec
<- ShuffleSpecCollection
```
`ShuffleSpec` is sealed, so those two are the only kinds. That is what lets
the flattening helper return `Seq[LeafShuffleSpec]` without a fallback case,
and it is the only thing giving the split teeth. Nothing outside
`partitioning.scala` extends `ShuffleSpec`. The helper moves out of
`EnsureRequirements` to `ShuffleSpecCollection.flatten`, beside the sealed
hierarchy that makes its match total and alongside the existing
`PartitioningCollection.flatten`.
`ShuffleSpecCollection` also takes over the `require(specs.nonEmpty, ...)`
that `numPartitions` used to carry, so the invariant outlives the method that
stated it.
### Why are the changes needed?
The two removed methods are not wrong in the same way.
`createPartitioning` has **no local answer at all**: the right member is the
one the *other* side matched, which is only visible to the caller comparing the
two sides. SPARK-59080 moved that decision into `EnsureRequirements`, so after
#58527 the method has no production caller, and only a runtime `require` stands
between a future caller and a wrong partitioning. The compiler should say that
instead.
`numPartitions` has **two consumers that want two different aggregations**,
which is why it should not be answered at all rather than answered as one of
them:
- `EnsureRequirements` ranks the children by parallelism to pick the
reference layout. It wants the **max**, and it now takes it over the flattened
members. That also makes the ranking agree with the member the shuffle is then
built from, which is already the finest one. Ranking on the head while building
from the max was incoherent.
- `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.
### Does this PR introduce _any_ user-facing change?
No.
Both behaviour changes are exact no-ops unless
`spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` is on,
which it is not by default. `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`. The only producer of a disagreeing member is the
`project(joinKeyPositions).toGrouped` branch that config gates.
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: the shape needs a collection
child, a `SinglePartition` sibling and an AQ
E shuffle-read rule firing on that pair, and those rules produce
`CoalescedHashPartitioning`, never `SinglePartition`. I kept `exists` because
that is how compatibility with a collection is defined everywhere else, and I
did not add a test for it, because pinning an answer I am not sure is right
would be worse than leaving it uncovered.
### How was this patch tested?
Three new tests, and one existing assertion deleted.
A plain fail-on-base measurement is not possible here: the new tests name
`LeafShuffleSpec`, so they cannot compile against the parent commit. Instead I
simulated the old head read under the new types, making the ranking read
`flatten(spec).head` and `SinglePartitionShuffleSpec`'s collection case read
`isCompatibleWith(specs.head)`. Both behavioural tests fail under that and pass
without it.
| test | under the simulated head read |
|---|---|
| `ShuffleSpecSuite`: a single-partition side matches a collection through
any member | fails |
| `EnsureRequirementsSuite`: a collection is ranked on its best member, not
on whichever came first | fails |
| `ShuffleSpecSuite`: flattening reaches the members of a nested collection
| n/a, covers the recursion |
Both behavioural tests 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.
Green: `ShuffleSpecSuite`, `DistributionSuite`, `EnsureRequirementsSuite`,
`ValidateRequirementsSuite`, `KeyGroupedPartitioningSuite`, `PlannerSuite`, 306
tests in all. `dev/lint-scala` is clean.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
