ulysses-you commented on code in PR #59863:
URL: https://github.com/apache/spark/pull/59863#discussion_r4408058934
##########
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:
Fixed in 94d5b1a66c7: the skip is now gated on `applyPartialClustering`, so
a pair planned with spread/replicate slots keeps both nodes even when the
grouping is the identity. That is narrower than gating on the config: with the
config on but replication not applicable there is no choice to reproduce, and
the skip still applies.
Added `SPARK-58996: partially clustered join keeps its replicate-side choice
when the alignment is the identity`. Bucketed sides, the larger one holding one
split per bucket, statistics pick the five-split side while the partition
counts pick the three-split one, and the UNION's other branch is converted to a
shuffled hash join. Without the gate the re-run reaches `SortExec` over the
dropped side, picks the other side, and the test fails.
##########
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:
Fixed in 94d5b1a66c7: both comments and messages now say the inner join
intersects the keys onto the side missing them, which is then the pass-through
while the kept node prunes.
##########
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:
Fixed in 94d5b1a66c7, with your wording.
##########
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:
Fixed in 94d5b1a66c7: the assertions now work on the second join - one
shuffle in the plan, the second join's marked side reporting the claim over its
three keys, and the other side padded to three. Renamed the test to drop the
regrouping node it no longer builds.
##########
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:
Fixed in 94d5b1a66c7: the comment now says the intersection lands on aqekb,
and the assertions require the kept node on the first join's side and none on
aqekb's.
--
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]
