peter-toth commented on code in PR #58279:
URL: https://github.com/apache/spark/pull/58279#discussion_r3917064306
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -732,6 +732,47 @@ case class EnsureRequirements(
case other => other
}
+ /**
+ * Finds the innermost `GroupPartitionsExec` in `plan`, rewrites it with
`f`, and drops any
+ * redundant grouping stacked above it. 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 --
`ConvertSortMergeJoinToShuffledHashJoin` and
+ * `OptimizeSkewedJoin` hand the whole tree back to it after rewriting some
other join -- so a
+ * join child arrives as `SortExec(GroupPartitionsExec(...))` rather than a
bare scan. The
+ * distribution step then 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 fresh outer node instead of the one below it
re-derives the
+ * alignment from an already-aligned layout: the inner node replicates an
input partition across
+ * the expected partitions and the outer one concatenates those replicas
back together before
+ * replicating again, duplicating rows. Descending to the innermost node and
dropping what sits
+ * above it reproduces exactly the plan a single pass would have produced.
+ *
+ * Only a *local* `SortExec` is traversed. A global one requires
`OrderedDistribution`, which a
Review Comment:
**[Non-blocking, docs]** This justifies the restriction through the
`OrderedDistribution` case, which
is the sharpest consequence, but it does not state the rule the descent
actually follows: descend only
through a node this rule itself inserted directly above this child. Without
that, a later reader has no
reason not to generalise the descent to a `ProjectExec`, and it looks like
an omission rather than a
decision.
It is worth stating because generalising it would be wrong. I instrumented
the descent over
`KeyGroupedPartitioningSuite`: a non-`SortExec` node hides a
`GroupPartitionsExec` from it 6 times, and
refusing to descend is right every time. They are `Project > SortMergeJoin >
Sort > GroupPartitions` and
`Project > Filter > Window > WindowGroupLimit`, where the hidden node
belongs to a different operator,
so reusing it would move another operator's alignment.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -745,32 +786,29 @@ case class EnsureRequirements(
mergedPartitionKeys: Seq[(InternalRowComparableWrapper, Int)],
reducers: Option[Seq[Option[Reducer[_, _]]]],
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 =>
Review Comment:
**[Non-blocking, altitude]** Not this PR's to fix, but worth having on the
record, because the reason
this fix is needed at all is that the rule has no single answer for "this
node is my own output". There
are four shapes today and this makes five:
| site | what it does with the rule's own output |
|---|---|
| the `ShuffleExchangeExec` case | `s.copy(outputPartitioning = ...)`,
reuses in place |
| the `GroupPartitionsExec` case beside it |
`ShuffleExchangeExec(newPartitioning, gpe.child)`, unwraps it |
| `unwrapGroupPartitions` | peels one level, only to reach `logicalLink` |
| the children loop plus `applyGroupPartitions` | did not account for it at
all, which is this bug |
| `rewriteGroupPartitions`, **new here** | descends through a local sort and
drops what is above |
Consolidating them is the follow-up I described in the body: stop inserting
a placeholder above a
co-partitioned child, and decide that child's node once where the merged
keys are known. That retires
both `rewriteGroupPartitions`' drop and `unwrapGroupPartitions`. I am not
asking for it here.
--
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]