dongjoon-hyun commented on code in PR #58279:
URL: https://github.com/apache/spark/pull/58279#discussion_r3864640499
##########
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
+ * `KeyedPartitioning` can satisfy (behind
`spark.sql.sources.v2.bucketing.sorting.enabled`)
+ * through a `GroupPartitionsExec` built to emit the partition keys in
sorted order; reusing that
+ * node for a join would overwrite its `expectedPartitionKeys` and clear
`distributePartitions`,
+ * destroying the ordering it exists to provide.
+ *
+ * Dropping a grouping is safe only because this is reached from
`checkKeyGroupCompatible`, which
+ * runs for joins alone. An operator with a single child (an aggregate or a
window over a
Review Comment:
**The safety argument here is not accurate.** `rewriteGroupPartitions` is
also reached from `withJoinKeyPositions` (line 806), which is called at line
258 from the generic `children.zip(requiredChildDistributions)` loop — gated
only on `childrenIndexes.length > 1`, not on the parent being a join.
`CoGroupExec` and `FlatMapCoGroupsInBatchExec` require `ClusteredDistribution`
on two children and qualify.
What actually keeps a non-grouped GPE away from that path today is a config
coincidence: `KeyedShuffleSpec.canCreatePartitioning` requires
`!v2BucketingPartiallyClusteredDistributionEnabled`
(partitioning.scala:1374-1376). But the `OrderedDistribution` branch (lines
100-106) produces a non-grouped GPE under `v2BucketingAllowSorting` alone, and
`RemoveRedundantSorts` can delete the global sort shielding it — leaving a
shape where line 258 under a cogroup could drop the outer grouping and leave
the required distribution unsatisfied.
Suggest rewording this paragraph to state the real invariant (and/or
guarding the drop, e.g. only drop when the inner node's partitioning still
satisfies the required distribution), so the safety condition is enforced
rather than documented.
##########
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:
**`joinKeyPositions` index-space mismatch when reusing the inner node.**
With
`spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled=true`,
pass 2's `KeyedShuffleSpec` computes `joinKeyPositions` against the GPE's
*reported* partitioning, whose expressions are already projected
(`GroupPartitionsExec.outputPartitioning` line 76:
`joinKeyPositions.fold(k.expressions)(_.map(k.expressions))`). But the node
this now rewrites applies the positions to its *child's raw*
`KeyedPartitioning` (`groupedPartitionsTuple`,
GroupPartitionsExec.scala:128-138).
Example: table partitioned by `(a, b)`, join on `b`. Pass 1: inner
`GPE.joinKeyPositions = Some(Seq(1))`, reported expressions `[b]`. Pass 2
computes `Some(Seq(0))` in projected space and this code writes it through the
local sort onto the inner node — which now projects position 0 of the raw `[a,
b]`, i.e. groups by `a` while `expectedPartitionKeys` hold `b`-values.
`alignToExpectedKeys`'s `keyMap.getOrElse(key, Seq.empty)` misses on every key:
empty partitions, silently dropped rows. Same exposure via
`withJoinKeyPositions` (line 258 path).
Pre-PR the fresh outer wrapper's child exposed `[b]` at index 0, so `Seq(0)`
was self-consistent (the bare-GPE case had the hazard already; the descent
extends it to the `SortExec` and stacked shapes). Suggestion: compose the
positions with the reused node's existing `joinKeyPositions` (or skip the
overwrite when they are defined) so the stored positions stay in the raw
child's index space. No existing test combines
`allowJoinKeysSubsetOfPartitionKeys` with a second `EnsureRequirements` pass,
so nothing in the suites would catch this.
##########
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
+ * `KeyedPartitioning` can satisfy (behind
`spark.sql.sources.v2.bucketing.sorting.enabled`)
+ * through a `GroupPartitionsExec` built to emit the partition keys in
sorted order; reusing that
+ * node for a join would overwrite its `expectedPartitionKeys` and clear
`distributePartitions`,
+ * destroying the ordering it exists to provide.
+ *
+ * Dropping a grouping is safe only because this is reached from
`checkKeyGroupCompatible`, which
+ * runs for joins alone. An operator with a single child (an aggregate or a
window over a
+ * partially clustered join, say) genuinely needs its non-grouped input
grouped, and never gets
+ * here -- see `KeyGroupedPartitioningSuite`'s partially-clustered aggregate
and window tests.
+ */
+ private[exchange] def rewriteGroupPartitions(
Review Comment:
**Re-run still flips the replicate-side decision, and the flipped flags now
land on the inner node over the raw scan — unequal join partition counts.**
`unwrapGroupPartitions` (line 730) was not updated alongside this helper: on
the re-run it returns the ER-inserted `SortExec`, whose `logicalLink` is `None`
(ER-built sorts carry no `LOGICAL_PLAN_TAG`, and AQE's inherited-tag
propagation stops at the join node). The stats branch at lines 614-629 is
therefore skipped, and the fallback `leftPartKeys.size < rightPartKeys.size`
compares two already-aligned key lists that are equal by construction
(`alignToExpectedKeys` emits exactly `numSplits` entries per key in both modes)
— so pass 2 deterministically picks `replicateLeftSide = false`, flipping pass
1 whenever stats chose `true`.
Concretely: left key `k` = 3 raw splits (small bytes), right = 1 split
(large). Pass 1: `replicateLeftSide = true`, expected counts `{k: 1}`, both
sides emit 1 partition. Pass 2: flip; `numExpectedPartitions` is read off the
*aligned* layout (`{k: 1}`), and this helper writes `distributePartitions =
true` into the inner GPE over the raw scan — `splits.map(Seq(_)).padTo(1, ...)`
never truncates, so the left emits **3** partitions while the replicating right
emits **1** → `Can't zip RDDs with unequal numbers of partitions` from
`SortMergeJoinExec` (or duplicated/misaligned rows in multi-key variants).
Pre-PR the flipped values landed on a fresh outer node over the aligned
child (1 vs 1, correct), so this shape is a regression. The new test avoids it
only because `sp1 > sp2` makes pass 1 already choose `replicateLeftSide =
false`. Suggestion: give `unwrapGroupPartitions` the same descent as this
helper so pass 2 reads stats and the original partitioning from the
pre-alignment plan (which also fixes the stale `numExpectedPartitions` source
at lines 664-675).
##########
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 =>
+ val newGroupPartitions = g.copy(
+ joinKeyPositions = joinKeyPositions,
+ expectedPartitionKeys = Some(mergedPartitionKeys),
+ reducers = reducers,
+ distributePartitions = distributePartitions)
+ newGroupPartitions.copyTagsFrom(g)
+ newGroupPartitions
+ }.getOrElse {
+ GroupPartitionsExec(plan, joinKeyPositions, Some(mergedPartitionKeys),
reducers,
+ distributePartitions)
}
}
/**
* Applies join key positions to a plan by wrapping or updating
GroupPartitionsExec.
*/
private def withJoinKeyPositions(plan: SparkPlan, positions: Seq[Int]):
SparkPlan = {
- plan match {
- case g: GroupPartitionsExec =>
- val newGroupPartitions = g.copy(joinKeyPositions = Some(positions))
- newGroupPartitions.copyTagsFrom(g)
- newGroupPartitions
- case _ => GroupPartitionsExec(plan, joinKeyPositions = Some(positions))
- }
+ rewriteGroupPartitions(plan) { g =>
Review Comment:
Both callers pass the identical three-line lambda (bind `g.copy(...)`,
`copyTagsFrom(g)`, return). Tag propagation is an invariant of "rewrite this
node", not of either caller — consider folding `copyTagsFrom(g)` into the
helper's `GroupPartitionsExec` branch so callers pass a bare copy:
```scala
rewriteGroupPartitions(plan)(_.copy(joinKeyPositions = Some(positions)))
.getOrElse(GroupPartitionsExec(plan, joinKeyPositions = Some(positions)))
```
A future third caller then cannot forget the tag copy (a bug class with no
compile error and no test).
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -2898,6 +2898,44 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
}
}
+ test("partially clustered join keeps its row count when EnsureRequirements
re-runs") {
Review Comment:
Two suggestions to keep this regression test from going vacuously green:
1. **The `ALTER TABLE` split trick is unnecessary.** The suite's own
`createTable` helper (line 261) passes `numRowsPerSplit = 1`, and
`InMemoryBaseTable.withData` splits per row (`rows.size >= numRowsPerSplit`),
so a single `INSERT INTO ... VALUES (1, 'aa'), (1, 'ab'), (2, 'bb')` yields two
splits for `id = 1` — the idiom every other partially clustered test here uses
(e.g. line 812). This is the only raw `CREATE TABLE testcat` in the suite, and
the current construction couples the test to an incidental schema-evolution
branch of the fixture: if that behavior changes, the test keeps passing while
no longer covering the fix.
2. **`checkAnswer` alone cannot fail on unfixed code if the setup drifts.**
`Seq(1, 1, 2, 7)` is the correct answer for any plan, including a plain shuffle
join. The test only exercises the fix while AQE stays on by default,
`ConvertSortMergeJoinToShuffledHashJoin` fires on the `np` branch, and the SPJ
branch stays shuffle-free — none of which is asserted. Adding
`assert(collectShuffles(df.queryExecution.executedPlan).isEmpty)` for the SPJ
side plus an assertion that no `GroupPartitionsExec` has a
`GroupPartitionsExec` descendant (the stacking this PR eliminates) would make
it fail loudly instead.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -1524,6 +1525,63 @@ class EnsureRequirementsSuite extends SharedSparkSession
{
}
}
+ test("only a local sort is looked through when reusing GroupPartitionsExec")
{
+ // A global `SortExec` 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. Reusing that node for
a join would
+ // overwrite its `expectedPartitionKeys` and clear `distributePartitions`,
destroying the
+ // ordering it exists to provide. Only a local sort may be looked through.
+ val leaf = DummySparkPlan(
+ outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(InternalRow(1),
InternalRow(2))))
+ val gpe = GroupPartitionsExec(leaf)
+ val ordering = Seq(SortOrder(exprA, Ascending))
+ def mark(g: GroupPartitionsExec): GroupPartitionsExec =
g.copy(distributePartitions = true)
+
+ // A bare GroupPartitionsExec is rewritten in place.
+ EnsureRequirements.rewriteGroupPartitions(gpe)(mark) match {
+ case Some(g: GroupPartitionsExec) => assert(g.distributePartitions)
+ case other => fail(s"expected a rewritten GroupPartitionsExec, got
$other")
+ }
+
+ // A local sort is looked through and the GroupPartitionsExec below it is
rewritten.
+ val localSort = SortExec(ordering, global = false, gpe)
+ EnsureRequirements.rewriteGroupPartitions(localSort)(mark) match {
+ case Some(SortExec(_, false, g: GroupPartitionsExec, _)) =>
assert(g.distributePartitions)
+ case other => fail(s"expected the local sort to be looked through, got
$other")
+ }
+
+ // A global sort is not looked through, so the caller wraps instead of
reusing.
+ val globalSort = SortExec(ordering, global = true, gpe)
+ assert(EnsureRequirements.rewriteGroupPartitions(globalSort)(mark).isEmpty,
+ "a GroupPartitionsExec below a global sort must never be reused")
+ }
+
+ test("a single-child operator over a partially clustered layout still gets
grouped") {
Review Comment:
Note that this test never reaches the changed code and passes
byte-identically on the base commit: with one child, the
`childrenIndexes.length > 1` block is skipped, so neither
`applyGroupPartitions` nor `withJoinKeyPositions` (and hence
`rewriteGroupPartitions`) runs. The grouping it observes comes from the
untouched wrap at line 118-119.
It is still a meaningful guard — if someone later "generalizes" the reuse
into that wrap
(`rewriteGroupPartitions(child)(identity).getOrElse(GroupPartitionsExec(child))`),
the returned child would stay non-grouped and this assertion fails, which is
exactly what the scaladoc argues must not happen. But as named, a reader will
assume it covers the drop logic in `rewriteGroupPartitions`. Suggest a comment
(or rename) making explicit that it pins the intentional non-idempotence of the
line-119 wrap for single-child operators, not the new helper.
##########
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
+ * `KeyedPartitioning` can satisfy (behind
`spark.sql.sources.v2.bucketing.sorting.enabled`)
+ * through a `GroupPartitionsExec` built to emit the partition keys in
sorted order; reusing that
+ * node for a join would overwrite its `expectedPartitionKeys` and clear
`distributePartitions`,
+ * destroying the ordering it exists to provide.
+ *
+ * Dropping a grouping is safe only because this is reached from
`checkKeyGroupCompatible`, which
+ * runs for joins alone. An operator with a single child (an aggregate or a
window over a
+ * partially clustered join, say) genuinely needs its non-grouped input
grouped, and never gets
+ * here -- see `KeyGroupedPartitioningSuite`'s partially-clustered aggregate
and window tests.
+ */
+ private[exchange] def rewriteGroupPartitions(
+ plan: SparkPlan)(f: GroupPartitionsExec => GroupPartitionsExec):
Option[SparkPlan] = {
+ plan match {
+ case g: GroupPartitionsExec =>
+ // A grouping over another grouping is one this rule added in an
earlier pass: drop it and
+ // rewrite the node below, which is the one that owns the alignment.
+ rewriteGroupPartitions(g.child)(f).orElse(Some(f(g)))
+ case s @ SortExec(_, false, _, _) =>
Review Comment:
Nit: the local-sort look-throughs elsewhere use the named form —
`AQEUtils.getRequiredDistribution`,
`ConvertSortMergeJoinToShuffledHashJoin.stripSort`, and `simpleCosting` all
write `case s: SortExec if !s.global`. The positional `SortExec(_, false, _,
_)` (mirrored in the new test) hard-codes the arity and field order: a new
`SortExec` parameter breaks this match at compile time in two places, and a
reordering of `global`/`child` would break it silently.
```scala
case s: SortExec if !s.global =>
rewriteGroupPartitions(s.child)(f).map(c => s.copy(child = c))
```
(with `copyTagsFrom(s)` if you switch from `withNewChildren` to `copy`,
since `withNewChildren` propagates tags).
--
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]
