ulysses-you commented on code in PR #59863:
URL: https://github.com/apache/spark/pull/59863#discussion_r4438941671
##########
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
Review Comment:
Updated in 5346f9f5880: the skip is limited to children whose leaves are all
shuffle query stages. Other children retain the node and its check against the
planned child partitioning. A controlled scan-rebuild test covers the guard and
shows the incorrect pairing when the node is removed.
##########
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 &&
Review Comment:
We initially changed `ConvertSortMergeJoinToShuffledHashJoin` to look
through the node. We then considered the pass-through node itself unnecessary:
it retains grouping and partition-key information in the driver plan and can be
misleading in the execution plan. This PR removes the node instead of teaching
only SHJ conversion to look through it.
The node is still constructed before the skip is decided, so this reduces
retention in the resulting plan, not construction cost. To preserve the
scan-rebuild guard, 5346f9f5880 restricts the skip to children whose leaves are
all shuffle query stages.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -2613,23 +2614,34 @@ class EnsureRequirementsSuite extends
SharedSparkSession {
// already wrapped the non-grouped side, and picked a different side
wherever one holds more
// than one split per key. Both shapes read the choice back off the
`distributePartitions`
// flags.
- def distribute(child: SparkPlan): Boolean = child.collectFirst {
- case g: GroupPartitionsExec => g
- }.get.distributePartitions
-
- def flags(smj: SparkPlan): (Boolean, Boolean) =
+ // A side picked to spread keeps its node; a side whose alignment is the
identity keeps its
+ // child and so has none, which reads as "not distributing". The node is
asserted on the side
+ // that is expected to distribute before the flags are read, so a node
that went missing
+ // cannot pass for the replicated side.
+ def flags(smj: SparkPlan, expected: (Boolean, Boolean)): Unit =
withSQLConf(
SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true") {
val Seq(newLeft, newRight) = EnsureRequirements.apply(smj).children
- (distribute(newLeft), distribute(newRight))
+ val nodes = Seq(newLeft, newRight).map(_.collectFirst {
+ case g: GroupPartitionsExec => g
+ })
+ nodes.zip(Seq(expected._1, expected._2)).foreach { case (node,
distributes) =>
+ if (distributes) {
+ assert(node.isDefined, "a side that distributes keeps its grouping
node")
+ }
+ }
+ val distributed = nodes.map(_.exists(_.distributePartitions))
Review Comment:
Updated in 5346f9f5880: these inputs do not have shuffle-stage leaves, so
both alignment nodes are retained. I restored the original assertions requiring
a node on each side before reading `distributePartitions`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -2688,9 +2702,15 @@ class EnsureRequirementsSuite extends SharedSparkSession
{
SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true") {
EnsureRequirements.apply(smj).children.map { child =>
- val gpe = child.collectFirst { case g: GroupPartitionsExec => g }.get
- (gpe.ungroupingOrigin,
- gpe.outputPartitioning.asInstanceOf[KeyedPartitioning].isGrouped)
+ child.collectFirst { case g: GroupPartitionsExec => g }.map { gpe =>
+ (gpe.ungroupingOrigin,
+ gpe.outputPartitioning.asInstanceOf[KeyedPartitioning].isGrouped)
+ }.getOrElse {
Review Comment:
Updated in 5346f9f5880: the test inputs do not have shuffle-stage leaves,
including the full outer case. I restored the original assertions requiring a
node on both sides, including for the declined left outer choice.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -726,11 +726,17 @@ class KeyGroupedPartitioningSuite
assert(collectAllShuffles(joins.head).isEmpty,
"should not add shuffle for both sides of the join")
val sides = Seq(joins.head.left,
joins.head.right).map(collectAllGroupPartitions)
- val distributed = sides.map(_.map(_.distributePartitions))
- assert(distributed == Seq(Seq(expectedLeftDistributed), Seq(false)),
+ // A side without a node distributes nothing either: its node would
have been a
+ // pass-through and is not built, and the side reports the layout it
already had.
Review Comment:
Updated in 5346f9f5880: both scan sides retain their nodes with the narrowed
skip, so I restored the original per-side grouping and partition-count
assertions for both configuration arms.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -913,7 +915,9 @@ case class EnsureRequirements(
}
}
- // Now we need to push-down the common partition information to the
`GroupPartitionsExec`s.
+ // Now we need to push-down the common partition information to the
sides that need it. A
+ // side whose alignment is the identity keeps its child instead (see
`regroupOrIdentity`
+ // below), which is common when one side was shuffled onto the other
side's layout.
Review Comment:
Updated in 5346f9f5880: the `outputPartitioning` documentation now notes
that a coalesced or skewed read over a keyed shuffle stage reports
`UnknownPartitioning` itself, allowing validation to reject it without the
alignment node.
--
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]
