peter-toth opened a new pull request, #58659:
URL: https://github.com/apache/spark/pull/58659
### What changes were proposed in this pull request?
**Stacked on apache#58531 (SPARK-59256) and apache#58552 (SPARK-59285).**
GitHub has no stacked pull
requests, so review only the last commit and do not merge this before those
two.
`GroupPartitionsExec` is how storage-partitioned join lines two scans up: it
coalesces, projects,
reorders and pads a scan's partitions. Today the planner decides that node
twice and the node then
re-derives it many times. This PR makes the planner decide once and the node
carry the answer.
**Where the node comes from.** `ensureDistributionAndOrdering` resolved
every child on its own first,
so it placed a `GroupPartitionsExec` over a co-partitioned child before
knowing the parent was a
join. `checkKeyGroupCompatible` then rewrote or replaced it, descending past
a local sort, merging
new parameters in and dropping any grouping stacked on top. The two sites
also worked in different
index spaces, the placeholder's positions against the raw partition keys and
the spec's against the
node's already projected report, which is what `joinKeyPositions.orElse` and
the reducer fast path
existed to paper over.
The two paths are now apart. The per-child loop resolves only the children
that answer for
themselves, through a lifted `resolveChild`. A new `coPartitionChildren`
owns the rest end to end:
it peels every grouping this rule inserted off both children, plans the
pairing from the sources'
own reports, and builds each node once. `ensureDistributionAndOrdering` is
72 lines.
**What the node holds.** `grouping` (the index groups it emits and what they
say about the layout)
and `outputPartitioning` are constructor fields now, derived once by
`GroupPartitionsExec.apply`,
the way `ShuffleExchangeExec` holds the partitioning it produces. They were
per-instance lazy vals,
so every `copy` and `withNewChildren` threw the memo away.
**What `satisfies` answers.** `KeyedPartitioning.satisfies` returned `true`
for a partitioning that
only satisfies after a projecting node, for the storage-partitioned join's
benefit alone. It is
strict now, and the loose question moved to `keysCanSatisfy`, which
`keysMaySatisfy` and
`mayGroupToSatisfy` compose with the strict one.
`EnsureRequirements.clusterKeyPositions` delegates
to a new `KeyedPartitioning.operationKeyPositions`, so the predicate that
decides whether a
projection is needed and the one that picks the positions cannot drift apart.
**What a second run sees.** Peeling and re-planning would re-decide a
pairing an earlier run already
settled, on an input that by then holds the keyed shuffle that run inserted.
So
`checkKeyGroupCompatible` asks first whether the two children satisfy their
distributions and line
up with each other as they arrive, and hands them back untouched when they
do. That question decides
nothing new on a plan this rule has not seen, since a source that needs a
node does not satisfy its
distribution as it stands.
EnsureRequirements
ensureDistributionAndOrdering the two paths, and the ordering
step
resolveChild a child that answers for itself
coPartitionChildren children that only answer together
checkKeyGroupCompatible try the join, build both nodes
once
alreadyCoPartitioned unless a previous run settled it
resolveCoPartitionedChildren else each on its own
(the shuffle onto the best spec)
GroupPartitionsExec(child, grouping, outputPartitioning, <recipe>,
enableSortedMerge)
GroupPartitionsExec.apply the only way to build one
computeGrouping pure, takes the child's
KeyedPartitioning
computeOutputPartitioning pure, takes the same
Retires `rewriteGroupPartitions`, `applyGroupPartitions` and
`innermostGroupPartition`'s
sort-rebuilding half, along with the `orElse` and the two index spaces they
served.
### Why are the changes needed?
**A join could commit to a pairing its own children then refuse.**
`GroupPartitionsExec` gives up
its keyed claim when it turns out to regroup a layout that pins undeclared
rows to
`hash(key) % numPartitions`, and only the node knows the permutation it
performs, so that answer
arrived after `checkKeyGroupCompatible` had skipped both shuffles. The
result is a plan
`ValidateRequirements` rejects, and every `AQEShuffleReadRule` and
`OptimizeSkewedJoin` drops its
result on such a stage, so partition coalescing, local read and skew join
are all off for it. Now
the site that builds the two nodes asks them before returning. This closes
SPARK-59272.
**The same derivation ran many times.** Measured on
`KeyGroupedPartitioningSuite`'s "join with two
partition keys and matching & sorted partitions": the base derives the
grouping 16 times over 48
node instances, this derives it 8 times over 64. A lazy val belongs to one
instance, and the
columnar and codegen rules rebuild the tree after this rule, so the
derivations followed the
instances. As a field they follow the planner's decisions instead.
**The caller was compensating for the predicate.**
`resolveKeyedPartitioning` asked in two steps
whether a member needed a node at all, the cheap full-coverage test and then
the projection that
decides whether a narrowing merges anything. `keysSatisfy` owns both now and
the caller is
`admitted.find(_._2)`. `ValidateRequirements` also becomes a real guard: on
a `KeyedPartitioning([a,
b])` with two partitions sharing `a = 1` under a
`ClusteredDistribution([a])`, `validate` was `true`
before and is `false` now, and the plan spreads rows sharing the operation
key either way.
**The rule was not idempotent.** Applying it to its own output changed the
plan, which matters
because AQE re-plans every query stage and two AQE rules hand the tree back
to it. Measured over
`KeyGroupedPartitioningSuite` by re-running the rule on each of its own
results, 2713 applications
in all: the base differs on 2, and it differs by adding a shuffle. Without
the check described
above this differs on 16, all of them a grouping node added over a shuffle.
With it, on none.
### Does this PR introduce _any_ user-facing change?
**Yes**, three plan changes. No API is added or removed outside `catalyst`,
which is in
`MimaExcludes`' `defaultExcludes` section and treated as internals.
**A join declines rather than leaving an unvalidatable plan.** On the
`SPARK-59050: SPJ: regrouping a
marked layout must not keep the unknown-keyed claim` query the second join
now takes a keyed
one-side shuffle, three shuffles to four, and
`ValidateRequirements.validate` on the executed plan
goes from `false` to `true`. The alternative was not "no shuffle", it was
"no shuffle and no AQE".
**`satisfies` accepts a partitioning whose expression *is* an operation
key**, such as a `years(ts)`
under a clustering naming `years(ts)`. The old reference-level test refused
it unless
`requireAllClusterKeys` was set, while the `requireAllClusterKeys` arm
accepted the same shape, so
this removes an inconsistency. It is a widening on the default configuration.
**A re-planned stage keeps the plan it had.** The rule is idempotent now,
where the base can add a
shuffle on the second pass, so an AQE stage's plan no longer drifts from the
one first planned.
### How was this patch tested?
**Nine new tests, each verified to fail with its own change reverted,
measured against this base
rather than a remembered one.** The `-` column is what was reinstated to
check that the test
measures something.
| test | reinstated | on the reinstated code |
|---|---|---|
| `DistributionSuite`: satisfies is strict about a projection that merges
partitions | the loose `keysSatisfy` | fails |
| `ShuffleSpecSuite`: a collection whose members all need grouping still
yields them | the collection filter on strict `satisfies` | fails |
| `EnsureRequirementsSuite`: SPARK-58996 only a local sort is looked through
| descending through a global sort | fails |
| a local sort with no grouping under it is left alone | peeling a sort with
nothing under it | fails |
| a grouped side is paired on its own key order | `toGrouped` for a grouped
source too | fails |
| pushing join key positions into a node re-derives its grouping | `copy`
instead of the rebuild | fails |
| single-partition children still honour a required partition count |
leaving both children alone | fails |
| a partitioning is never projected onto no position | projecting onto the
empty position set | fails |
| a second pass leaves a pairing this rule already made alone | re-deciding
a settled pairing | fails |
One existing expectation moved, `Seq(true, true, true)` to `Seq(true, true,
true, true)` in the
SPARK-59050 regrouping test, and it now asserts
`ValidateRequirements.validate` ahead of the shuffle
count because that names the reason the shuffle exists.
**A differential sweep is the strongest evidence here.** 307200 generated
storage-partitioned join
plans, 40 partitioning shapes a side crossed with six join types over 32
configuration cells, run
against the base and against this change. Each plan is checked twice, that
`ValidateRequirements`
passes and that two sides whose shuffles were both skipped really declare
the same key sequence.
The cases are keyed and set-differenced, not just counted.
| | base | this change |
|---|---|---|
| plans that skip both shuffles yet disagree on the keys | 2024 | **0** |
| plans `ValidateRequirements` rejects | 6724 | 5296 |
| plans the planner cannot build at all | 29952 | 29356 |
**Nothing is worse on any of the three, and the sets line up exactly.** The
2024 co-partitioning
violations are the same 2024 cases as the removed validation failures, which
is SPARK-59272 in one
number. The 596 fewer planner failures are the same 596 cases that now
appear among the rejected
plans, so a shape the base could not plan at all is now planned, badly. Zero
cases are added to any
of the three.
The remaining 29356 failures are shared with the base and are the
generator's, not the planner's.
They come from `Partitioning.createShuffleSpec`, which throws when a
co-partitioned child reports
`UnknownPartitioning`, and that needs a marked layout that is not grouped.
Nothing builds one. The
marker is only ever put on a layout `KeyedShuffleSpec.createPartitioning`
has just laid out one
partition per key, and both paths that would regroup a marked layout refuse
it.
Green on this base: `DistributionSuite`, `ShuffleSpecSuite`,
`EnsureRequirementsSuite`,
`ValidateRequirementsSuite`, `PlannerSuite`, `GroupPartitionsExecSuite`,
`ProjectedOrderingAndPartitioningSuite`, `KeyGroupedPartitioningSuite`,
`DataSourceV2CatalystRuntimeFilterSuite`, `AdaptiveQueryExecSuite`,
`ExchangeSuite` and
`DataSourceV2Suite`, 753 tests. `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]
