spark-test commented on code in PR #59863:
URL: https://github.com/apache/spark/pull/59863#discussion_r4404496294
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -10530,8 +10591,15 @@ class KeyGroupedPartitioningSuite
// The shape the test exists for. Without these the assertions above
would pass on a
// plan that had lost the storage-partitioned join entirely.
val groupings = collectAllGroupPartitions(plan)
- assert(groupings.size == 2,
- s"aqe=$aqeEnabled: the second join groups both of its
sides:\n$plan")
+ // A side whose alignment is the pass-through keeps its shuffle and
gets no node, so the
Review Comment:
**Nit (P3):** The sides in this comment are reversed, and the replacement
assertions don't pin a side either.
The inner first join clears the keyed shuffle's unknown-keys marker (the
`InnerLike` arm of `ShuffledJoin.outputPartitioning`), and partition filtering
is on by default. So the inner second join intersects onto aqekb's [1, 2, 3].
The aqekb scan, which has no shuffle and is missing key 4, is the pass-through.
The one kept node sits over the first join and prunes key 4.
The new checks don't distinguish the two sides:
- `groupings.size === 1` and `expectedKeyCount.isDefined` hold for a pushed
node on either side.
- `collect(outerJoin)` walks everything below the root-most join, so it
already follows from the first check.
A plan whose only node had moved to the aqekb scan would pass every
assertion here. Yet the header says this test covers the grouping over the
first join, whose right leg is the keyed shuffle stage. Asserting that the kept
node sits over the first join's subtree, and correcting the comment, would
restore that.
**Verification:**
- **Regression:** Under both AQE settings, the test fails if no grouping
sits over the first join's subtree, which holds the keyed shuffle, for example
if the kept node moved to the aqekb scan.
- **Inspection:** The test's comment describes the plan under default
partition filtering: the inner second join intersects onto aqekb's keys, and
the kept node prunes the first join's extra key.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -1066,15 +1066,21 @@ class EnsureRequirementsSuite extends
SharedSparkSession {
// simple case
var smjExec = SortMergeJoinExec(
exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner,
None, plan1, plan2)
+ // A side whose alignment is the pass-through keeps its own layout, so
the regrouping node
+ // is optional: peel the sort and the node to reach the plan that was
planned over.
Review Comment:
**Nit (P3):** This contradicts the helper's code and its own inner comment.
The helper peels only the `SortExec`. For a `GroupPartitionsExec` it returns
`g.outputPartitioning`, the node's own report, and the inner comment says 'The
node's own report is what these assertions are about'. It never reaches the
`DummySparkPlan` under a node. This line still describes the removed pattern
match, which bound the dummy plan's partitioning beneath the node, so a reader
cannot tell which partitioning the SPARK-41413 assertions check.
```suggestion
// A side whose alignment is the pass-through keeps its own layout, so
the regrouping node
// is optional: peel the sort, then read the node's own report, or the
dummy plan's if none.
```
**Verification:**
- **Inspection:** The helper's lead and inner comments agree with each other
and with the code that returns a GroupPartitionsExec's own outputPartitioning.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -2200,8 +2257,10 @@ class KeyGroupedPartitioningSuite
val groupPartitions =
collectGroupPartitions(df.queryExecution.executedPlan)
if (pushDownValues) {
assert(shuffles.isEmpty, "should not add shuffle when partition
values mismatch")
- assert(groupPartitions.size === 2,
- "should add group partitions when partition values mismatch")
+ // Only the side that is missing keys pads; the other already holds
the merged keys, so
Review Comment:
**Nit (P3):** This comment, its twin at line 2301, and the two assertion
messages describe the opposite of the plan these tests produce. Both tests use
`createJoinTestDF`'s default inner join and leave
`spark.sql.sources.v2.bucketing.partition.filter.enabled` at its default
(true). Neither scan is marked, so `mergeAndDedupPartitions` intersects. The
merged keys are purchases' own: its two (bucket, day) keys here, and the
buckets for ids 1 and 3 in the SPARK-41413 test.
So purchases, the side missing keys, is the pass-through and gets no node.
The one node sits on items and prunes its extra key, and nothing pads. The
assertions still hold, but the text will mislead whoever next updates them.
Please reword both comments and both messages: the keys are intersected onto
purchases' keys, so purchases has no node and items' node prunes the extra key.
**Verification:**
- **Inspection:** In both partition-key mismatch tests, the comment and
assertion message describe the plan produced under default partition filtering:
the side missing keys is the pass-through, and the remaining node prunes the
other side.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -924,15 +926,19 @@ case class EnsureRequirements(
// them. Each side stamps why it is left ungrouped (`UngroupingOrigin`),
so a reader of the
// finished plan judges the pair on the layouts it holds; the stamp also
decides the node's
// routing (`GroupPartitionsExec.distributePartitions`).
- (
- GroupPartitionsExec(rawLeft, leftSpec.joinKeyPositions,
- Some(mergedPartitionKeys), leftReducers,
+ 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 (replicateLeftSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN)),
- GroupPartitionsExec(rawRight, rightSpec.joinKeyPositions,
- Some(mergedPartitionKeys), rightReducers,
- ungroupingOrigin = Option.when(applyPartialClustering)(
- if (replicateRightSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN)))
+ if (replicates) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN))
+ if (gpe.outputPartitioning == child.outputPartitioning) child else gpe
Review Comment:
**Non-blocking (P2):** Under partially clustered distribution, this skip
also removes the node that a later EnsureRequirements re-run needs to recover
the pre-alignment plan.
Take an inner sort merge join of two `bucket(4, id)` tables with partial
clustering on. S is smaller by statistics and reports splits `[b0, b0, b1, b2,
b3]`. L is larger and reports one split per bucket.
**First pass.** Statistics pick S to replicate, and the expected counts come
from L, so they are all 1. Because that output is grouped, `computeGrouping`
clears L's `SPLIT_FOR_JOIN` stamp. L's node therefore reports exactly L's
partitioning, and this line returns L. The derived scan ordering leaves out
bucket transforms, so the ordering step adds `SortExec(L)` directly over the
scan.
**Re-run.** After `ConvertSortMergeJoinToShuffledHashJoin` converts another
join in the same AQE round, it re-applies the rule to the whole plan. The
SPARK-58996 union tests drive exactly this; `OptimizeSkewedJoin` does the same
under force-apply. Then:
- `keepArrivedPairing` is off under partial clustering.
- `peelGroupPartitions` strips a local sort only when a
`GroupPartitionsExec` sits under it, so `SortExec(L)` stays as the raw side.
- That sort has no logical link, so the side choice falls back to partition
counts. 5 < 4 is false, so L is replicated and S is spread.
L's `b0` is now read once per S split, and AQE adopts the equal-cost but
different plan.
Rows stay correct, because the counts still come from the raw keys. What
breaks is that the statistics-based choice no longer survives the re-run, which
is the property 'SPARK-58996: partially clustered join keeps its replicate-side
choice when EnsureRequirements re-runs' pins. That test misses this case
because both of its sides keep nodes. Before this change, L kept its identity
node, both sides peeled back to the scans, and the re-run repeated the
statistics-based choice. With partial clustering disabled the skip is safe:
`keepArrivedPairing`, or deterministic re-planning from the same reported
partitionings, reproduces the pairing.
**Recommended change:** Skip the pass-through GroupPartitionsExec only while
spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled is off,
and keep building it while that config is on. Add a storage-partitioned join
test under AQE, in the shape of the existing partially clustered re-run join
tests, where the larger side has one split per bucket and needs a sort.
**Why this works:** With the node kept, the local sort the ordering step
adds sits over a GroupPartitionsExec. A re-run then strips both, reaches the
scans with their logical links, and repeats the first pass's statistics-based
replicate-side choice. Outside partial clustering the skip is unchanged,
including the AQE one-side keyed shuffle shape the change targets, because
keyed shuffles are never created while partial clustering is enabled.
**Scope:** Keep the alignment node while partially clustered distribution is
enabled, so that EnsureRequirements re-runs reproduce the first pass's
replicate-side choice.
**Compatibility:** While partially clustered distribution is disabled, the
pass-through skip, the final pairing gate and every reported partitioning stay
exactly as this change makes them.
**Risks:** While partially clustered distribution is enabled, a pass-through
side keeps its 1:1 GroupPartitionsExec as in the merge target, so EXPLAIN still
shows that node there.
**Constraints:** The skip must stay in effect while partial clustering is
disabled, so that the AQE one-side keyed shuffle join still reaches the join
without a node over the shuffle stage.
**Exclusions:**
- Changing ConvertSortMergeJoinToShuffledHashJoin, peelGroupPartitions, or
the statistics and partition-count heuristic that picks the replicate side.
**Success:** While partially clustered distribution is enabled, every side
of a pushed storage-partitioned join alignment keeps its GroupPartitionsExec,
and an EnsureRequirements re-run over the planned tree reproduces the first
pass's partially clustered decision. While partially clustered distribution is
disabled, a side whose node would report exactly its child's partitioning still
gets no node.
**Verification:**
- **Regression:** A partially clustered storage-partitioned sort merge join
over bucketed tables, planned in an AQE round where another join in the same
plan is converted to a shuffled hash join, keeps the replicate-side choice that
statistics made before the re-run. The smaller side by statistics holds several
splits for some bucket, and the larger side holds one split per bucket. The
test also checks the answers.
- **Behavior:** With partial clustering disabled, the AQE one-side keyed
shuffle join still reaches the join with no GroupPartitionsExec directly over
its shuffle stage.
- **Compatibility:** Existing partially clustered storage-partitioned joins
keep their split and replicate stamps, partition counts and answers with
partial clustering enabled, and keep the pass-through skip with it disabled.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -10759,26 +10827,14 @@ class KeyGroupedPartitioningSuite
val plan = df.queryExecution.executedPlan
assert(ValidateRequirements.validate(plan),
s"the executed plan must satisfy every operator's required
distribution:\n$plan")
- // Locate the marked side's node: no reducers, and output partition
i holds exactly input
- // partition i. That is weaker than `identityGrouping`, which also
asks for one output per
- // input, so the assertions below are what pin the claim and the
three keys.
- val identityGpe = collectAllGroupPartitions(plan).find { g =>
- g.reducers.isEmpty && g.groupedPartitions.zipWithIndex.forall {
- case ((_, inputIndices), outputIndex) =>
- inputIndices.lengthCompare(1) == 0 && inputIndices.head ==
outputIndex
- }
- }
- assert(identityGpe.isDefined,
- s"expected a reducer-free identity GroupPartitionsExec,
got:\n$plan")
- identityGpe.get.outputPartitioning match {
- case k: KeyedPartitioning =>
- assert(k.mayContainUnknownPartitionKeys,
- "the identity grouping must keep the unknown-keyed claim")
- assert(k.numPartitions === 3,
- s"the marked side keeps all three of its declared
keys:\n$plan")
- case other =>
- fail(s"expected a KeyedPartitioning output, got $other")
- }
+ // The marked side's alignment is the pass-through, so no node is
built for it. The
+ // claim it would have re-stated has to reach the plan on the side
itself, still over
+ // its three declared keys.
+ val markedKeyed = collect(plan) { case p: SparkPlan =>
p.outputPartitioning }
+ .flatMap(PartitioningCollection.representativeOf)
+ .filter(_.mayContainUnknownPartitionKeys)
+ assert(markedKeyed.exists(_.numPartitions == 3),
Review Comment:
**Non-blocking (P2):** This check passes without looking at the second join.
`collect(plan)` reads every node, and the first join's one-side keyed shuffle
of `pt` is always stamped `HASH_ROUTING` over pa's three keys by
`KeyedShuffleSpec.createPartitioning`. `RIGHT OUTER` forwards
`right.outputPartitioning`, so that exchange and the first join already supply
a marked three-partition member.
Now suppose the second join declined its pairing and shuffled both inputs.
That could come from a change that let partition filtering intersect a marked
pair down to [1, 2], the case the SPARK-59272 note above guards against, or
from the marked input losing its claim. `checkAnswer`, `ValidateRequirements`
and this check would all still pass, and with AQE off nothing counts shuffles.
The removed check failed in both cases.
Asserting on the second join itself restores the signal: no shuffle beyond
the first join's one, the marked input reporting the claim over three keys, and
pb padded to three. The test name still describes 'a reducer-free identity
regrouping', which this change no longer builds, so it needs renaming too.
**Verification:**
- **Regression:** On both partition-filter arms, the test fails if the
second join over the marked RIGHT OUTER join output is shuffled instead of
storage-partitioned, or if the second join's marked input no longer reports the
unknown-keyed claim over its three declared keys.
--
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]
