ulysses-you commented on code in PR #58279:
URL: https://github.com/apache/spark/pull/58279#discussion_r3930412603
##########
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:
Done -- one descent per side: `leftGrouping`/`rightGrouping` are computed
once at the top of the block, and the unwrapped plans, the statistics, the
original keys and the positions all derive from them. `unwrapGroupPartitions`
keeps its one remaining caller at the shuffle site. 94afe6e4401
##########
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:
Done -- the fallback compares the pre-alignment split counts, read through a
new `PartitioningCollection.numKeyedPartitions` (the representative keyed
member, collections included). Pinned by a unit test that forces the fallback
-- dummy plans carry no `logicalLink` -- and reads the choice back off the
`distributePartitions` flags; reverting the fix fails its first arm (`(true,
false)` against the expected `(false, true)`). 94afe6e4401
##########
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:
Done -- comment added verbatim. 94afe6e4401
--
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]
