peter-toth commented on code in PR #58279:
URL: https://github.com/apache/spark/pull/58279#discussion_r3924540478
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -657,23 +660,30 @@ case class EnsureRequirements(
log"distribution.")
replicateRightSide = false
} else {
- // In partially clustered distribution, we should use un-grouped
partition values
- val (partiallyClusteredChild, partiallyClusteredSpec) = if
(replicateLeftSide) {
- (unwrappedRight, rightSpec)
- } else {
- (unwrappedLeft, leftSpec)
- }
- // Original `KeyedPartitioning` can be obtained from the child
directly if the child
- // satisfied the distribution requirement; or from the child's
child if it didn't as
- // the child must be a `GroupPartitionsExec` inserted by
`EnsureRequirement`
- // to satisfy the distribution requirement.
+ // In partially clustered distribution, we should use un-grouped
partition values.
+ // The positions projecting them come from the innermost
grouping when there is
+ // one: like in `applyGroupPartitions`, they were computed
against the raw
+ // partition keys, while the spec's were computed against the
node's already
+ // projected report on a re-run.
+ val (partiallyClusteredChild, partiallyClusteredPositions) =
Review Comment:
**Finding 9.** `unwrapGroupPartitions(right)` at `:618` and
`innermostGroupPartition(right)` here run the same descent and land on the same
node. That is what makes the pair sound. The keys come from that node's child,
the positions from that node itself. The code derives each one separately and
pairs them by hand, so the pairing holds because two call sites agree rather
than because it is one value.
That is the shape findings 3 and 4 came out of. The plan moved to the
pre-alignment node, the positions stayed on the spec, and nothing failed to
compile. One value removes the class. At `:617-618`:
```scala
val leftGrouping = innermostGroupPartition(left)
val rightGrouping = innermostGroupPartition(right)
val unwrappedLeft = leftGrouping.map(_._1.child).getOrElse(left)
val unwrappedRight =
rightGrouping.map(_._1.child).getOrElse(right)
```
and here:
```scala
val (partiallyClusteredChild, partiallyClusteredPositions) =
if (replicateLeftSide) {
(unwrappedRight,
rightGrouping.flatMap(_._1.joinKeyPositions).orElse(rightSpec.joinKeyPositions))
} else {
(unwrappedLeft,
leftGrouping.flatMap(_._1.joinKeyPositions).orElse(leftSpec.joinKeyPositions))
}
```
The descent then runs twice per join instead of three times.
`unwrapGroupPartitions` is left with its one remaining caller, the
`ShuffleExchangeExec` site at `:288`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -657,23 +660,30 @@ case class EnsureRequirements(
log"distribution.")
replicateRightSide = false
} else {
- // In partially clustered distribution, we should use un-grouped
partition values
- val (partiallyClusteredChild, partiallyClusteredSpec) = if
(replicateLeftSide) {
- (unwrappedRight, rightSpec)
- } else {
- (unwrappedLeft, leftSpec)
- }
- // Original `KeyedPartitioning` can be obtained from the child
directly if the child
- // satisfied the distribution requirement; or from the child's
child if it didn't as
- // the child must be a `GroupPartitionsExec` inserted by
`EnsureRequirement`
- // to satisfy the distribution requirement.
+ // In partially clustered distribution, we should use un-grouped
partition values.
Review Comment:
**Finding 10.** This is about `:644`, which sits just above the hunk, so
this is the nearest line I can anchor on.
The statistics read at `:623-645` now comes from the pre-alignment plan. The
branch it falls back to does not. `leftPartKeys` and `rightPartKeys` at
`:551-552` are `leftSpec`/`rightSpec`'s keys, and those specs are built at
`:167` from the aligned node's report. Both sides align to the same
`mergedPartitionKeys` with the same per-key counts, so on a re-run the two
lists have equal length and `leftPartKeys.size < rightPartKeys.size` is `false`
whatever pass 1 chose. That is the flip
[r3864640486](https://github.com/apache/spark/pull/58279#discussion_r3864640486)
describes. The fix reached the trigger and left the flip.
I could not reach it on this head and the argument here is from reading.
After the fix `unwrapGroupPartitions` lands on the join child pass 1 planned,
which carries a `logicalLink`, so the statistics branch wins unless a side
reports `sizeInBytes <= 1`. I instrumented this fallback and ran
`KeyGroupedPartitioningSuite` plus `EnsureRequirementsSuite`. It did not fire
once in 187 tests.
So this is about holding the invariant at both branches, not about a defect
today. The pre-alignment counts are one `collectFirst` away, and `:684` already
does that extraction for one side:
```scala
// The pre-alignment split counts, for the same reason the
statistics above read the
// pre-alignment plan. On a re-run both aligned reports hold
the same number of keys.
def rawNumKeys(plan: SparkPlan, aligned: Int): Int =
plan.outputPartitioning match {
case e: Expression => e
.collectFirst { case k: KeyedPartitioning =>
k.numPartitions }
.getOrElse(aligned)
case _ => aligned
}
rawNumKeys(unwrappedLeft, leftPartKeys.size) <
rawNumKeys(unwrappedRight, rightPartKeys.size)
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -731,38 +741,97 @@ case class EnsureRequirements(
}
/**
- * Unwraps a GroupPartitionsExec to get the underlying child plan.
+ * The innermost `GroupPartitionsExec` reachable from `plan` by descending
only through nodes
+ * this rule itself inserted above it, together with a function rebuilding
the traversed local
+ * sorts over a replacement node. `None` when no `GroupPartitionsExec` is
reachable.
+ *
+ * The descent only traverses a `GroupPartitionsExec` and a *local*
`SortExec`. That bound is a
+ * decision, not an omission: a `GroupPartitionsExec` hidden behind any
other node belongs to a
+ * different operator, and reusing it would move that operator's alignment.
Instrumenting the
+ * descent over `KeyGroupedPartitioningSuite`, the non-`SortExec` shapes
hiding a node are
+ * `Project > SortMergeJoin > Sort > GroupPartitions` and `Project > Filter
> Window >
+ * WindowGroupLimit > GroupPartitions`, where refusing to descend is right
every time. A global
+ * `SortExec` also stops the descent: it requires `OrderedDistribution`,
which a
+ * `KeyedPartitioning` can satisfy (behind
`spark.sql.sources.v2.bucketing.sorting.enabled`)
+ * through a `GroupPartitionsExec` built to emit the partition keys in
sorted order, and
+ * reusing that node for a join would destroy the ordering it exists to
provide.
*/
- private def unwrapGroupPartitions(plan: SparkPlan): SparkPlan = plan match {
- case g: GroupPartitionsExec => g.child
- case other => other
+ private def innermostGroupPartition(
+ plan: SparkPlan): Option[(GroupPartitionsExec, SparkPlan => SparkPlan)]
= plan match {
+ case g: GroupPartitionsExec =>
+ // A grouping over another grouping is one this rule added in an earlier
pass: keep the
+ // descent below it and drop this one.
+ innermostGroupPartition(g.child).orElse(Some((g, identity[SparkPlan])))
+ case s: SortExec if !s.global =>
+ innermostGroupPartition(s.child).map { case (g, rebuild) =>
+ (g, (newChild: SparkPlan) => s.withNewChildren(Seq(rebuild(newChild))))
+ }
+ case _ => None
}
+ /**
+ * Rewrites the innermost `GroupPartitionsExec` in `plan` with `f` and drops
any redundant
+ * grouping stacked above it, per the descent of
[[innermostGroupPartition]]. Returns `None`
+ * when `plan` holds no `GroupPartitionsExec`, leaving it to the caller to
create one.
+ *
+ * This is what makes the rule idempotent for storage-partitioned joins.
`EnsureRequirements`
+ * is re-run on plans it already produced: `AdaptiveSparkPlanExec` builds
one instance of this
+ * rule, and `ConvertSortMergeJoinToShuffledHashJoin` and
`OptimizeSkewedJoin` hand the whole
+ * tree back to it after rewriting some other join, all within one
+ * `queryStagePreparationRules` pass. A join child then arrives as
+ * `SortExec(GroupPartitionsExec(...))` rather than a bare scan, and the
distribution step adds
+ * a plain `GroupPartitionsExec` on top, because a partially clustered
`KeyedPartitioning`
+ * reports `isGrouped = false` by design and so is only satisfied "after
grouping". Rewriting
+ * that outer node instead of the one below it re-derives the alignment from
an already-aligned
+ * layout and duplicates rows; descending to the innermost node and dropping
what sits above it
+ * reproduces the plan a single pass would have produced.
+ *
+ * Dropping a grouping is safe because only `applyGroupPartitions` calls
this, reached from
+ * `checkKeyGroupCompatible`, which runs for joins alone: every
`GroupPartitionsExec` a join
+ * child carries is this rule's own. A single-child operator genuinely needs
its non-grouped
+ * input grouped and takes the wrap in the children loop instead;
`withJoinKeyPositions`, which
+ * other multi-child operators reach, does not reuse at depth.
+ */
+ private[exchange] def rewriteGroupPartitions(plan: SparkPlan)(
+ f: GroupPartitionsExec => GroupPartitionsExec): Option[SparkPlan] =
+ innermostGroupPartition(plan).map { case (g, rebuild) =>
+ val rewritten = f(g)
+ rewritten.copyTagsFrom(g)
+ rebuild(rewritten)
+ }
+
+ /**
+ * Unwraps the `GroupPartitionsExec` nodes this rule inserted over a join
child, down to the
+ * pre-alignment plan, per the descent of [[innermostGroupPartition]].
+ *
+ * The statistics-based replicate-side choice and the original partition
keys below must read
+ * from the pre-alignment plan on every pass: the local sort one level down
carries no
+ * `logicalLink` and reports the aligned layout instead.
+ */
+ private def unwrapGroupPartitions(plan: SparkPlan): SparkPlan =
+ innermostGroupPartition(plan).map(_._1.child).getOrElse(plan)
+
/**
* Applies or updates `GroupPartitionsExec` with the given parameters.
*
- * `GroupPartitionsExec` can be either the given plan node (child of the
join inserted by
- * `EnsureRequirement`) if the original child didn't satisfy the
distribution requirement; or we
- * can create a new one specifically for this join.
+ * Reuses the node this rule inserted over the join child in an earlier
pass, per the descent
+ * of [[innermostGroupPartition]], and creates a new one when the child
carries none.
*/
private def applyGroupPartitions(
plan: SparkPlan,
joinKeyPositions: Option[Seq[Int]],
mergedPartitionKeys: Seq[(InternalRowComparableWrapper, Int)],
reducers: Option[Seq[Option[KeyReducer]]],
distributePartitions: Boolean): SparkPlan = {
- plan match {
- case g: GroupPartitionsExec =>
- val newGroupPartitions = g.copy(
- joinKeyPositions = joinKeyPositions,
- expectedPartitionKeys = Some(mergedPartitionKeys),
- reducers = reducers,
- distributePartitions = distributePartitions)
- newGroupPartitions.copyTagsFrom(g)
- newGroupPartitions
- case _ =>
- GroupPartitionsExec(plan, joinKeyPositions, Some(mergedPartitionKeys),
reducers,
- distributePartitions)
+ rewriteGroupPartitions(plan) { g =>
+ g.copy(
+ joinKeyPositions = g.joinKeyPositions.orElse(joinKeyPositions),
+ expectedPartitionKeys = Some(mergedPartitionKeys),
+ reducers = reducers,
Review Comment:
**Finding 11.** `joinKeyPositions` one line up keeps what the reused node
holds, because the incoming value was computed against the node's already
projected report. `reducers` is computed the same way.
`leftSpec.reducersBothWays(rightSpec)` at `:564` runs over the reported
expressions, which on a re-run are the reduced ones this node produced. It is
written straight through.
Nothing goes wrong on this head, and I traced why. A re-run with reducers
never reaches here. Both sides report the same reduced keys and the same
expressions after pass 1, so `isCompatible` at `:525-528` is true.
`v2BucketingPartiallyClusteredDistributionEnabled` is false, because a reducer
needs `canReduceKeys` and that excludes it. So the whole `if` at `:529` is
skipped.
I measured it rather than only tracing it. Printing `g.reducers` and the
incoming value here and running `KeyGroupedPartitioningSuite` plus
`EnsureRequirementsSuite`, it fires 86 times, and every one is `existing=false,
incoming=true, expectedPartitionKeys set=false`. The node is always the one the
children loop just created, never an aligned one from an earlier pass.
So this is a request for the reason in the source, not a fix. Next to a
defended `orElse` the plain assignment reads like an oversight. And if that
short-circuit ever weakens, the reused node silently loses its reducer, its
keys go back to the raw space, and every expected key misses.
```suggestion
// Unlike `joinKeyPositions`, these need no `orElse`. A re-run with
reducers never reaches
// here. Both sides then report the same reduced keys, so
`isCompatible` above is true and
// the whole block is skipped.
reducers = reducers,
```
--
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]
