Skip to content

Commit db78554

Browse files
committed
[SPARK-59289][SQL] Decide a co-partitioned child's GroupPartitionsExec once, where the join's merged keys are known
### What changes were proposed in this pull request? `GroupPartitionsExec` is how storage-partitioned join lines two scans up: it coalesces, projects, reorders and pads a scan's partitions. Today the planner decides that node twice and the node then re-derives it many times. This PR makes the planner decide once and the node carry the answer. **Where the node comes from.** `ensureDistributionAndOrdering` resolved every child on its own first, so it placed a `GroupPartitionsExec` over a co-partitioned child before knowing the parent was a join. `checkKeyGroupCompatible` then rewrote or replaced it, descending past a local sort, merging new parameters in and dropping any grouping stacked on top. The two sites also worked in different index spaces, the placeholder's positions against the raw partition keys and the spec's against the node's already projected report, which is what `joinKeyPositions.orElse` and the reducer fast path existed to paper over. The two paths are now apart. The per-child loop resolves only the children that answer for themselves, through a lifted `resolveChild`. A new `coPartitionChildren` owns the rest end to end: it peels every grouping this rule inserted off both children, plans the pairing from the sources' own reports, and builds each node once. **How the children are carried.** The co-partitioned set used to be a `Seq[Int]` of child indexes beside an `isCoPartitioned` flag, threaded alongside `Seq[Distribution]`. It is one value now: a `Seq[Option[ClusteredDistribution]]` aligned with `children`, holding what each child has to satisfy and nothing for a child that answers for itself. Readers `zip` it instead of indexing, and it carries the type the split point already knew, so three asserts, four `asInstanceOf[ClusteredDistribution]` and a dead branch in `candidatesFor` go with it. `ensureDistributionAndOrdering` is 53 lines and holds no `var`. **Who decides what.** `coPartitionChildren` recognises only the three shapes that need nothing arranged between the children, each at the same altitude: every side a single small partition, a storage-partitioned join, or two compatible pass-through specs. The shuffle is `shuffleToCoPartition`, and the layout it drags everyone onto is `pickCoPartitionTarget`, returning a `CoPartitionTarget` that names the winning member and, per child, the member pairing with it. `ShuffleSpecCollection` answers `isCompatibleWith` existentially and so never names the member the two sides agreed on, which is why that pairing used to be rediscovered at three levels. It is decided once and carried. **What the node holds.** `grouping` (the index groups it emits and what they say about the layout) and `outputPartitioning` are constructor fields now, derived once by `GroupPartitionsExec.apply`, the way `ShuffleExchangeExec` holds the partitioning it produces. They were per-instance lazy vals, so every `copy` and `withNewChildren` threw the memo away. **What `satisfies` answers.** `KeyedPartitioning.satisfies` returned `true` for a partitioning that only satisfies after a projecting node, for the storage-partitioned join's benefit alone. It is strict now, and the loose question moved to `keysCanSatisfy`, which `keysMaySatisfy` and `mayGroupToSatisfy` compose with the strict one. `EnsureRequirements.resolveKeyedPartitioning` picks the positions to project onto through a new `KeyedPartitioning.positionsCoveringClusterKeys`, which `keysSatisfy` also reads, so the predicate that decides whether a projection is needed and the one that picks the positions cannot drift apart. Whether a projection is permitted at all is one `mayProjectToClusterKeys`, where it used to be spelled twice in opposite polarity. **A small vocabulary on `ClusteredDistribution`.** Five `Partitioning.satisfies0` implementations spelled out the same `if (requireAllClusterKeys) areAllClusterKeysMatched(exprs) else exprs.forall(x => clustering.exists(_.semanticEquals(x)))`, so the flag was read in five places and the membership test was written out nine times across the file. The pair is now `matchesClusterKeys`, beside the `areAllClusterKeysMatched` it wraps, and the membership test it is built from is `isClusterKey`. `allClusterKeysAmong` replaces the mirror fold, which asks whether a set of attributes covers the clustering, and `AQEUtils` reaches for it too. `areAllClusterKeysMatched` on both distributions, and the `StatefulOpClusteredDistribution` arm that spelled the same thing out at a call site, are one `Seq.corresponds` each. The five callers are `HashPartitioningLike`, `NullAwareHashPartitioning`, `CoalescedNullAwareHashPartitioning`, `RangePartitioning` and `ShufflePartitionIdPassThrough`, and each rewrite is one call over the same expressions in the same order. The prose in these files called the same thing an "operation key" in about forty places, which was a second name for what the API already calls a cluster key, so that is gone too. **What a second run sees.** Peeling and re-planning would re-decide a pairing an earlier run already settled, on an input that by then holds the keyed shuffle that run inserted. So `checkKeyGroupCompatible` asks first whether the two children satisfy their distributions and line up with each other as they arrive, and hands them back untouched when they do. That question decides nothing new on a plan this rule has not seen, since a source that needs a node does not satisfy its distribution as it stands. EnsureRequirements ensureDistributionAndOrdering the two paths, and the ordering step resolveChild a child that answers for itself coPartitionChildren children that only answer together checkKeyGroupCompatible try the join, build both nodes once alreadyCoPartitioned unless a previous run settled it resolveEachChild else each on its own shuffleToCoPartition and shuffle whoever is left over pickCoPartitionTarget onto one layout, picked once GroupPartitionsExec(child, grouping, outputPartitioning, <recipe>, enableSortedMerge) GroupPartitionsExec.apply the only way to build one computeGrouping pure, takes the child's KeyedPartitioning computeOutputPartitioning pure, takes the same Retires `rewriteGroupPartitions`, `applyGroupPartitions` and `innermostGroupPartition`'s sort-rebuilding half, along with the `orElse` and the two index spaces they served. ### Why are the changes needed? **A join could commit to a pairing its own children then refuse.** `GroupPartitionsExec` gives up its keyed claim when it turns out to regroup a layout that pins undeclared rows to `hash(key) % numPartitions`, and only the node knows the permutation it performs, so that answer arrived after `checkKeyGroupCompatible` had skipped both shuffles. The result is a plan `ValidateRequirements` rejects, and every `AQEShuffleReadRule` and `OptimizeSkewedJoin` drops its result on such a stage, so partition coalescing, local read and skew join are all off for it. Now the site that builds the two nodes asks them before returning. This closes SPARK-59272. **The same derivation ran many times.** Measured on `KeyGroupedPartitioningSuite`'s "join with two partition keys and matching & sorted partitions": the base derives the grouping 16 times over 48 node instances, this derives it 8 times over 64. A lazy val belongs to one instance, and the columnar and codegen rules rebuild the tree after this rule, so the derivations followed the instances. As a field they follow the planner's decisions instead. **The caller was compensating for the predicate.** `resolveKeyedPartitioning` asked in two steps whether a member needed a node at all, the cheap full-coverage test and then the projection that decides whether a narrowing merges anything. `keysSatisfy` owns both now and the caller is `admitted.find(_.satisfiesAsIs)`. `ValidateRequirements` also becomes a real guard: on a `KeyedPartitioning([a, b])` with two partitions sharing `a = 1` under a `ClusteredDistribution([a])`, `validate` was `true` before and is `false` now, and the plan spreads rows sharing the cluster key either way. **The rule was not idempotent.** Applying it to its own output changed the plan, which matters because AQE re-plans every query stage and two AQE rules hand the tree back to it. Measured over `KeyGroupedPartitioningSuite` by re-running the rule on each of its own results, 2713 applications in all: the base differs on 2, and it differs by adding a shuffle. Without the check described above this differs on 16, all of them a grouping node added over a shuffle. With it, on none. ### Does this PR introduce _any_ user-facing change? **Yes**, three plan changes. No public API changes. What does move is `GroupPartitionsExec`'s case class signature and the predicates on `KeyedPartitioning` and `ClusteredDistribution`, in `sql.execution` and `sql.catalyst`, and `MimaExcludes` treats both packages as internals in its `defaultExcludes` section. **A join declines rather than leaving an unvalidatable plan.** On the `SPARK-59050: SPJ: regrouping a marked layout must not keep the unknown-keyed claim` query the second join now takes a keyed one-side shuffle, three shuffles to four, and `ValidateRequirements.validate` on the executed plan goes from `false` to `true`. The alternative was not "no shuffle", it was "no shuffle and no AQE". **`satisfies` accepts a partitioning whose expression *is* a cluster key**, such as a `years(ts)` under a clustering naming `years(ts)`. The old reference-level test refused it unless `requireAllClusterKeys` was set, while the `requireAllClusterKeys` arm accepted the same shape, so this removes an inconsistency. It is a widening on the default configuration. **A re-planned stage keeps the plan it had.** The rule is idempotent now, where the base can add a shuffle on the second pass, so an AQE stage's plan no longer drifts from the one first planned. ### How was this patch tested? Eleven tests, ten of them new and one a rewrite of the existing peel test. Each fails once the decision it pins is put back, so none of them is vacuous. | test | the decision put back | |---|---| | `DistributionSuite`: satisfies is strict about a projection that merges partitions | the loose `keysSatisfy` | | `DistributionSuite`: satisfies accepts an expression that is itself a cluster key | the reference-only coverage test | | `ShuffleSpecSuite`: a collection whose members all need grouping still yields them | the collection filter on strict `satisfies` | | `EnsureRequirementsSuite`: SPARK-58996 only a local sort is looked through | descending through a global sort | | a local sort with no grouping under it is left alone | peeling a sort with nothing under it | | a grouped side is paired on its own key order | `toGrouped` for a grouped source too | | pushing join key positions into a node re-derives its grouping | `copy` instead of the rebuild | | single-partition children still honour a required partition count | leaving both children alone | | a partitioning is never projected onto no position | projecting onto the empty position set | | a second pass leaves a pairing this rule already made alone | re-deciding a settled pairing | | pushed positions compose with the ones a node already carries | replacing a node's positions instead of composing | One existing expectation moved, `Seq(true, true, true)` to `Seq(true, true, true, true)` in the SPARK-59050 regrouping test, and it now asserts `ValidateRequirements.validate` ahead of the shuffle count because that names the reason the shuffle exists. Three test names lost the word "operation" with the rest of the prose. **A differential sweep is the strongest evidence here.** 307200 generated storage-partitioned join plans, 40 partitioning shapes a side crossed with six join types over 32 configuration cells, run against the base and against this change. Each plan is checked twice, that `ValidateRequirements` passes and that two sides whose shuffles were both skipped really declare the same key sequence. The cases are keyed and set-differenced, not just counted. | | base | this change | |---|---|---| | plans that skip both shuffles yet disagree on the keys | 2024 | **0** | | plans `ValidateRequirements` rejects | 6724 | 5296 | | plans the planner cannot build at all | 29952 | 29356 | **Nothing is worse on any of the three, and the sets line up exactly.** The 2024 co-partitioning violations are the same 2024 cases as the removed validation failures, which is SPARK-59272 in one number. The 596 fewer planner failures are the same 596 cases that now appear among the rejected plans, so a shape the base could not plan at all is now planned, badly. Zero cases are added to any of the three. The remaining 29356 failures are shared with the base and are the generator's, not the planner's. They come from `Partitioning.createShuffleSpec`, which throws when a co-partitioned child reports `UnknownPartitioning`, and that needs a marked layout that is not grouped. Nothing builds one. The marker is only ever put on a layout `KeyedShuffleSpec.createPartitioning` has just laid out one partition per key, and both paths that would regroup a marked layout refuse it. Green: `DistributionSuite`, `ShuffleSpecSuite`, `EnsureRequirementsSuite`, `ValidateRequirementsSuite`, `PlannerSuite`, `GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite`, `KeyGroupedPartitioningSuite`, `DataSourceV2CatalystRuntimeFilterSuite`, `AdaptiveQueryExecSuite`, `ExchangeSuite`, `JoinSuite`, `CoalesceShufflePartitionsSuite` and `DataSourceV2Suite`, 1624 tests. `dev/lint-scala` is clean. All seven plan-stability suites are green as well, 322 tests, so no golden TPCDS or TPCH plan changed. That is the net for a stray plan change out of the `EnsureRequirements` restructuring, since the rule runs on every query and those two workloads have no storage-partitioned join. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code
1 parent 2870b66 commit db78554

9 files changed

Lines changed: 1547 additions & 1045 deletions

File tree

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala

Lines changed: 243 additions & 116 deletions
Large diffs are not rendered by default.

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala

Lines changed: 61 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,23 @@ package org.apache.spark.sql.catalyst
2020
import org.apache.spark.SparkFunSuite
2121
/* Implicit conversions */
2222
import org.apache.spark.sql.catalyst.dsl.expressions._
23-
import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, CollationAwareMurmur3Hash, Expression, Literal, Pmod, SortOrder}
23+
import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, CollationAwareMurmur3Hash, Expression, Literal, Pmod, SortOrder, TransformExpression}
2424
import org.apache.spark.sql.catalyst.plans.SQLHelper
2525
import org.apache.spark.sql.catalyst.plans.physical._
26+
import org.apache.spark.sql.connector.catalog.functions.ScalarFunction
2627
import org.apache.spark.sql.internal.SQLConf
27-
import org.apache.spark.sql.types.IntegerType
28+
import org.apache.spark.sql.types.{DataType, IntegerType}
2829

2930
class DistributionSuite extends SparkFunSuite with SQLHelper {
3031

32+
/** A bound function with a stable canonical name, enough to build a `TransformExpression`. */
33+
private object TestYearsFunction extends ScalarFunction[Int] {
34+
override def inputTypes(): Array[DataType] = Array(IntegerType)
35+
override def resultType(): DataType = IntegerType
36+
override def name(): String = "years"
37+
override def canonicalName(): String = "test.years"
38+
}
39+
3140
protected def checkSatisfied(
3241
inputPartitioning: Partitioning,
3342
requiredDistribution: Distribution,
@@ -413,6 +422,56 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
413422
checkSatisfied(groupedKP, ClusteredDistribution(Seq(x)), true)
414423
}
415424

425+
test("SPARK-59289: satisfies is strict about a projection that merges partitions") {
426+
val a = AttributeReference("a", IntegerType)()
427+
val b = AttributeReference("b", IntegerType)()
428+
val clustered = ClusteredDistribution(Seq(a))
429+
430+
// Partitioned by [a, b], clustered on [a] alone. Both are grouped, so the only question is
431+
// whether rows sharing `a` share a partition.
432+
val merging = KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 1), InternalRow(1, 2)))
433+
val notMerging = KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 1), InternalRow(2, 2)))
434+
assert(merging.isGrouped && notMerging.isGrouped)
435+
436+
withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
437+
// a = 1 sits on two partitions, so this satisfies nothing until a GroupPartitionsExec
438+
// projects the keys onto [a]. `keysMaySatisfy` is the question that says so.
439+
checkSatisfied(merging, clustered, false)
440+
assert(merging.keysMaySatisfy(clustered))
441+
442+
// Projecting here would merge no partition, so every `a` already sits on one and the
443+
// partitioning satisfies as it stands. Keeping it beats projecting, since it still names `b`
444+
// for a downstream operator to co-partition on.
445+
checkSatisfied(notMerging, clustered, true)
446+
}
447+
}
448+
449+
test("SPARK-59289: satisfies accepts a partition expression that is itself a cluster key") {
450+
val ts = AttributeReference("ts", IntegerType)()
451+
val other = AttributeReference("other", IntegerType)()
452+
val years = TransformExpression(TestYearsFunction, Seq(ts))
453+
val partitioning = KeyedPartitioning(Seq(years), Seq(InternalRow(1), InternalRow(2)))
454+
assert(partitioning.isGrouped)
455+
456+
// The clustering names the transform rather than the column it reads, so rows sharing
457+
// `years(ts)` already sit on one partition and nothing has to be re-partitioned. The test that
458+
// decides this reads the expression as well as its references, and it runs on the default
459+
// configuration.
460+
checkSatisfied(partitioning, ClusteredDistribution(Seq(years)), true)
461+
462+
// The `requireAllClusterKeys` arm has always accepted this shape, since it compares
463+
// expressions rather than references, so the two arms agree now.
464+
checkSatisfied(
465+
partitioning, ClusteredDistribution(Seq(years), requireAllClusterKeys = true), true)
466+
467+
// Clustered on the column the transform reads, rows sharing `ts` share a partition too, since
468+
// the partition expression is a function of it.
469+
checkSatisfied(partitioning, ClusteredDistribution(Seq(ts)), true)
470+
471+
// Neither the expression nor its reference is a cluster key here, so this is refused.
472+
checkSatisfied(partitioning, ClusteredDistribution(Seq(other)), false)
473+
}
474+
416475
test("SPARK-56877: fromPartitionings reuses already-consistent nested collections") {
417476
val x = AttributeReference("x", IntegerType)()
418477
val y = AttributeReference("y", IntegerType)()

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import org.apache.spark.sql.internal.SQLConf
2727
import org.apache.spark.sql.types.{DataType, IntegerType, LongType, StructType}
2828

2929
class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
30+
3031
private val passThrough_a_10 = ShufflePartitionIdPassThrough(DirectShufflePartitionID($"a"), 10)
3132
private val passThrough_b_10 = ShufflePartitionIdPassThrough(DirectShufflePartitionID($"b"), 10)
3233
private val passThrough_c_10 = ShufflePartitionIdPassThrough(DirectShufflePartitionID($"c"), 10)
@@ -39,6 +40,19 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
3940
override def canonicalName(): String = "test.bucket"
4041
}
4142

43+
test("SPARK-59289: createShuffleSpec keeps members that only satisfy after grouping") {
44+
val a = AttributeReference("a", IntegerType)()
45+
// Filtering the members on the strict `satisfies` would leave nothing here, and an empty
46+
// `ShuffleSpecCollection` is rejected by its own `require`.
47+
val ungroupedKeyed = KeyedPartitioning(
48+
Seq(a), Seq(InternalRow(1), InternalRow(1), InternalRow(2)))
49+
assert(!ungroupedKeyed.isGrouped, "test setup: no member serves the distribution as it stands")
50+
val collection = PartitioningCollection(Seq(ungroupedKeyed, ungroupedKeyed))
51+
val spec = collection.createShuffleSpec(ClusteredDistribution(Seq(a)))
52+
assert(spec.asInstanceOf[ShuffleSpecCollection].specs.size == 2)
53+
assert(spec.flatten.map(_.numPartitions) == Seq(3, 3))
54+
}
55+
4256
protected def checkCompatible(
4357
left: ShuffleSpec,
4458
right: ShuffleSpec,

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ object AQEUtils {
4545
case p: ProjectExec =>
4646
getRequiredDistribution(p.child).flatMap {
4747
case h: ClusteredDistribution =>
48-
if (h.clustering.forall(e => p.projectList.exists(_.semanticEquals(e)))) {
48+
if (h.allClusterKeysAmong(p.projectList)) {
4949
Some(h)
5050
} else {
5151
// It's possible that the user-specified repartition is effective but the output

0 commit comments

Comments
 (0)