ulysses-you commented on code in PR #58279:
URL: https://github.com/apache/spark/pull/58279#discussion_r3921907835
##########
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:
Confirmed -- a regression of this PR's descent. `unwrapGroupPartitions` now
shares the rewrite descent (`innermostGroupPartition`), so the statistics and
the original partition keys read from the pre-alignment plan on every pass. New
regression test has the replicated side hold more splits for id = 1 than the
other side, so a flipped decision overflows `padTo` on the distribute side and
the join sides end up with an unequal number of partitions. 5af8667631b
##########
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:
Confirmed. `applyGroupPartitions` keeps the positions a reused node already
holds: they were computed against the child's raw partition keys and stay
authoritative for this join, while the incoming ones are in the node's already
projected report. The new regression test partitions the left table by `(extra,
id)` and joins on the second partition key -- joining on the first column would
not discriminate, since the positions agree on both passes there. 5af8667631b
##########
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:
Confirmed. The descent and drop are now confined to `applyGroupPartitions`,
which is reached from `checkKeyGroupCompatible` for joins alone;
`withJoinKeyPositions` reuses only a topmost node again, since every
multi-child operator reaches it. The scaladoc states that invariant, and a unit
test pins the top-most behavior. 5af8667631b
##########
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:
Done -- `case s: SortExec if !s.global`. Kept `withNewChildren`, which
already propagates the sort's tags. 5af8667631b
##########
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:
Done -- `copyTagsFrom` moved into the helper's `GroupPartitionsExec` branch,
callers pass a bare `copy`. A unit test now pins tag retention on the bare,
local-sort and top-most rewrite paths. 5af8667631b
##########
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:
Done both. The test is rebuilt on the suite's `createTable` idiom
(`numRowsPerSplit = 1`), and it asserts the storage-partitioned side stays
shuffle-free and no `GroupPartitionsExec` is stacked over another. 5af8667631b
##########
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:
Done -- the test's comment now states that it pins the children loop's wrap
for single-child operators and never runs `rewriteGroupPartitions` (with one
child the multi-child block is skipped). 5af8667631b
--
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]
