dongjoon-hyun commented on code in PR #59863:
URL: https://github.com/apache/spark/pull/59863#discussion_r4425921038
##########
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:
**[1] Returning the bare child also drops the node's guard against a child
rebuilt after `EnsureRequirements`.**
A `GroupPartitionsExec` reports `UnknownPartitioning` once its child no
longer reports the partitioning it was planned over, and
`checkChildStillMatches` then fails the query at execution
(GroupPartitionsExec.scala:137-142 and 353-358). Its doc names the rules that
rebuild such a child: `PlanAdaptiveDynamicPruningFilters` and
`ReuseAdaptiveSubquery` with AQE, and `ReuseExchangeAndSubquery` without it,
which rebuild a `BatchScanExec` when they rewrite a subquery in its runtime or
storage filters (GroupPartitionsExec.scala:126-131). The rebuilt scan plans its
input partitions again, since `batch` and `inputPartitions` are lazy vals
(BatchScanExec.scala:70 and 111).
With the identity node, a different split set from the connector on that
second call fails with an internal error. Without it, the join zips the rebuilt
scan with the other side's node, which was aligned to the old keys. That fails
with "Can't zip RDDs with unequal numbers of partitions" when the counts
differ, and silently pairs the wrong partitions when they match.
As that doc says, no rule is known to change the splits today, so this is
defense in depth. Could the comment at L918-920 say that a skipped side gives
up this guard, or could the node stay, as in [2]?
##########
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:
**[5] `checkPartiallyClusteredJoin` accepts a missing node in both arms, but
with partial clustering on, both sides always keep one.**
With the config on, `pushCommonValues` is always true
(EnsureRequirements.scala:732-735), and both join types here may duplicate the
right side (`ShuffledJoin.canDuplicateRightSide`), so a replicate side is
always chosen and EnsureRequirements.scala:943 keeps both nodes. That includes
the two calls where the left side is picked and the join type declines it. So
the comment at L729-730 holds only for the `partiallyClustered = false` arm. In
the other arm, a skipped identity node would still pass: it reads as not
distributing, and the partition count falls back to the side's own (L735-739),
which is the same for an identity node. Could that arm require both nodes?
```suggestion
// With partial clustering on, both join types here may duplicate
the right side, so a
// replicate side is chosen and both sides keep a node. With it off,
a side without a node
// distributes nothing: its node would have been a pass-through and
is not built.
if (partiallyClustered) {
assert(sides.forall(_.nonEmpty), s"both sides keep a grouping
node:\n$plan")
}
```
##########
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:
**[3] This accepts a missing node on the side that is expected not to
distribute, although every case in this test now keeps both nodes.**
All four cases are inner joins with partial clustering on. An inner join may
duplicate either side, so a replicate side is always chosen, and
EnsureRequirements.scala:943 keeps both nodes, an identity one included. So the
comment at L2617-2620, which dates from the first commit, no longer holds here.
Because of the `exists` at L2634, the test also does not pin that gate. In the
last case (left `[1, 2]`, right `[1, 2, 3]`), the inner join intersects the
keys onto `[1, 2]`, so the replicated left node is the identity. Without the
`replicateSideChosen` term, as in the first commit, that node would be skipped,
read as "not distributing", and match the expected `false`.
Master's `distribute` read the node with `.get`. Could both nodes be
required again, and the comment at L2617-2620 say that every case here chooses
a replicate side?
```suggestion
assert(nodes.forall(_.isDefined), "a chosen replicate side keeps
both grouping nodes")
val distributed = nodes.map(_.get.distributePartitions)
```
##########
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:
**[6] With this, the keyed-shuffle side often has no `GroupPartitionsExec`,
so the `outputPartitioning` doc in GroupPartitionsExec.scala:116-119 no longer
names what protects it.**
That doc says that an `AQEShuffleReadExec` landing over a keyed shuffle
stage relies on the node's fallback to `UnknownPartitioning`, which makes
`ValidateRequirements` reject the rewrite. Without the node, the rejection
comes from `AQEShuffleReadExec`'s own `UnknownPartitioning` for a coalesced or
skew-split read of a keyed layout (AQEShuffleReadExec.scala:102-110).
GroupPartitionsExec.scala is outside this diff, so I left this here. Nit: could
that paragraph say that a side without a node relies on `AQEShuffleReadExec`
instead?
##########
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:
**[4] With the `getOrElse` fallback, the mirrored LEFT OUTER case passes
whether its identity node is kept or skipped, so this test does not pin
374a7c3's keep for a declined choice.**
In `smj(LeftOuter, Seq(InternalRow(1)), threeSplits)` at L2734, the counts
pick the left side, which a left outer join may not duplicate. Since 374a7c3,
that declined choice keeps both nodes. With 94d5b1a's gate on
`applyPartialClustering`, the left identity node would be skipped, and this
fallback would return `(None, true)`, the same pair as the kept node, so the
assertion holds either way. Only the FULL OUTER case at L2738 needs the
fallback: neither side may be duplicated there, so no replicate side is chosen,
and the identity right side has no node.
Could the fallback apply to that case only, so that the other three calls
require a node on both sides? For example, `stampsOf(smj: SparkPlan,
nodeOnBothSides: Boolean = true)` with `assert(!nodeOnBothSides, ...)` at the
top of the `getOrElse` branch, and `nodeOnBothSides = false` on the FULL OUTER
call.
##########
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:
**[2] Could the shuffled-hash-join conversion look through an identity
`GroupPartitionsExec`, instead of the alignment not building one?**
The conversion this PR enables is blocked in
`ConvertSortMergeJoinToShuffledHashJoin.findShuffleStage`
(ConvertSortMergeJoinToShuffledHashJoin.scala:187-210), which looks through the
join's own `SortExec`, and through a few row-wise operators behind a config,
but not through a `GroupPartitionsExec`. Removing the node where it is built
has costs elsewhere:
- A re-run peels a local sort only through a `GroupPartitionsExec`
(L940-942). That is what `replicateSideChosen` works around, and what
[r4404496222](https://github.com/apache/spark/pull/59863#discussion_r4404496222)
and
[r4411054861](https://github.com/apache/spark/pull/59863#discussion_r4411054861)
went through.
- The plan changes against 4.2.0, which builds the node on both sides
(`EnsureRequirements.scala:683-687` at `v4.2.0`).
- The skipped side loses the guard in [1] and the node's metrics.
A look-through next to the `SortExec` case
(ConvertSortMergeJoinToShuffledHashJoin.scala:193) would keep the plan as it is:
```scala
// An identity regrouping reads each stage partition exactly once, so
the stage's
// statistics still bound every build partition.
case g: GroupPartitionsExec if g.grouping.isIdentity =>
findShuffleStage(g.child, lookThroughOperatorsEnabled)
```
`PartitionGrouping.isIdentity` (GroupPartitionsExec.scala:837-841) holds
only when output partition i is input partition i and no key was rewritten,
which rules out coalescing, padding, pruning and replication, so `mapStats`
stays a per-partition bound. `stripSort`
(ConvertSortMergeJoinToShuffledHashJoin.scala:163-166) then leaves the node
under the shuffled hash join, and the `EnsureRequirements` re-run after the
conversion (ConvertSortMergeJoinToShuffledHashJoin.scala:155) finds the node to
peel. Could you consider this, or say in the description why removing the node
is preferred?
--
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]
