ulysses-you commented on code in PR #59863:
URL: https://github.com/apache/spark/pull/59863#discussion_r4438940180
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -824,8 +825,9 @@ case class EnsureRequirements(
// Otherwise, query result could be incorrect.
val canReplicateLeft = ShuffledJoin.canDuplicateLeftSide(joinType)
val canReplicateRight = ShuffledJoin.canDuplicateRightSide(joinType)
+ replicateSideChosen = canReplicateLeft || canReplicateRight
Review Comment:
Fixed in 5346f9f5880: `replicateSideChosen` is set in the branch that
chooses a side. It stays true if the join type subsequently declines the choice.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -1066,15 +1066,22 @@ 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, then read the node's own report, or the
dummy plan's if none.
+ def plannedPartitioning(plan: SparkPlan): Partitioning = plan match {
+ case SortExec(_, _, child, _) => plannedPartitioning(child)
+ // The node's own report is what these assertions are about. Only a
side whose alignment
+ // is the identity has no node, and there the child's layout is the
aligned one.
+ case g: GroupPartitionsExec => g.outputPartitioning
+ case d: DummySparkPlan => d.outputPartitioning
+ case other => fail(s"expected a dummy plan, got $other")
+ }
+
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
- case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _,
- GroupPartitionsExec(DummySparkPlan(_, _, left:
KeyedPartitioning, _, _),
- _, _, _, _, _, _, _, _), _),
- SortExec(_, _,
- GroupPartitionsExec(DummySparkPlan(_, _, right:
KeyedPartitioning, _, _),
- _, _, _, _, _, _, _, _), _),
- _) =>
+ case smj: SortMergeJoinExec =>
+ assert(groupPartitionsNodes(smj).size == 1, "the pruning side is
regrouped")
Review Comment:
Updated in 5346f9f5880: the skip now requires shuffle-stage leaves, so both
sides retain their nodes in these three cases. I restored the original
plan-shape assertions requiring a node on each side.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -10759,26 +10909,18 @@ 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 has to reach the second join on the side itself. Assert on
that join, since the
+ // first join's keyed shuffle already carries a marked three-key
member of its own.
+ val secondJoin = collect(plan) { case j: SortMergeJoinExec => j
}.head
+ assert(collectAllShuffles(plan).size == 1,
+ s"only the first join shuffles:\n$plan")
+ val marked =
PartitioningCollection.representativeOf(secondJoin.left.outputPartitioning)
+ .filter(_.mayContainUnknownPartitionKeys)
+ assert(marked.exists(_.numPartitions == 3),
+ s"the second join's marked side keeps its claim over its three
keys:\n$plan")
Review Comment:
Updated in 5346f9f5880: the marked input is not a shuffle-stage subtree, so
it keeps its identity grouping. The test again requires that node and checks
its claim on the second join input.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -924,15 +928,27 @@ 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))
+ // Keep both alignment nodes when a replicate side is chosen, even if
the join type
+ // declines it: a re-run peels a local sort only through a
`GroupPartitionsExec`, and
+ // that peel is how the replicate-side choice is reproduced.
+ if (!replicateSideChosen &&
+ gpe.outputPartitioning == child.outputPartitioning) {
+ child
+ } else {
+ gpe
+ }
+ }
+ (regroupOrIdentity(rawLeft, leftSpec, leftReducers, replicateLeftSide),
Review Comment:
Fixed in 5346f9f5880: the push message is logged only when the alignment is
committed and at least one node is retained. Added a test for both cases.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1158,6 +1164,61 @@ class KeyGroupedPartitioningSuite
testWithCustomersAndOrders(customers_partitions, Array.empty, 2, 0)
}
+ test("SPARK-60618: Skip the pass-through GroupPartitionsExec in a
storage-partitioned " +
+ "join's alignment") {
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ // `items` reports more partition keys than the join uses, so the
aligned layout has to be
+ // projected onto the join keys. The shuffled side is laid out on that
same projection, so
+ // its `GroupPartitionsExec` would only re-state the keys its shuffle
already has and must
+ // not be built.
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
Review Comment:
Fixed in 5346f9f5880: moved the plan-shape comment above `withSQLConf`.
--
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]
