ulysses-you opened a new pull request, #59863:
URL: https://github.com/apache/spark/pull/59863
### What changes were proposed in this pull request?
`EnsureRequirements.checkKeyGroupCompatible` pushes the merged partition
keys to both sides of a
co-partitioned join by wrapping each side in a `GroupPartitionsExec`. When
one side was shuffled
onto the other side's keyed layout, it already holds those merged keys, so
the node it would get
only re-states the partitioning its child already reports: it maps every
partition onto itself,
coalesces nothing, and reports the same keys, marker and routing.
This PR skips that node. The two sides are built through a new local
`regroupOrIdentity`, which
returns the child itself when the `GroupPartitionsExec` it would build
reports exactly the
child's partitioning:
def regroupOrIdentity(
child: SparkPlan,
spec: KeyedShuffleSpec,
reducers: Option[Seq[Option[KeyReducer]]],
replicates: Boolean): SparkPlan = {
val gpe = GroupPartitionsExec(child, spec.joinKeyPositions,
Some(mergedPartitionKeys), reducers,
ungroupingOrigin = Option.when(applyPartialClustering)(
if (replicates) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN))
if (gpe.outputPartitioning == child.outputPartitioning) child else gpe
}
Comparing the reported partitioning is what makes this one question.
Equality covers the
expressions, the key rows and their order, the data types,
`isGrouped`/`isCollapsed`,
`unknownKeysOrigin`, and the layout's `ungroupingOrigin`, so it also pins
the mapping: the keys are
the child's in the child's order, and each output partition can only hold
the input partition
named for the same key. A node that coalesces several splits per key, pads
missing keys, prunes,
rewrites keys through reducers, spreads or replicates a side, or gives up
the keyed claim or its
routing, reports something else and is still built.
Tests: the assertions that required the node on *both* sides now read the
same fact from the
side's own partitioning (or accept that the pass-through side keeps the
layout it has). No
assertion was weakened; see "How was this patch tested?".
### Why are the changes needed?
The alignment wraps both sides even when one of them already holds the
merged keys, and that node
is a pass-through. A side shuffled onto the other side's layout is exactly
such a case, since
`KeyedShuffleSpec.createPartitioning` lays it out on the same projection the
alignment is computed
from. Downstream rules then have to look through a node that carries no
information, and
`ConvertSortMergeJoinToShuffledHashJoin` does not: it stops at the node, so
a sort merge join that
could be a shuffled hash join keeps its child sorts.
The shuffled side of a storage-partitioned `SortMergeJoin` reaches the join
as
Sort
+- GroupPartitions (ExpectedPartitionKeys: 3)
+- ShuffleQueryStage
+- Exchange keyedpartitioning(item_id, ...)
where the node maps every partition onto itself. A real `MERGE INTO` whose
join was a `FullOuter`
`SortMergeJoin` over a storage-partitioned target and a shuffled source had
the same shape
(`GroupPartitions (JoinKeyPositions: [0])` over the stage); the join stayed
a sort merge join and
the target side's four-key sort spilled 10.1 TiB. With this change the join
converts to a shuffled
hash join and that sort disappears.
### Does this PR introduce _any_ user-facing change?
Yes, in the sense that `EXPLAIN` output changes: a storage-partitioned join
whose alignment is the
identity no longer shows a `GroupPartitions` node on the side that already
holds the merged keys.
Query results, configurations and APIs are unchanged. Both the alignment and
its configurations
are unreleased on master, so this is not a change against a released version.
### How was this patch tested?
New test `SPARK-60618: Skip the pass-through GroupPartitionsExec in a
storage-partitioned join's
alignment` in `KeyGroupedPartitioningSuite`: `items` is partitioned by `id`
and `name`, the join
uses only `id`, and `purchases` is unpartitioned, so the shuffled side is
laid out on the same
projection the alignment is computed from. It checks that the one-side keyed
shuffle still happens,
that no `GroupPartitionsExec` sits over the shuffle stage (or over the
`AQEShuffleReadExec` above
it), that the source side keeps its regrouping node, and the query's
answers. Without the change
the test finds the node over the stage and fails.
Existing suites, all passing:
build/sbt 'sql/testOnly
org.apache.spark.sql.connector.KeyGroupedPartitioningSuite' # 240
build/sbt 'sql/testOnly
org.apache.spark.sql.execution.exchange.EnsureRequirementsSuite' # 81
build/sbt 'sql/testOnly
org.apache.spark.sql.execution.datasources.v2.GroupPartitionsExecSuite' # 38
build/sbt 'sql/testOnly
org.apache.spark.sql.execution.adaptive.AdaptiveQueryExecSuite' # 146
`KeyGroupedPartitioningSuite` runs many shapes under both AQE arms.
Assertions that required the
node on both sides were updated as follows, each keeping what it verified:
- `SPARK-59688`: only the reduced side keeps its node; the reducer count and
`ValidateRequirements` are unchanged.
- partitioned join with two keys / `SPARK-41413` (subset): the side that is
missing keys pads,
asserted as exactly one regrouped side.
- `SPARK-59294`: a side without a node is not distributing; partition counts
fall back to the join
child.
- `SPARK-59289`: the second join groups exactly one side, the pushed
alignment node, under the
second join, and the keyed shuffle is still there.
- `SPARK-59050`: the marked side's unknown-keyed claim is asserted on the
plan's side itself,
still over its three keys.
- `EnsureRequirementsSuite` (`SPARK-41413`, `SPARK-59256`, `SPARK-59272`,
`SPARK-59671`,
`SPARK-58996`): read the same facts from the node when it is there (or
from the side's own
partitioning when the alignment is the identity), and assert that a side
expected to distribute
keeps its node, so a missing node cannot pass for the replicated side.
Ordering-related cases were re-run on their own as well: the
`tryEnableSortedMerge` tests
(`SPARK-56549`, `SPARK-55715`, `SPARK-58996`, `SPARK-59279`, `SPARK-60044`)
all pass. A pass-through
node coalesces nothing, so it was never a k-way-merge carrier
(`GroupPartitionsExec.mergeIsPlanned`
requires coalescing).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (deepseek-flash[1m])
🤖 Generated with [Claude Code](https://claude.com/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]
