This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new a615688368b6 [SPARK-59671][SQL] Validate a co-partitioned pair by its
children's pairing
a615688368b6 is described below
commit a615688368b6ad87a31f4c6a6bb25a6c1e4752a1
Author: Xiduo You <[email protected]>
AuthorDate: Thu Oct 8 17:48:19 2026 +0800
[SPARK-59671][SQL] Validate a co-partitioned pair by its children's pairing
### What changes were proposed in this pull request?
`ValidateRequirements` judges a multi-child co-partitioned operator (a
join, a cogroup) on the pairing of its children: every side offers the layouts
it actually reports (`PartitioningCollection.specsForPairing` — no key deduped,
none re-sorted), and one member of the first side has to pair with every other
side (`KeyedShuffleSpec.isCompatibleWith`). `validateInternal` keeps the base
loop; for a multi-child clustered operator the per-child distribution clause is
delegated to that pair [...]
For that question to be decidable, a keyed layout that repeats partition
keys has to say where the repetition comes from, so `KeyLayout` carries
`ungroupingOrigin: Option[UngroupingOrigin]`, a sealed ADT with three values:
- `SPLIT_FOR_JOIN` / `REPLICATED_FOR_JOIN` — the two complementary sides of
a partially clustered alignment. For a key, a partition holding part of it on
the spread side holds all of it on the repeating side, so pairing the two index
by index loses no match, and that split is the whole of the pair's soundness.
- `SPLIT_FOR_ORDERING` — the spread `EnsureRequirements.resolveChild`
builds for an `OrderedDistribution`. It pairs with nothing.
Producers and carriers:
- Only `EnsureRequirements` stamps, at the two sites that build these
shapes; the alignment stamp stays gated by the existing
`canDuplicateLeftSide`/`canDuplicateRightSide` checks, so a `FullOuter` join
gets no spread pair and a `LeftOuter` repeats its right side or groups. A
`SortMergeAsOfJoinExec` gets no keyed alignment at all: its matches are read in
the order within a partition, which neither side preserves.
- `GroupPartitionsExec` drops the `distributePartitions` constructor field
and derives the routing from the origin (`ungroupingOrigin.exists(_ !=
REPLICATED_FOR_JOIN)`); `computeGrouping` writes the node's *own* origin into
the layout it builds, replacing any claim the child's layout carried, and only
while the computed shape actually stays ungrouped: a grouping that settles
every key spends the stamp the way `toGrouped` does, so a grouped layout never
carries a claim.
- Carriers keep one agreed claim: `toGrouped` spends the origin, `concat`
and `fromPartitionings` keep it only where every member carries the same, and
`withKeyPositionsFor` forwards its node's own.
Consumers:
- `KeyedPartitioning.satisfies`: a clustering asks that a key's rows share
a partition, which no ungrouped layout gives on its own, so the clustering arm
stays the base `isGrouped && keysSatisfy`, and the stamp is read only where a
pair is judged (`pairsUngrouped`, feeding `specsForPairing`; the role check in
`isCompatibleWith`). An ordering pairs nothing, so an ungrouped layout answers
it on the keys alone, with the keys checked to ascend in the sequence the
ordering reads (`keysSort [...]
- `KeyedShuffleSpec.isCompatibleWith`: an ungrouped pair is compatible only
as one spread side against one repeating side.
### Why are the changes needed?
https://issues.apache.org/jira/browse/SPARK-59671
The validator compared children through `createShuffleSpec` — the
planner-facing view (keys deduped and re-sorted) of what a node about to be
inserted would report. Two failure modes follow:
1. **Sound plans fail as a whole.** A partially clustered alignment leaves
both of its sides ungrouped on purpose, so a finished plan never matches the
planner-facing view; `AdaptiveSparkPlanExec.optimizeQueryStage` validates the
whole candidate plan and one unsatisfied node reverts every
`AQEShuffleReadRule` application in it. An aggregate's shuffle read then stays
uncoalesced because a join elsewhere in the plan is aligned, e.g. (with
`spark.sql.sources.v2.bucketing.partiallyCluster [...]
2. **Admitting ungrouped layouts without the producer's fact is a
wrong-results hazard.** `AQEShuffleReadExec` with mapper partition specs
reports `s.child.outputPartitioning` — the layout from *before* the shuffle,
e.g. a scan whose key `1` sits in two splits (`[1, 1, 2]`). Two such sides
agree key by key and on the count while neither holds a key's whole group; a
join consuming them as aligned zips mapper `i` against mapper `i` and loses
every match that crosses mappers (3 rows inst [...]
### Does this PR introduce _any_ user-facing change?
No change to query results. Plans regain the AQE shuffle-read optimizations
(partition coalescing, local reads where sound) that a plan holding a
deliberately ungrouped layout used to revert stage-wide; one such layout — an
ungrouped keyed one under an `OrderedDistribution`, which pairs no partition —
now satisfies it, where grouping was required before. Ungrouped layouts are
otherwise consumed no more widely than on base: a clustering still refuses one
unless a producer vouched for t [...]
### How was this patch tested?
Tests were added/rewritten at four levels:
- End-to-end (`KeyGroupedPartitioningSuite`): `a partially clustered join
leaves AQE's shuffle coalescing alone`; `a three-table chain keeps its AQE
coalescing` (also pins the complementary stamps below the joins and the
claim-free join outputs); `an AQE local read over a split scan does not lose
matches` (the data-consistency regression, with final-plan assertions: three
shuffle stages, no local read reporting a keyed layout, no grouping node over a
partial read); `an ungrouped sort [...]
- Producer gates (`EnsureRequirementsSuite`): `an ordering spread carries
its own reason and serves only the ordering`; `the alignment stamps one
repeating side, only where the join type may` (Inner, both LeftOuter picks,
FullOuter); `an as-of join gets no keyed alignment, so its sides shuffle`.
- Pairing rules (`ValidateRequirementsSuite`, 11 SPARK-59671 tests): the
role rules, count/order/hashed-side refusals, cogroup, collapsed layouts,
ordering-versus-pairing.
- Layout mechanics (`ShuffleSpecSuite`): `the specs a side offers are the
layouts it reports`; `carriers keep one agreed origin and grouping spends it`;
`an ordering reads sorted keys, a clustering reads the pairing` (both halves of
the ordering gate: an unsorted ungrouped layout and an unsorted grouped one
answer false). `GroupPartitionsExecSuite`: `the output claim is the node's own,
never inherited`; `an alignment that settles every key spends its stamp`.
Full runs: catalyst 34/34 (`ShuffleSpecSuite`, `DistributionSuite`); sql
377/377 (`KeyGroupedPartitioning*` — all three classes —
`ValidateRequirementsSuite`, `EnsureRequirementsSuite`,
`GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite`);
`dev/lint-scala` clean. Each correctness gate was additionally
mutation-checked: relaxing it (admitting an unstamped ungrouped layout to a
pairing, letting a stamped layout serve a clustering on its own, dropping the
key-order check, [...]
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude Code (session model DeepSeek Flash) — analysis,
implementation and test iteration under the author's direction; the author
reviewed every change.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Closes #58942 from ulysses-you/spj-validate-pairing.
Authored-by: Xiduo You <[email protected]>
Signed-off-by: Xiduo You <[email protected]>
(cherry picked from commit 91f015743a1c48d96ed2e2fb76ec86eefe84ef98)
Signed-off-by: Xiduo You <[email protected]>
---
.../sql/catalyst/plans/physical/partitioning.scala | 251 +++++++++++++++--
.../spark/sql/catalyst/ShuffleSpecSuite.scala | 144 +++++++++-
.../datasources/v2/GroupPartitionsExec.scala | 62 ++--
.../execution/exchange/EnsureRequirements.scala | 45 +--
.../execution/exchange/ValidateRequirements.scala | 54 +++-
.../connector/KeyGroupedPartitioningSuite.scala | 287 ++++++++++++++++++-
.../datasources/v2/GroupPartitionsExecSuite.scala | 52 +++-
.../exchange/EnsureRequirementsSuite.scala | 95 ++++++-
.../exchange/ValidateRequirementsSuite.scala | 312 ++++++++++++++++++++-
9 files changed, 1212 insertions(+), 90 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
index f14410eb0ae5..ca9cadcbcfcd 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
@@ -433,6 +433,20 @@ case class CoalescedNullAwareHashPartitioning(
copy(from = from.copy(expressions = newChildren))
}
+// Describes why a keyed layout reports repeated partition keys on purpose,
rather than leaving them
+// for the grouping node that would settle them.
+sealed trait UngroupingOrigin
+
+// One side of a partially clustered alignment: it keeps its splits and
spreads them across the
+// slots its key is expected in.
+case object SPLIT_FOR_JOIN extends UngroupingOrigin
+
+// The other side of that alignment: it repeats each key's whole group in
every slot of it.
+case object REPLICATED_FOR_JOIN extends UngroupingOrigin
+
+// A spread an operator's ordering requirement is read through, which pairs
with nothing.
+case object SPLIT_FOR_ORDERING extends UngroupingOrigin
+
/**
* The physical layout of the partitions a [[KeyedPartitioning]] describes,
which is everything
* about them except the expressions naming them.
@@ -487,13 +501,23 @@ case class CoalescedNullAwareHashPartitioning(
* one, clears the markers first. That is
precision, not the
* guarantee: without it the merge would
spread the spurious
* marker onto the accurate side.
+ * @param ungroupingOrigin Why the keys repeat on purpose, when a producer
built them that way
+ * rather than leaving them for the grouping node that
would settle them.
+ * A clustering refuses an ungrouped layout whatever
this says: on its
+ * own such a layout never holds a key's rows
together. The two
+ * alignment roles are read only where a pair is judged
+ * (`KeyedPartitioning.pairsUngrouped`). An ordering
pairs nothing,
+ * which `None` and [[SPLIT_FOR_ORDERING]] answer.
+ * See [[UngroupingOrigin]].
*/
+
case class KeyLayout(
@transient partitionKeys: Seq[InternalRowComparableWrapper],
dataTypes: Seq[DataType],
isGrouped: Boolean,
isCollapsed: Boolean,
- mayContainUnknownPartitionKeys: Boolean = false) {
+ mayContainUnknownPartitionKeys: Boolean = false,
+ ungroupingOrigin: Option[UngroupingOrigin] = None) {
// The rows carry the types they are compared at, so the pair is checked
against itself rather
// than argued about. One row is enough: a layout's rows come from one
wrapper factory, and the
@@ -565,13 +589,18 @@ case class KeyLayout(
*
* == Distribution Satisfaction and Grouping ==
* Besides the default `satisfies()`, `KeyedPartitioning` answers a family of
questions. They differ
- * in what they let happen to the data before the distribution counts as met.
Only
* `keysMaySatisfy()` is asked from outside the class. The rest build it and
`satisfies()` up.
+ * A layout a producer left ungrouped says so on the layout itself, and
`satisfies()` reads it
+ * there (`UngroupingOrigin`).
*
* - `keysSatisfy()`: do the keys as they stand co-locate every cluster key,
with nothing left for
- * a node to project away? This is the strict question, and `satisfies()` is
it plus `isGrouped`.
- * Strict for a `ClusteredDistribution`; its `OrderedDistribution` arm is a
gate on what may claim
- * a global ordering, and the order itself is still the caller's to check.
+ * a node to project away? This is the strict question, and `satisfies()` is
it plus `isGrouped`
+ * under a `ClusteredDistribution`, which an ungrouped layout serves only
through the pairing
+ * its producer's stamp vouches for (`pairsUngrouped()`), never on its own.
Its
+ * `OrderedDistribution` arm is a gate on what may claim a global ordering:
an ordering reads
+ * the keys without pairing anything, so settling is not asked, but the keys
do have to ascend
+ * (`keysSortedFor()`), and the two pairing roles are refused, built as they
are for a
+ * clustering.
* - `keysCanSatisfy()`: `keysSatisfy()`, or the keys co-locate them after a
node has projected
* away the expressions that carry none. Only
* `spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys`
admits the second half.
@@ -678,6 +707,7 @@ case class KeyedPartitioning(
def isGrouped: Boolean = layout.isGrouped
def isCollapsed: Boolean = layout.isCollapsed
def mayContainUnknownPartitionKeys: Boolean =
layout.mayContainUnknownPartitionKeys
+ def ungroupingOrigin: Option[UngroupingOrigin] = layout.ungroupingOrigin
/** This partitioning over a changed layout, e.g.
`withLayout(_.copy(isGrouped = false))`. */
def withLayout(f: KeyLayout => KeyLayout): KeyedPartitioning = copy(layout =
f(layout))
@@ -696,7 +726,8 @@ case class KeyedPartitioning(
* string is the only place a reader sees why a marked partitioning still
shuffles.
*/
override protected def stringArgs: Iterator[Any] =
- Iterator(expressions, partitionKeys, isGrouped, isCollapsed,
mayContainUnknownPartitionKeys)
+ Iterator(expressions, partitionKeys, isGrouped, isCollapsed,
+ mayContainUnknownPartitionKeys) ++ ungroupingOrigin
override protected def withNewChildrenInternal(
newChildren: IndexedSeq[Expression]): KeyedPartitioning =
@@ -795,7 +826,9 @@ case class KeyedPartitioning(
def toGrouped: KeyedPartitioning = {
// Unique keys need no dedup, only the sort.
val uniqueKeys = if (isGrouped) partitionKeys else partitionKeys.distinct
- withLayout(_.copy(partitionKeys = uniqueKeys.sorted(keyOrdering),
isGrouped = true))
+ // The keys are unique afterwards, so a reason to be ungrouped has been
spent.
+ withLayout(_.copy(partitionKeys = uniqueKeys.sorted(keyOrdering),
isGrouped = true,
+ ungroupingOrigin = None))
}
/**
@@ -855,10 +888,73 @@ case class KeyedPartitioning(
reducers: Seq[Option[KeyReducer]]): (Seq[DataType],
Seq[InternalRowComparableWrapper]) =
KeyedPartitioning.reduceKeys(partitionKeys, keyDataTypes, reducers)
- override def satisfies0(required: Distribution): Boolean = {
- super.satisfies0(required) || (isGrouped && keysSatisfy(required))
+ /**
+ * The ordering `o` reads these keys at, shared by `keysSortedFor` and
+ * `EnsureRequirements.resolveChild`, which sorts the expected keys with it.
+ *
+ * A key row holds the values of the partition expressions, so sort order
`i` reads field `i` of
+ * the row, at the type the row holds it. Binding the ordering to the
expressions' references
+ * instead would evaluate them once more on the key rows, which is right
only while every
+ * partition expression is a bare column: a join can report the other side's
join key, e.g.
+ * `100 - b`, and the keys hold that expression's values rather than `b`'s.
Both callers ask
+ * after `keysSatisfy`'s `OrderedDistribution` arm, which admits a layout
only while its
+ * expressions are the ordering's, position by position
+ * (`OrderedDistribution.areAllClusterKeysMatched`), and `resolveChild`
asserts that it holds.
+ */
+ private[sql] def keyOrderingFor(
+ o: OrderedDistribution): Ordering[InternalRowComparableWrapper] =
+ RowOrdering.create(
+ o.ordering.zip(keyDataTypes).zipWithIndex.map { case ((order, dataType),
i) =>
+ order.copy(child = BoundReference(i, dataType, nullable = true),
+ sameOrderExpressions = Seq.empty)
+ },
+ Nil).on((t: InternalRowComparableWrapper) => t.row)
+
+ /**
+ * Whether the keys ascend in the sequence `o` reads them: the half of an
ordering claim the
+ * key space does not carry. A grouped layout's keys are sorted by
construction (`toGrouped`
+ * and the sources both sort with `groupedKeyRowOrdering`), but an ungrouped
one has no
+ * producer that sorted it, so this check is what keeps `satisfies` from
answering for a
+ * layout that cannot back the answer. `EnsureRequirements.resolveChild`
asks the same
+ * question to decide keep-versus-spread, on this same construction. An
empty zip is
+ * vacuously sorted, the answer for a single key.
+ */
+ private[sql] def keysSortedFor(o: OrderedDistribution): Boolean = {
+ val keyOrdering = keyOrderingFor(o)
+ partitionKeys.zip(partitionKeys.drop(1)).forall { case (k1, k2) =>
keyOrdering.lteq(k1, k2) }
}
+ override def satisfies0(required: Distribution): Boolean = required match {
+ // An ordering reads the keys without pairing partitions, so settling them
is not what it
+ // asks: an ungrouped layout answers it on the keys alone, whether a
producer stamped it for
+ // the ordering or nothing stamped it at all. The keys, though, have to
ascend, and no
+ // producer vouches for that on an unstamped layout, so the check is here.
The two pairing
+ // roles answer no: they were built for a clustering, not for an ordering.
+ case o: OrderedDistribution =>
+ super.satisfies0(o) ||
+ (ungroupingOrigin.forall(_ == SPLIT_FOR_ORDERING) &&
+ keysSatisfy(o) && keysSortedFor(o))
+ // A clustering asks that rows sharing a cluster key share a partition,
which an ungrouped
+ // layout never gives on its own. Not even an alignment side: its claim is
about the pair it
+ // was built for, read where a pair is judged (`pairsUngrouped`) and not
here.
+ case _ =>
+ super.satisfies0(required) || (isGrouped && keysSatisfy(required))
+ }
+
+ /**
+ * Whether this layout may pair as it stands despite being ungrouped: a
producer stamped it as
+ * one side of an alignment, and the pairing being judged is the one the
stamp vouches for. The
+ * caller holds the other side, and `KeyedShuffleSpec.isCompatibleWith` is
where the two roles
+ * are asked to complement. `satisfies` refuses this layout, and rightly so:
on its own it does
+ * not hold a key's rows together, so every consumer but that pairing needs
the node that
+ * settles it. The count is asked here as it is asked there, since this is
an admission of the
+ * member as it stands.
+ */
+ private[sql] def pairsUngrouped(required: ClusteredDistribution): Boolean =
+ ungroupingOrigin.exists(r => r == SPLIT_FOR_JOIN || r ==
REPLICATED_FOR_JOIN) &&
+ required.requiredNumPartitions.forall(_ == numPartitions) &&
+ keysSatisfy(required)
+
/**
* The positions of the partition expressions that cover a cluster key of
`required`, and so
* have to survive a projection. An expression covers one in two ways:
@@ -1096,7 +1192,9 @@ object KeyedPartitioning {
*/
def concat(kps: Seq[KeyedPartitioning]): KeyedPartitioning = {
val concatenatedKeys = kps.flatMap(_.partitionKeys)
+ val ungroupingOrigin =
PartitioningCollection.concatUngrouping(kps.map(_.ungroupingOrigin))
kps.head.withLayout(_.copy(
+ ungroupingOrigin = ungroupingOrigin,
partitionKeys = concatenatedKeys,
// A child that has duplicates of its own puts them in the concatenation
too, which answers
// this without walking the keys.
@@ -1373,8 +1471,9 @@ case class PartitioningCollection(partitionings:
Seq[Partitioning])
// operation's still can, through the projection a `GroupPartitionsExec`
performs. That is the
// admission set this filter had before `satisfies` became strict, up to
one shape it now also
// keeps, a partition expression that *is* a cluster key, which
`areKeysCompatible` turns away
- // anyway. The set matters because `ValidateRequirements` builds a spec
from a finished plan
- // through here.
+ // anyway. The set matters because the planner's `shuffleToCoPartition`
reads it through here to
+ // pick the layout the other children are laid out on, and a member this
filter drops is one the
+ // collection cannot offer as that layout.
//
// Every admitted member stays, because `isCompatibleWith` answers for any
of them and the
// collection cannot know which one the other side matched.
@@ -1412,16 +1511,31 @@ object PartitioningCollection {
private[sql] def numKeyedPartitions(partitioning: Partitioning): Option[Int]
=
representativeOf(partitioning).map(_.numPartitions)
+ /**
+ * The origin every one of `origins` carries, for a carrier merging them
into one layout: the
+ * claim survives a merge only where all members make the same one. A
disagreement, or no member
+ * at all, agrees on nothing.
+ */
+ private[sql] def concatUngrouping(
+ origins: Seq[Option[UngroupingOrigin]]): Option[UngroupingOrigin] =
+ origins.distinct match {
+ case Seq(role) => role
+ case _ => None
+ }
+
/**
* Whether `p` can serve `required` once a `GroupPartitionsExec` has
projected a
* [[KeyedPartitioning]]'s keys down to the cluster keys. `satisfies` asks
whether it serves as it
* stands, and only a keyed partitioning answers the two differently.
*
* A partitioning that is not grouped is not admitted, even though a node
would also group it.
- * This is the admission set `satisfies` gave the one caller before it
became strict, and the
- * caller feeds `ValidateRequirements` as well as the planner, so it does
not widen what a
- * finished plan is checked against.
`EnsureRequirements.createKeyedShuffleSpecs` is where the
- * planner asks the wider question, for a child it is about to group itself.
+ * This is the admission set `satisfies` gave the caller before it became
strict, and it is what
+ * `PartitioningCollection.createShuffleSpec` answers with, which the
planner's
+ * `shuffleToCoPartition` reads to lay the other children out on what this
one reports.
+ * `ValidateRequirements` does not read it: `specsForPairing` is its
admission, and it admits an
+ * ungrouped member only where partially clustered distribution reports one.
+ * `EnsureRequirements.createKeyedShuffleSpecs` is where the planner asks
the wider question, for
+ * a child it is about to group itself.
*
* The partition count clause is here for consistency with the strict
question. A node changes the
* count, so the pre-grouping one is no prediction of it.
@@ -1435,6 +1549,61 @@ object PartitioningCollection {
case other => other.satisfies(required)
}
+ /**
+ * The specs `p` offers for `distribution`, one per member that may serve
it, each of them the
+ * layout that member reports (`reportedSpecOf`): the partitions it has, in
the order it reports
+ * them, with the partition expressions that carry no cluster key left out.
Leaving those out is
+ * what lets two members pair when the operation clusters on part of what
they are partitioned
+ * on, and it invents nothing: a kept key is that partition's own, the count
is the member's own,
+ * and the order is the member's own.
+ *
+ * A keyed member is admitted as it stands: on `satisfies`, the strict
question, or, since a
+ * clustering is what a pair owes together, on the alignment stamp that
vouches for exactly
+ * this pairing (`KeyedPartitioning.pairsUngrouped`; `isCompatibleWith` then
asks the two
+ * sides' roles to complement). A member that is not keyed is asked on
`satisfies` alone and
+ * has no projection to make. Compare the planner's admission of a member
+ * (`EnsureRequirements.createKeyedShuffleSpecs`): it asks the wider
`keysMaySatisfy`, because
+ * the node it inserts settles the member afterwards, and it additionally
requires the coverage
+ * of every operation key (`spark.sql.requireAllClusterKeysForCoPartition`),
which this drops:
+ * a skew heuristic, since a member whose partitioning keys cover only a
subset of the
+ * operation's keys is a sound pairing.
+ */
+ private[sql] def specsForPairing(
+ p: Partitioning,
+ distribution: ClusteredDistribution): Seq[ShuffleSpec] =
+ flatten(p).flatMap {
+ case k: KeyedPartitioning =>
+ Option.when(k.satisfies(distribution) ||
k.pairsUngrouped(distribution))(
+ reportedSpecOf(k, distribution))
+ case other =>
+
Option.when(other.satisfies(distribution))(other.createShuffleSpec(distribution))
+ }
+
+ /**
+ * The spec for a member a finished plan reports: its own layout with the
partition expressions
+ * that carry no cluster key left out. No key is deduped and no key is
re-sorted, so a side offers
+ * a view of the partitions the plan has rather than the layout
+ * `KeyedPartitioning.createShuffleSpec` builds for a node a planner would
insert: that one dedups
+ * and sorts, and its count is the grouped one, which the plan does not hold.
+ *
+ * A projection here is a relabelling of the member's partitions, and this
leans on the caller for
+ * that: only `keysSatisfy`'s projecting half admits such a member, and it
holds the projection to
+ * one that merges no partition. The projecting half excludes a marked
member, so a marked one is
+ * admitted only with its keys covering the clustering structurally, and the
position count below
+ * comes out full for it: it leaves here unprojected, which its claim needs,
since it routes the
+ * rows of an undeclared key by a hash over these keys in this order and
dropping one breaks it.
+ */
+ private def reportedSpecOf(
+ k: KeyedPartitioning,
+ distribution: ClusteredDistribution): KeyedShuffleSpec = {
+ val positions = k.positionsCoveringClusterKeys(distribution).toSeq
+ if (positions.size == k.expressions.length) {
+ KeyedShuffleSpec(k, distribution)
+ } else {
+ KeyedShuffleSpec(k.project(positions), distribution, Some(positions))
+ }
+ }
+
/**
* The uniform `mayContainUnknownPartitionKeys` marker of `p`'s keyed
members, read from its
* first keyed member, the same one `checkKeyedPartitioningInvariant`
compares against.
@@ -1469,11 +1638,17 @@ object PartitioningCollection {
// already holds the canonical layout is returned as it is. That is what
keeps repeated
// `outputPartitioning` computations over deeply nested collections (e.g.
chains of same-key
// joins) O(1) per level.
+ // Interning below retypes the members that disagree with the agreed
origin, and a retyped
+ // claim would be one nothing built.
+ val concatUngrouping =
+
PartitioningCollection.concatUngrouping(representatives.map(_.layout.ungroupingOrigin))
val canonicalLayout = representatives.map(_.layout)
.find(l => l.isCollapsed == anyCollapsed &&
- l.mayContainUnknownPartitionKeys == anyUnknownKeys)
+ l.mayContainUnknownPartitionKeys == anyUnknownKeys &&
+ l.ungroupingOrigin == concatUngrouping)
.orElse(representatives.headOption.map(_.layout.copy(
- isCollapsed = anyCollapsed, mayContainUnknownPartitionKeys =
anyUnknownKeys)))
+ isCollapsed = anyCollapsed, mayContainUnknownPartitionKeys =
anyUnknownKeys,
+ ungroupingOrigin = concatUngrouping)))
.orNull
def intern(p: Partitioning): Partitioning = representativeOf(p) match {
@@ -1645,12 +1820,14 @@ case object SinglePartitionShuffleSpec extends
LeafShuffleSpec {
// disagree when the subset config projects them onto different key sets.
//
// `EnsureRequirements` never reaches this: a spec whose
`canCreatePartitioning` is false is
- // never the best one. The one production caller that can put a collection
on the `other` side
- // is `ValidateRequirements`'
`specs.tail.forall(_.isCompatibleWith(specs.head))`, and there the
- // stricter answer is the safer one. It does leave this direction stricter
than the collection's
- // own `exists`, against the symmetry this trait's doc assumes, but
nothing observable follows:
- // only `KeyedShuffleSpec` can make members disagree on `numPartitions`,
and it has no
- // `SinglePartitionShuffleSpec` case, so that direction is already false.
+ // never the best one. Nor does any other production caller put a
collection on the `other`
+ // side: `pickCoPartitionTarget` flattens before it pairs a member, and
`ValidateRequirements`
+ // offers leaf specs and pairs those, where it used to compare every child
against the first
+ // one's whole spec. So what is left is the answer `ShuffleSpecSuite`
pins. It stays `forall`
+ // rather than the `exists` a collection answers with, against the
symmetry the trait doc
+ // assumes, but nothing observable follows: only `KeyedShuffleSpec` can
make members disagree on
+ // `numPartitions`, and it has no `SinglePartitionShuffleSpec` case, so
that direction is
+ // already false.
case ShuffleSpecCollection(specs) => specs.forall(isCompatibleWith)
}
@@ -1930,6 +2107,13 @@ case class IdentityReducer(transform:
TransformExpression) extends Reducer[Any,
* does not: one `keyPositions` entry is empty there,
which
* `canCreatePartitioning` and `areKeysCompatible`
both turn away. See
* `KeyedPartitioning.createShuffleSpec`.
+ * The field's other producer is
`PartitioningCollection.reportedSpecOf`,
+ * and there `Some` records a projection of the
*report*: a view of the
+ * partitions the child already holds, which no node
is owed. Such a spec
+ * is only ever read by `isCompatibleWith`, which
ignores the field, so
+ * the "insert or compose a grouping node" reading the
+ * `EnsureRequirements` consumers give `Some` stays
about planner-produced
+ * specs alone.
*/
case class KeyedShuffleSpec(
partitioning: KeyedPartitioning,
@@ -1985,19 +2169,30 @@ case class KeyedShuffleSpec(
// Here we check:
// 1. both distributions have the same number of clustering keys
// 2. both partitioning have the same number of partitions
- // 3. partition expressions from both sides are compatible, which means:
- // 3.1 both sides have the same number of partition expressions
- // 3.2 for each pair of partition expressions at the same index, the
corresponding
+ // 3. where either side is ungrouped on purpose, the pair is one of each:
the side that
+ // spreads a key's splits (`SPLIT_FOR_JOIN`) and the side that repeats
the group
+ // (`REPLICATED_FOR_JOIN`). An unstamped ungrouped layout serves no
clustering, so only
+ // `compatibleAsIs` can offer one here, and its builder settles the
side with a grouping
+ // node.
+ // 4. partition expressions from both sides are compatible, which means:
+ // 4.1 both sides have the same number of partition expressions
+ // 4.2 for each pair of partition expressions at the same index, the
corresponding
// partition keys must share overlapping positions in their
respective clustering keys.
- // 3.3 each pair of partition expressions at the same index must
describe one key space: two
+ // 4.3 each pair of partition expressions at the same index must
describe one key space: two
// bare references, the same transform function, or two transforms
reduced through the
// same pairing. A transform of an expression counts only against
one over the same
// argument shape. A pair the join would first reduce onto one key
space does not count,
// see `areKeysCompatible`'s `allowReduce`.
- // 4. the partition values from both sides are following the same order.
+ // 5. the partition values from both sides are following the same order.
case otherSpec @ KeyedShuffleSpec(otherPartitioning, otherDistribution, _)
=>
distribution.clustering.length == otherDistribution.clustering.length &&
numPartitions == otherSpec.numPartitions &&
+ ((partitioning.ungroupingOrigin, otherPartitioning.ungroupingOrigin)
match {
+ case (None, None) => true
+ case (Some(l), Some(r)) =>
+ Set(l, r) == Set(SPLIT_FOR_JOIN, REPLICATED_FOR_JOIN)
+ case _ => false
+ }) &&
areKeysCompatible(otherSpec, allowReduce = false) &&
// The reason the types are asked as well as the rows is on
`describesSameKeys`, so the
// next site comparing keys cannot forget the type clause.
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
index a73235118c57..d300ee7db61c 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst
import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException}
import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference, DirectShufflePartitionID, Expression, GetStructField,
TransformExpression}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute,
AttributeReference, DirectShufflePartitionID, Expression, GetStructField,
SortOrder, TransformExpression}
import org.apache.spark.sql.catalyst.plans.SQLHelper
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.connector.catalog.functions.{FlipLowBitFunction,
Reducer, ReducibleFunction, ScalarFunction}
@@ -93,9 +93,8 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
test("SPARK-59289: createShuffleSpec drops a keyed member that is not
grouped") {
val a = AttributeReference("a", IntegerType)()
val clustered = ClusteredDistribution(Seq(a))
- // A node would group this one too, but the admission set here is what a
finished plan is
- // checked against by `ValidateRequirements`, so it stays what the strict
question admitted
- // before a projection was allowed to answer it.
+ // A node would group this one too, but a collection is a layout the plan
holds, so it stays
+ // what the strict question admitted before a projection was allowed to
answer it.
val ungrouped = KeyedPartitioning(Seq(a), Seq(InternalRow(1),
InternalRow(1), InternalRow(2)))
val grouped = KeyedPartitioning(Seq(a), Seq(InternalRow(1),
InternalRow(2), InternalRow(3)))
assert(!ungrouped.isGrouped && grouped.isGrouped, "test setup")
@@ -114,6 +113,143 @@ class ShuffleSpecSuite extends SparkFunSuite with
SQLHelper {
assert(specs.head.isInstanceOf[HashShuffleSpec], s"and it is the hash one,
got $specs")
}
+ test("SPARK-59671: the specs a side offers are the layouts it reports") {
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val cd = ClusteredDistribution(Seq(a))
+ // An ungrouped member is offered where its layout says where the
ungrouping comes from: one
+ // side of an alignment keeps its splits and spreads them. A finished plan
is judged on the
+ // layouts it holds, so the reason is read off them rather than off a
configuration.
+ val spread = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1),
InternalRow(2)))
+ .withLayout(_.copy(ungroupingOrigin = Some(SPLIT_FOR_JOIN)))
+ val specs = PartitioningCollection.specsForPairing(spread, cd)
+ assert(specs.size == 1, s"one member, one spec, got $specs")
+ val spec = specs.head.asInstanceOf[KeyedShuffleSpec]
+ assert((spec.partitioning eq spread) && spec.joinKeyPositions.isEmpty,
+ s"its own layout, no projection to make, got $specs")
+
+ // A member that does not say why it is ungrouped is waiting for the node
that would settle it,
+ // and a finished plan has none left to insert.
+ val unsettled = spread.withLayout(_.copy(ungroupingOrigin = None))
+ assert(PartitioningCollection.specsForPairing(unsettled, cd).isEmpty,
+ "an ungrouped member is a plan only where a producer built one")
+
+ // An origin is read for the distribution it is good for. A spread an
ordering requirement is
+ // read through is no side of a clustering, and a clustering is what asks
here.
+ val forOrdering = spread.withLayout(_.copy(ungroupingOrigin =
Some(SPLIT_FOR_ORDERING)))
+ assert(PartitioningCollection.specsForPairing(forOrdering, cd).isEmpty,
+ "an origin a clustering does not read is not offered to it")
+
+ // The side that repeats a key's whole group is the other half of that
alignment, and it is
+ // offered the same way.
+ val repeat = spread.withLayout(_.copy(ungroupingOrigin =
Some(REPLICATED_FOR_JOIN)))
+ assert(PartitioningCollection.specsForPairing(repeat, cd).size == 1)
+
+ // A count the operation pinned is asked as it stands, which `satisfies`
asks for every member,
+ // the ones whose layout says why they are ungrouped included.
+ val four = KeyedPartitioning(Seq(a), (1 to 4).map(InternalRow(_)))
+ .withLayout(_.copy(ungroupingOrigin = Some(SPLIT_FOR_JOIN)))
+ assert(PartitioningCollection
+ .specsForPairing(four, ClusteredDistribution(Seq(a),
requiredNumPartitions = Some(4)))
+ .size == 1, "the size the operation asks for is the one the member
reports")
+ assert(PartitioningCollection
+ .specsForPairing(four, ClusteredDistribution(Seq(a),
requiredNumPartitions = Some(3)))
+ .isEmpty, "a member of another size does not serve the distribution")
+
+ // The subset permission applies where the operation's keys are a subset
of the member's
+ // partitioning keys, so a member may carry an expression the operation
does not cluster on. It
+ // is offered as its own partitions under the key the operation clusters
on: that expression is
+ // left out, and no key is deduped or re-sorted for it.
+ withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key
-> "true") {
+ val source = KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 1),
InternalRow(2, 2)))
+ val offered = PartitioningCollection.specsForPairing(source, cd)
+ .head.asInstanceOf[KeyedShuffleSpec]
+ assert(offered.partitioning.expressions === Seq(a) &&
+ offered.partitioning.numPartitions === source.numPartitions &&
+ offered.joinKeyPositions === Some(Seq(0)),
+ s"the member's own partitions under the operation's key, got $offered")
+
+ // A marked layout is never offered through a projection: the projecting
half of `keysSatisfy`
+ // excludes a marked member, so the keys an offered marked one carries
are the keys its claim
+ // is over.
+ val marked = source.withLayout(_.copy(mayContainUnknownPartitionKeys =
true))
+ assert(PartitioningCollection.specsForPairing(marked, cd).isEmpty,
+ "a marked member is not offered through a projection")
+ }
+
+ // A member whose keys do not cover the clustering is not offered, and a
member that is not
+ // keyed is asked for its own spec.
+ val elsewhere = KeyedPartitioning(Seq(b), Seq(InternalRow(1),
InternalRow(2)))
+ assert(PartitioningCollection.specsForPairing(elsewhere, cd).isEmpty)
+ assert(PartitioningCollection
+ .specsForPairing(HashPartitioning(Seq(a), 2),
cd).head.isInstanceOf[HashShuffleSpec])
+ }
+
+ test("SPARK-59671: carriers keep one agreed origin and grouping spends it") {
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val keys = Seq(InternalRow(1), InternalRow(1), InternalRow(2))
+ def stamped(attr: AttributeReference, role: UngroupingOrigin):
KeyedPartitioning =
+ KeyedPartitioning(Seq(attr), keys).withLayout(_.copy(ungroupingOrigin =
Some(role)))
+ val spread = stamped(a, SPLIT_FOR_JOIN)
+ val repeat = stamped(a, REPLICATED_FOR_JOIN)
+
+ // Grouping spends the claim: the keys are unique afterwards, so there are
no splits left to
+ // spread and no group left to repeat, and a kept stamp would turn a sound
pair of grouped
+ // sides away in `isCompatibleWith`, which reads a stamp as one side of an
ungrouped pair.
+ assert(spread.toGrouped.isGrouped &&
spread.toGrouped.ungroupingOrigin.isEmpty)
+
+ // A concatenation stands for the claim all of its members make: one role
survives it, and
+ // two roles stand for nothing.
+ assert(KeyedPartitioning.concat(Seq(spread, spread)).ungroupingOrigin
+ .contains(SPLIT_FOR_JOIN))
+ assert(KeyedPartitioning.concat(Seq(spread,
repeat)).ungroupingOrigin.isEmpty)
+
+ // A collection holds one canonical layout its members share, and it
agrees the same way. An
+ // aligned pair's joined output is the disagreeing case: the zip of a
spread side and a
+ // repeating side is a layout neither claim describes, so it comes out
inert.
+ val agreed = PartitioningCollection.fromPartitionings(
+ Seq(spread, stamped(b, SPLIT_FOR_JOIN)))
+ assert(agreed.partitionings.forall(p => p.asInstanceOf[KeyedPartitioning]
+ .ungroupingOrigin.contains(SPLIT_FOR_JOIN)), "one role, one canonical
layout")
+ val joined = PartitioningCollection.fromPartitionings(Seq(spread, repeat))
+ assert(joined.partitionings.forall(_.asInstanceOf[KeyedPartitioning]
+ .ungroupingOrigin.isEmpty), "two roles agree on nothing")
+ }
+
+ test("SPARK-59671: an ordering reads sorted keys, a clustering reads the
pairing") {
+ val a = AttributeReference("a", IntegerType)()
+ def layout(keys: Seq[Int], role: Option[UngroupingOrigin]):
KeyedPartitioning =
+ KeyedPartitioning(Seq(a), keys.map(InternalRow(_)))
+ .withLayout(_.copy(ungroupingOrigin = role))
+ val sorted = Seq(1, 1, 2)
+ val ordering = OrderedDistribution(Seq(SortOrder(a, Ascending)))
+ val clustered = ClusteredDistribution(Seq(a))
+
+ withSQLConf(SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
+ // An ordering pairs no partition with another, so it asks no settling,
but it does ask
+ // the keys to ascend, which a grouped layout has by construction and an
ungrouped one
+ // only where something sorted it.
+ assert(layout(sorted, None).satisfies(ordering))
+ assert(layout(sorted, Some(SPLIT_FOR_ORDERING)).satisfies(ordering))
+ assert(!layout(Seq(1, 2, 1, 2), None).satisfies(ordering),
+ "an ungrouped layout nothing sorted does not claim an ordering")
+ assert(!layout(Seq(2, 1), None).satisfies(ordering),
+ "a grouped layout whose keys do not ascend does not claim an ordering")
+ // The two pairing roles were built for a clustering, not for an
ordering.
+ assert(!layout(sorted, Some(SPLIT_FOR_JOIN)).satisfies(ordering))
+ assert(!layout(sorted, Some(REPLICATED_FOR_JOIN)).satisfies(ordering))
+ // A clustering asks that a key's rows share a partition, which no
ungrouped layout gives
+ // on its own: an alignment side's stamp is read where the pair is judged
+ // (`specsForPairing`, pinned above), not here.
+ Seq(None, Some(SPLIT_FOR_ORDERING), Some(SPLIT_FOR_JOIN),
Some(REPLICATED_FOR_JOIN)).foreach {
+ role => assert(!layout(sorted, role).satisfies(clustered))
+ }
+ assert(layout(Seq(1, 2), None).satisfies(clustered),
+ "a grouped layout answers a clustering as it stands")
+ }
+ }
+
protected def checkCompatible(
left: ShuffleSpec,
right: ShuffleSpec,
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
index 2cdc6b9d410b..9f4ceb93dab5 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
@@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.codegen.GenerateOrdering
import org.apache.spark.sql.catalyst.plans.QueryPlan
-import org.apache.spark.sql.catalyst.plans.physical.{IdentityReducer,
KeyedPartitioning, KeyLayout, KeyReducer, Partitioning, PartitioningCollection,
UnknownPartitioning}
+import org.apache.spark.sql.catalyst.plans.physical.{IdentityReducer,
KeyedPartitioning, KeyLayout, KeyReducer, Partitioning, PartitioningCollection,
REPLICATED_FOR_JOIN, UngroupingOrigin, UnknownPartitioning}
import org.apache.spark.sql.catalyst.util.{truncatedString,
InternalRowComparableWrapper}
import org.apache.spark.sql.execution.{SafeForKWayMerge, SparkPlan,
SQLExecution, UnaryExecNode}
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
@@ -69,9 +69,13 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
* reducers rather than as their display names. A name embeds
the exprId of the
* attribute it transforms, so it would break canonical
equality the same way the
* reducers do. See `doCanonicalize`.
- * @param distributePartitions When true, splits for a key are distributed
across the expected
- * partitions (padding with empty partitions).
When false, all splits
- * are replicated to every expected partition for
that key.
+ * @param ungroupingOrigin Why this node's output is left ungrouped on
purpose, stamped into
+ * the layout it reports. It also decides the routing:
+ * `SPLIT_FOR_JOIN` and `SPLIT_FOR_ORDERING`
distribute each key's
+ * splits across the expected slots (padding with
empty partitions),
+ * [[REPLICATED_FOR_JOIN]] groups the splits first and
repeats each group
+ * into its expected slots, and `None` groups them,
one partition per key.
+ * See [[UngroupingOrigin]].
* @param enableSortedMerge When true, uses [[SortedMergeCoalescedRDD]] to
perform a k-way merge
* of the coalesced partitions, preserving the
child's output ordering
* end-to-end. Set by [[EnsureRequirements]] when a
parent operator
@@ -89,10 +93,16 @@ case class GroupPartitionsExec(
@transient joinKeyPositions: Option[Seq[Int]],
@transient expectedKeyCount: Option[Int],
@transient reducers: Option[Seq[Option[KeyReducer]]],
- @transient distributePartitions: Boolean,
+ @transient ungroupingOrigin: Option[UngroupingOrigin],
@transient enableSortedMerge: Boolean
) extends UnaryExecNode {
+ /**
+ * Whether each key's splits are distributed across its expected slots, or
grouped into one
+ * partition per key first. Derived from the origin; see `@param
ungroupingOrigin`.
+ */
+ def distributePartitions: Boolean = ungroupingOrigin.exists(_ !=
REPLICATED_FOR_JOIN)
+
/**
* The layout this node was planned to produce, but only while its child
still reports the one it
* was planned over.
@@ -451,11 +461,15 @@ case class GroupPartitionsExec(
return None
}
val positions =
joinKeyPositions.fold(plannedInNewChild)(_.map(plannedInNewChild))
+ // Only a grouping node reaches here today: the one caller
(`CombineAdjacentAggregation`)
+ // regroups the node between two aggregates, which is built unstamped,
so forwarding the
+ // origin preserves rather than claims. A stamped node would also need
its expected slots
+ // re-derived, which this path does not do.
val regrouped = GroupPartitionsExec(
child = newChild,
joinKeyPositions = Option.when(positions !=
newChildKp.expressions.indices)(positions),
reducers = reducers,
- distributePartitions = distributePartitions,
+ ungroupingOrigin = ungroupingOrigin,
enableSortedMerge = enableSortedMerge)
regrouped.copyTagsFrom(this)
Some(regrouped)
@@ -523,7 +537,7 @@ private[sql] object GroupPartitionsExec {
*
* **Both are derived, and neither `copy` nor the generated `apply`
re-derives them**, so a change
* to `child`, `joinKeyPositions`, `expectedPartitionKeys` (stored as
`expectedKeyCount`),
- * `reducers` or `distributePartitions` has to come back through here.
`enableSortedMerge` is not
+ * `reducers` or `ungroupingOrigin` has to come back through here.
`enableSortedMerge` is not
* an input to either, which is why `tryEnableSortedMerge` may `copy` it.
*
* Two other `copy` calls in this file are deliberate.
`withNewChildInternal` carries both fields
@@ -537,7 +551,7 @@ private[sql] object GroupPartitionsExec {
joinKeyPositions: Option[Seq[Int]] = None,
expectedPartitionKeys: Option[Seq[(InternalRowComparableWrapper, Int)]]
= None,
reducers: Option[Seq[Option[KeyReducer]]] = None,
- distributePartitions: Boolean = false,
+ ungroupingOrigin: Option[UngroupingOrigin] = None,
enableSortedMerge: Boolean = false): GroupPartitionsExec = {
// There must be a `KeyedPartitioning` in the child's output partitioning,
as a
// `GroupPartitionsExec` node is added to a plan only in that case.
@@ -548,11 +562,11 @@ private[sql] object GroupPartitionsExec {
// `representativeOf` found one, so the partitioning is an expression tree.
val childExpr = childPartitioning.asInstanceOf[Partitioning with
Expression]
val grouping = computeGrouping(
- childKp, joinKeyPositions, expectedPartitionKeys, reducers,
distributePartitions)
+ childKp, joinKeyPositions, expectedPartitionKeys, reducers,
ungroupingOrigin)
GroupPartitionsExec(child, grouping,
- computeOutputPartitioning(childExpr, grouping, joinKeyPositions,
reducers), childPartitioning,
- joinKeyPositions, expectedPartitionKeys.map(_.size), reducers,
distributePartitions,
- enableSortedMerge)
+ computeOutputPartitioning(childExpr, grouping, joinKeyPositions,
reducers),
+ childPartitioning, joinKeyPositions, expectedPartitionKeys.map(_.size),
reducers,
+ ungroupingOrigin, enableSortedMerge)
}
/**
@@ -585,13 +599,14 @@ private[sql] object GroupPartitionsExec {
// commits to a pairing and falls back to a shuffle when it happens.
//
// The marker read off `p` is the same bit the `else` branch publishes
through
- // `grouping.layout`, since `computeGrouping` copies the child's layout
without touching it. A
- // change to that copy has to move this guard with it.
+ // `grouping.layout`, since `computeGrouping` copies the child's layout
and carries the
+ // marker over. A change to that copy has to move this guard with it.
if (PartitioningCollection.keyedMarkerOf(p).contains(true) &&
!grouping.isIdentity) {
UnknownPartitioning(grouping.partitions.size)
} else {
// One instance for every member, so they share it by reference. It
already carries the
- // child's marker, and the guard above is what lets that carry over.
+ // child's marker, and the guard above is what lets that carry over, and
the origin this
+ // node's producer stamped, which `computeGrouping` wrote into it.
val layout = grouping.layout
p.transform {
case k: KeyedPartitioning =>
@@ -687,7 +702,9 @@ private[sql] object GroupPartitionsExec {
joinKeyPositions: Option[Seq[Int]],
expectedPartitionKeys: Option[Seq[(InternalRowComparableWrapper, Int)]],
reducers: Option[Seq[Option[KeyReducer]]],
- distributePartitions: Boolean): PartitionGrouping = {
+ ungroupingOrigin: Option[UngroupingOrigin]): PartitionGrouping = {
+ // The routing follows the origin, the same derivation the node itself
reports.
+ val distributePartitions = ungroupingOrigin.exists(_ !=
REPLICATED_FOR_JOIN)
// Project partition keys if join key positions are specified
val (projectedDataTypes, projectedKeys) =
joinKeyPositions.fold(
@@ -747,16 +764,21 @@ private[sql] object GroupPartitionsExec {
case ((_, Seq(single)), outputIndex) => single == outputIndex
case _ => false
}
- // A `copy` of the child's layout rather than a fresh one, so the marker
and anything the layout
- // grows later are carried without this having to name them. Built here,
where the child is in
- // scope, so `computeOutputPartitioning` has nothing left to re-apply.
+ // A `copy` of the child's layout rather than a fresh one, so the marker
and anything the
+ // layout grows later are carried without this having to name them. The
origin is the one
+ // field not carried: the output claims what this node's own producer
stamped, never what
+ // the child's claimed, and only while the shape holds the claim -- a
grouping that settles
+ // every key leaves nothing repeating, so it spends the claim the way
`toGrouped` does.
+ // Built here, where the child is in scope, so `computeOutputPartitioning`
has nothing left
+ // to re-apply.
PartitionGrouping(
partitions,
childKp.layout.copy(
partitionKeys = partitions.map(_._1),
dataTypes = reducedDataTypes,
isGrouped = isGrouped,
- isCollapsed = isCollapsed),
+ isCollapsed = isCollapsed,
+ ungroupingOrigin = if (isGrouped) None else ungroupingOrigin),
isIdentity, numPrunedPartitions, numReplicatedPartitionReads)
}
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
index af186f4d0a62..0480b38c7ec9 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
@@ -136,8 +136,8 @@ case class EnsureRequirements(
case Some(resolution) =>
(distribution, resolution) match {
case (o: OrderedDistribution, _) =>
- // OrderedDistribution requires grouped KeyedPartitioning with
sorted keys
- // according to the distribution's ordering.
+ // OrderedDistribution requires a KeyedPartitioning with keys
sorted according to
+ // the distribution's ordering.
val satisfyingKeyedPartitioning = resolution.fold(identity, _._1)
// A key row holds the values of the partition expressions.
`keysSatisfy` admits a
// partitioning here only when those expressions are the
ordering's, position by
@@ -146,26 +146,18 @@ case class EnsureRequirements(
// `100 - b`.
assert(o.areAllClusterKeysMatched(satisfyingKeyedPartitioning.expressions),
"the partition expressions must be the ordering's, position by
position")
- val keyRowOrdering = RowOrdering.create(
-
o.ordering.zip(satisfyingKeyedPartitioning.keyDataTypes).zipWithIndex.map {
- case ((order, dataType), i) =>
- order.copy(child = BoundReference(i, dataType, nullable =
true),
- sameOrderExpressions = Seq.empty)
- },
- Nil)
- val keyOrdering = keyRowOrdering.on((t:
InternalRowComparableWrapper) => t.row)
- val keys = satisfyingKeyedPartitioning.partitionKeys
- // An empty zip is vacuously sorted, which is the answer for a
single key.
- if (keys.zip(keys.drop(1)).forall { case (k1, k2) =>
keyOrdering.lteq(k1, k2) }) {
+ if (satisfyingKeyedPartitioning.keysSortedFor(o)) {
child
} else {
- // Use distributePartitions to spread splits across expected
partitions
- val sortedGroupedKeys = keys
+ // Spread the splits across the expected partitions, in the
ordering's sequence
+ val sortedGroupedKeys =
satisfyingKeyedPartitioning.partitionKeys
.groupBy(identity).view.mapValues(_.size)
- .toSeq.sortBy(_._1)(keyOrdering)
+
.toSeq.sortBy(_._1)(satisfyingKeyedPartitioning.keyOrderingFor(o))
GroupPartitionsExec(child,
expectedPartitionKeys = Some(sortedGroupedKeys),
- distributePartitions = true
+ // The keys stay ungrouped so that the ordering the operator
reads is the one
+ // derived from them, which is no side of a pairing.
+ ungroupingOrigin = Some(SPLIT_FOR_ORDERING)
)
}
@@ -607,6 +599,9 @@ case class EnsureRequirements(
right: SparkPlan,
rightRequired: ClusteredDistribution): Option[Seq[SparkPlan]] = {
parent match {
+ // A keyed alignment is planned for these two operators only. A
`SortMergeAsOfJoinExec` is
+ // a `ShuffledJoin` and gets none: its matches are read in the order
within a partition,
+ // which neither a spread side nor a repeating one preserves.
case smj: SortMergeJoinExec =>
checkKeyGroupCompatible(left, leftRequired, right, rightRequired,
smj.joinType)
case sj: ShuffledHashJoinExec =>
@@ -919,13 +914,25 @@ case class EnsureRequirements(
}
// Now we need to push-down the common partition information to the
`GroupPartitionsExec`s.
+ //
+ // Where `applyPartialClustering` holds, exactly one side repeats:
`replicateRightSide` is
+ // the negation of `replicateLeftSide`, and the branch above is taken
only when the side it
+ // picked may replicate for the join type. That split is the whole of
the pairing's
+ // soundness: for a key, a partition holding part of it on the spread
side holds all of it
+ // on the repeating side, so pairing the two index by index loses no
match, and
+ // `ValidateRequirements` cannot tell such a pair from two sides that
split the key between
+ // them. Each side stamps why it is left ungrouped (`UngroupingOrigin`),
so a reader of the
+ // finished plan judges the pair on the layouts it holds; the stamp also
decides the node's
+ // routing (`GroupPartitionsExec.distributePartitions`).
(
GroupPartitionsExec(rawLeft, leftSpec.joinKeyPositions,
Some(mergedPartitionKeys), leftReducers,
- distributePartitions = applyPartialClustering && !replicateLeftSide),
+ ungroupingOrigin = Option.when(applyPartialClustering)(
+ if (replicateLeftSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN)),
GroupPartitionsExec(rawRight, rightSpec.joinKeyPositions,
Some(mergedPartitionKeys), rightReducers,
- distributePartitions = applyPartialClustering &&
!replicateRightSide))
+ ungroupingOrigin = Option.when(applyPartialClustering)(
+ if (replicateRightSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN)))
}
// The pairing is only worth committing to if both children still declare
the same aligned key
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala
index 1ac6b809fd25..4336b0847bc7 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala
@@ -45,22 +45,33 @@ object ValidateRequirements extends Logging {
assert(requiredChildDistributions.length == children.length)
assert(requiredChildOrderings.length == children.length)
+ // A `ClusteredDistribution` is the one distribution an operator can owe
its children
+ // together rather than one by one, so an operator whose children all owe
one is judged on
+ // their mutual layout: the pairing below decides it, and the per-child
distribution clause
+ // is skipped for it. Nothing is lost by the skip: `specsForPairing`
admits a member only
+ // on its own answer, which is the clause's plus the alignment stamp a
producer vouched for,
+ // and a side offering nothing fails the pairing by itself. Every other
child, an operator
+ // with a single clustered child included, answers the clause for itself:
a stamped layout
+ // is one side of a pair, and a consumer that is not the pairing's other
side owes it the
+ // node that settles it. That is every such operator, not only a join: a
cogroup zips
+ // corresponding partitions too.
+ val clusteredMultiChild = children.length > 1 &&
+ requiredChildDistributions.forall(_.isInstanceOf[ClusteredDistribution])
+
val satisfied =
children.zip(requiredChildDistributions.zip(requiredChildOrderings)).forall {
case (child, (distribution, ordering))
- if !child.outputPartitioning.satisfies(distribution)
+ if (!clusteredMultiChild &&
!child.outputPartitioning.satisfies(distribution))
|| !SortOrder.orderingSatisfies(child.outputOrdering, ordering) =>
logDebug(s"ValidateRequirements failed: $distribution,
$ordering\n$plan")
false
case _ => true
}
- if (satisfied && children.length > 1 &&
-
requiredChildDistributions.forall(_.isInstanceOf[ClusteredDistribution])) {
- // Check the co-partitioning requirement.
- val specs =
children.map(_.outputPartitioning).zip(requiredChildDistributions).map {
- case (p, d) =>
p.createShuffleSpec(d.asInstanceOf[ClusteredDistribution])
- }
- if (specs.tail.forall(_.isCompatibleWith(specs.head))) {
+ if (satisfied && clusteredMultiChild) {
+ // Check the co-partitioning requirement. A pair aligned without
grouping is one each
+ // side's layout answers for (`UngroupingOrigin`), so a pair no producer
built is one
+ // whose sides offer nothing to pair and fails here.
+ if (satisfiesForPairing(children, requiredChildDistributions)) {
true
} else {
logDebug(s"ValidateRequirements failed: children not co-partitioned
in\n$plan")
@@ -70,4 +81,31 @@ object ValidateRequirements extends Logging {
satisfied
}
}
+
+ /**
+ * Whether the sides of a multi-child clustered operator line up: every side
offers the layouts
+ * it reports ([[PartitioningCollection.specsForPairing]]), and one member
of the first side
+ * pairs with every other side. A plan holds what its members report, so no
key is deduped and
+ * none is re-sorted to make a pair: a side is judged on the partitions it
has, under the key
+ * the operation clusters on.
+ *
+ * `EnsureRequirements` asks the same question of a pair it takes as it
stands by a different
+ * predicate: its `compatibleAsIs` reads two unprojected specs, while a
member here may be
+ * relabelled onto the operation's cluster key (`reportedSpecOf`), and its
`agreeingPairs` /
+ * `committed` path commits on the sides a reduce rebuilt rather than on the
pair it picked. A
+ * finished plan is asked the strict question alone, which the reduced pair
answers
+ * (`hasSameReducedKeys`). The coverage of every operation key
+ * (`spark.sql.requireAllClusterKeysForCoPartition`) is a skew heuristic and
part of neither.
+ */
+ private def satisfiesForPairing(
+ children: Seq[SparkPlan],
+ distributions: Seq[Distribution]): Boolean = {
+ val specs = children.zip(distributions).map { case (child, distribution) =>
+ PartitioningCollection.specsForPairing(
+ child.outputPartitioning,
distribution.asInstanceOf[ClusteredDistribution])
+ }
+ specs.headOption.exists { firstSide =>
+ firstSide.exists(head => specs.tail.forall(side =>
side.exists(_.isCompatibleWith(head))))
+ }
+ }
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
index fb92a17161e0..2b767c452939 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.{Ascending,
AttributeReference,
import org.apache.spark.sql.catalyst.expressions.aggregate.Complete
import org.apache.spark.sql.catalyst.plans.{Cross, ExistenceJoin, Inner,
JoinType, LeftAnti, LeftSemi, LeftSingle}
import org.apache.spark.sql.catalyst.plans.physical
-import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning
+import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning,
PartitioningCollection, REPLICATED_FOR_JOIN, SPLIT_FOR_JOIN}
import org.apache.spark.sql.connector.catalog.{Column, Identifier,
InMemoryBaseTable, InMemoryCatalystRuntimeFilterCatalog, InMemoryTable,
InMemoryTableCatalog}
import org.apache.spark.sql.connector.catalog.constraints.Constraint
import org.apache.spark.sql.connector.catalog.functions._
@@ -50,6 +50,7 @@ import org.apache.spark.sql.execution.{
SparkPlan,
UnionExec,
WholeStageCodegenExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec,
ResultQueryStageExec}
import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec,
HashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2ScanRelation, GroupPartitionsExec}
import org.apache.spark.sql.execution.exchange.{EnsureRequirements,
ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike,
ValidateRequirements}
@@ -8317,8 +8318,8 @@ class KeyGroupedPartitioningSuite
test("SPARK-59120: reduced partition keys are read at the types they were
built with") {
// The join reduces the identity side onto the year key space, so its keys
become `IntegerType`
// years while the partitioning still reports `identity(ts)`, declaring
`TimestampType`. With
- // the subset opt-in on, AQE re-runs `createShuffleSpec` on the already
reduced children through
- // `ValidateRequirements`, which projects and sorts those keys. The
mechanism is in
+ // the subset opt-in on, AQE re-runs `EnsureRequirements` over the already
reduced children,
+ // and its `createShuffleSpec` projects and sorts those keys. The
mechanism is in
// `ShuffleSpecSuite`'s "createShuffleSpec sorts the projected keys at
their built-with types".
withTable("t_identity", "t_years") {
createTsTable("t_identity", Array(identity("ts")))
@@ -9892,6 +9893,286 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-59671: a partially clustered join leaves AQE's shuffle
coalescing alone") {
+ // AQE validates a stage's whole candidate plan before accepting a
shuffle-read change: on the
+ // base, a storage-partitioned join whose sides are aligned but not
grouped kept every shuffle
+ // in its stage uncoalesced, unrelated ones included. Partially clustered
distribution plans
+ // such a pair: the side that keeps its splits spreads them, and the other
replicates its
+ // group across them, so both report repeated keys on purpose. The
assertions below pin the
+ // read that the join's presence must leave alone.
+ val idCols = Array(Column.create("id", IntegerType), Column.create("data",
StringType))
+ createTable("pc1", idCols, Array(identity("id")))
+ createTable("pc2", idCols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.pc1 VALUES (1, 'a1'), (2, 'a2'), (3, 'a3')")
+ // Key 1 twice: the side holding it keeps and spreads its two splits,
which is what makes the
+ // pair ungrouped while its keys still line up.
+ sql("INSERT INTO testcat.ns.pc2 VALUES (1, 'b1'), (1, 'b1b'), (2, 'b2'),
(4, 'b4')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1") {
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(a) */ a.id, b.data
+ |FROM testcat.ns.pc1 a JOIN testcat.ns.pc2 b ON a.id = b.id
+ |UNION ALL
+ |SELECT count(*), cast(id % 2 AS STRING)
+ |FROM testcat.ns.pc1 GROUP BY id % 2
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(1, "b1"), Row(1, "b1b"), Row(2, "b2"), Row(2,
"1"), Row(1, "0")))
+
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ val joins = collect(plan) { case j: ShuffledJoin => j }
+ assert(joins.size == 1, s"test setup: one storage-partitioned
join:\n$plan")
+ assert(collectGroupPartitions(plan).exists { g =>
+
PartitioningCollection.representativeOf(g.outputPartitioning).exists(!_.isGrouped)
+ }, s"test setup: the pair is spread, so a side repeats its keys:\n$plan")
+ // The join's side shuffles nothing (this suite's `collectShuffles`
counts the exchanges a
+ // join reads through), and the chain shuffles once, for the aggregate:
the stage holding
+ // both is the one whose read the join's presence could have kept
uncoalesced.
+ assert(collectShuffles(plan).isEmpty, s"the join shuffles
nothing:\n$plan")
+ assert(collectAllShuffles(plan).size === 1,
+ s"and the chain shuffles once, for the aggregate:\n$plan")
+ val aqeReads = collect(df.queryExecution.executedPlan) { case r:
AQEShuffleReadExec => r }
+ assert(aqeReads.size === 1 && aqeReads.head.hasCoalescedPartition,
+ s"the aggregate's shuffle read must
coalesce:\n${df.queryExecution.executedPlan}")
+
+ // And the two share a stage, which is what makes the coalesce a
decision the join can
+ // block: a read in a stage of its own would coalesce whatever the join
did.
+ val finalStage = collect(df.queryExecution.executedPlan) {
+ case s: ResultQueryStageExec => s
+ }
+ assert(finalStage.size === 1 && finalStage.head.plan.exists {
+ case j: ShuffledJoin => true
+ case _ => false
+ }, s"test setup: the join is in the stage the read belongs to:\n" +
+ s"${df.queryExecution.executedPlan}")
+ }
+ }
+
+ test("SPARK-59671: a three-table chain keeps its AQE coalescing") {
+ val idCols = Array(Column.create("id", IntegerType), Column.create("data",
StringType))
+ createTable("p3a", idCols, Array(identity("id")))
+ createTable("p3b", idCols, Array(identity("id")))
+ createTable("p3c", idCols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.p3a VALUES (1, 'a1'), (2, 'a2')")
+ sql("INSERT INTO testcat.ns.p3b VALUES (1, 'b1'), (1, 'b1b'), (2, 'b2')")
+ sql("INSERT INTO testcat.ns.p3c VALUES (1, 'c1'), (2, 'c2')")
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1") {
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(a, b), MERGE(a, c) */ a.id AS aid, b.id AS bid,
c.data
+ |FROM testcat.ns.p3a a JOIN testcat.ns.p3b b ON a.id = b.id
+ |JOIN testcat.ns.p3c c ON a.id = c.id
+ |UNION ALL
+ |SELECT count(*), 0, 'x' FROM testcat.ns.p3a GROUP BY id % 2
+ |""".stripMargin)
+ // The union coerces the join's `id` and the aggregate's `count(*)` to
one type, bigint.
+ checkAnswer(df, Seq(Row(1L, 1, "c1"), Row(1L, 1, "c1"), Row(2L, 2, "c2"),
+ Row(1L, 0, "x"), Row(1L, 0, "x")))
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(collectShuffles(plan).isEmpty, s"the chain must not
shuffle:\n$plan")
+ assert(plan.exists(p => keyedPartitioningsOf(Seq(p)).size >= 2),
+ s"test setup: a side reports one keyed member per join key
column:\n$plan")
+ assert(ValidateRequirements.validate(plan), s"the chain's pairing holds
up:\n$plan")
+ // Each join zips a spread side against a repeating side: the collection
it reports agrees
+ // on no origin, so nothing above may read the pair's claim off it,
while the groupings
+ // below it hold the two complementary stamps.
+ val joins = collect(plan) { case j: ShuffledJoin => j }
+ assert(joins.nonEmpty, s"test setup: a storage-partitioned
chain:\n$plan")
+ assert(joins.forall(j =>
keyedPartitioningsOf(Seq(j)).forall(_.ungroupingOrigin.isEmpty)),
+ s"a zipped pair reports no claim of its own:\n$plan")
+ val origins =
collectAllGroupPartitions(plan).flatMap(_.ungroupingOrigin).distinct
+ assert(origins.contains(SPLIT_FOR_JOIN) &&
origins.contains(REPLICATED_FOR_JOIN),
+ s"one side spreads and the other repeats, got $origins:\n$plan")
+ assert(collect(df.queryExecution.executedPlan) { case r:
AQEShuffleReadExec => r }
+ .exists(_.hasCoalescedPartition),
+ s"the aggregate's shuffle read must
coalesce:\n${df.queryExecution.executedPlan}")
+ }
+ }
+
+ test("SPARK-59671: an AQE local read over a split scan does not lose
matches") {
+ // A lower join with a non-bucketed side cannot align, so it hash-shuffles
the bucketed scan
+ // as its source reports it: ungrouped, with one key's rows across splits.
AQE broadcasts the
+ // narrow side on its runtime size, and the wide side's shuffle, now a
broadcast join's probe
+ // side, becomes a local read. A local read reports the partitioning from
*before* the shuffle,
+ // which is that ungrouped keyed layout. The upper join then reads two
layouts that agree key
+ // by key while neither holds a key's whole group: consuming them as an
aligned pair would zip
+ // split i against split i and drop every match that crosses mappers.
Validation has to refuse
+ // the pair so that AQE keeps the full shuffle read instead.
+ val idCols = Array(Column.create("id", IntegerType), Column.create("data",
StringType))
+ createTable("lrt", idCols, Array(identity("id")))
+ createTable("lrs1", idCols, Array.empty[Transform])
+ createTable("lrs2", idCols, Array.empty[Transform])
+ // Key 1 in two splits, so that no single mapper holds all of its rows.
The padding keeps the
+ // relation over the runtime broadcast threshold, which is what keeps the
upper join a sort
+ // merge join reading both lower joins.
+ val pad = "x" * 4096
+ sql(s"INSERT INTO testcat.ns.lrt VALUES (1, 't1a$pad')")
+ sql(s"INSERT INTO testcat.ns.lrt VALUES (1, 't1b$pad')")
+ sql("INSERT INTO testcat.ns.lrt VALUES (2, 't2a')")
+ sql("INSERT INTO testcat.ns.lrs1 VALUES (1, 's1a'), (2, 's1b')")
+ sql("INSERT INTO testcat.ns.lrs2 VALUES (1, 's2a'), (2, 's2b')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ // Nothing broadcasts at plan time, so both lower joins plan as sort
merge joins and
+ // shuffle their sides; AQE broadcasts the narrow sides on their
runtime size alone.
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "1k",
+ // Coalescing would merge the wide side's shuffle into one slot first,
and a local read
+ // over a single slot holds every row and loses nothing. Off, the read
stays one slot per
+ // mapper, which is the shape that reports the scan's own ungrouped
layout.
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false",
+ // More reducer slots than the wide side has mappers, which is when a
local read keeps one
+ // partition per mapper.
+ SQLConf.SHUFFLE_PARTITIONS.key -> "4") {
+ val df = sql(
+ """
+ |SELECT substring(a.data, 1, 3), substring(b.data, 1, 3)
+ |FROM (SELECT t1.id AS id, t1.data AS data FROM testcat.ns.lrt t1
+ | JOIN testcat.ns.lrs1 s1 ON t1.id = s1.id) a
+ |JOIN (SELECT t2.id AS id, t2.data AS data FROM testcat.ns.lrt t2
+ | JOIN testcat.ns.lrs2 s2 ON t2.id = s2.id) b
+ |ON a.id = b.id
+ |""".stripMargin)
+ // Every pair of key 1's rows across the two relations, and key 2's one
pair. A read that
+ // kept one partition per mapper and paired them index by index would
find only the two
+ // same-split pairs of key 1 and key 2's one.
+ checkAnswer(df, Seq(Row("t1a", "t1a"), Row("t1a", "t1b"), Row("t1b",
"t1a"),
+ Row("t1b", "t1b"), Row("t2a", "t2a")))
+
+ // The shape the answer is about, so that it is not read as an accident:
the wide side's
+ // shuffle carries a scan that reports one key across splits, and both
narrow sides ended up
+ // broadcast, which is what turned the wide side's shuffle into a
local-read candidate.
+ val plan = df.queryExecution.executedPlan
+ val shuffles = collectAllShuffles(plan)
+ // The two wide sides are one scan, so they share one exchange by reuse;
each narrow side
+ // keeps its own, read whole as its broadcast's build input.
+ assert(shuffles.size == 3,
+ s"one shared wide-side shuffle and one per narrow
side:\n${plan.treeString}")
+ val splitScans = shuffles
+ .flatMap(s => keyedPartitioningsOf(Seq(s.child))).filter(!_.isGrouped)
+ assert(splitScans.nonEmpty,
+ s"test setup: a shuffled scan reports a key across
splits:\n${plan.treeString}")
+ assert(collect(plan) { case j: BroadcastHashJoinExec => j }.size == 2,
+ s"test setup: both lower joins broadcast at
runtime:\n${plan.treeString}")
+
+ // The refusal this test exists for: no local read may hand out the wide
side's ungrouped
+ // keyed layout. A local read building a broadcast reads the whole
relation and reports the
+ // scan's own unknown partitioning, which pairs nothing and is harmless.
+ assert(collect(plan) { case r: AQEShuffleReadExec if r.isLocalRead => r }
+ .forall(r => keyedPartitioningsOf(Seq(r)).isEmpty),
+ s"no local read reports a keyed layout:\n${plan.treeString}")
+ // And nothing re-pairs a partial read: the plan holds no grouping node
at all.
+ assert(collectAllGroupPartitions(plan).isEmpty,
+ s"nothing re-pairs a partial read:\n${plan.treeString}")
+ }
+ }
+
+ test("SPARK-59671: an ungrouped sort branch leaves AQE's shuffle coalescing
alone") {
+ // The sort branch keeps the scan's splits as they stand: the catalog
hands the partition keys
+ // over sorted, so no node settles them and the layout stays ungrouped
under the sort. An
+ // ordering reads the keys without pairing anything, so that layout
answers for itself.
+ // Before this, the finished plan failed validation on it and AQE kept the
aggregate
+ // branch's shuffle read uncoalesced: the stage-wide symptom this PR is
motivated by.
+ val idCols = Array(Column.create("id", IntegerType), Column.create("data",
StringType))
+ createTable("sot", idCols, Array(identity("id")))
+ // Key 1 across two splits, so the layout the sort reads is ungrouped.
+ sql("INSERT INTO testcat.ns.sot VALUES (2, 'o2')")
+ sql("INSERT INTO testcat.ns.sot VALUES (1, 'o1a')")
+ sql("INSERT INTO testcat.ns.sot VALUES (1, 'o1b')")
+
+ withTable("sos") {
+ sql("CREATE TABLE sos (k INT) USING parquet")
+ sql("INSERT INTO sos VALUES (1), (1), (2), (2), (3)")
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1") {
+ val df = sql(
+ """
+ |(SELECT id FROM testcat.ns.sot ORDER BY id)
+ |UNION ALL
+ |(SELECT count(*) FROM sos GROUP BY k)
+ |""".stripMargin)
+ // The union coerces the branch outputs to bigint.
+ val collected = df.collect().map(_.getLong(0)).toSeq
+ // The ordered branch comes first and stays ordered, its key's splits
apart.
+ assert(collected.take(3) == Seq(1L, 1L, 2L),
+ s"the ordered branch keeps its
order:\n${df.queryExecution.executedPlan}")
+ assert(collected.drop(3).sorted == Seq(1L, 2L, 2L), "one row per key
group")
+
+ val plan = df.queryExecution.executedPlan
+ // The symptom: the aggregate's shuffle read coalesces even though the
finished plan
+ // holds an ungrouped layout.
+ assert(collect(plan) { case r: AQEShuffleReadExec => r
}.exists(_.hasCoalescedPartition),
+ s"the aggregate's shuffle read coalesces:\n$plan")
+ // The shape it coalesces on: the sort reads the scan's own ungrouped
layout with no
+ // grouping node in between, which is what the old validation refused
the plan for.
+ assert(collectScans(plan).exists(s => s.outputPartitioning match {
+ case k: KeyedPartitioning => !k.isGrouped
+ case _ => false
+ }), s"test setup: the sort branch keeps an ungrouped layout:\n$plan")
+ assert(collectAllGroupPartitions(plan).isEmpty,
+ s"test setup: nothing settles the sort branch's keys:\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-59671: an aggregate over a left outer join's spread side groups
it") {
+ // A left outer join reports its left side verbatim, and the producer's
duplication gate
+ // makes that side the spread one: its layout holds a key's rows across
partitions and
+ // carries the alignment stamp. The stamp vouches for the pair the join
consumed and for
+ // nothing else, so the aggregate above must not read the layout as
clustered: it gets a
+ // grouping node first, or key 1's two splits come out as two groups.
+ val idCols = Array(Column.create("id", IntegerType), Column.create("data",
StringType))
+ createTable("lo1", idCols, Array(identity("id")))
+ createTable("lo2", idCols, Array(identity("id")))
+ // The left side holds three splits against two, so the heuristic
replicates the right and
+ // the left spreads, the only side a LeftOuter may duplicate.
+ sql("INSERT INTO testcat.ns.lo1 VALUES (1, 'l1a')")
+ sql("INSERT INTO testcat.ns.lo1 VALUES (1, 'l1b')")
+ sql("INSERT INTO testcat.ns.lo1 VALUES (2, 'l2')")
+ sql("INSERT INTO testcat.ns.lo2 VALUES (1, 'r1'), (2, 'r2')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1") {
+ val df = sql(
+ s"""
+ |SELECT id, count(*) FROM (
+ | ${selectWithMergeJoinHint("l", "r")} l.id AS id
+ | FROM testcat.ns.lo1 l LEFT OUTER JOIN testcat.ns.lo2 r ON l.id =
r.id
+ |) GROUP BY id
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(1, 2L), Row(2, 1L)))
+
+ // The shape: the join reports its spread side verbatim, stamp included,
and a grouping
+ // node settles it before the aggregate reads it.
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ val joins = collect(plan) { case j: ShuffledJoin => j }
+ assert(joins.size == 1, s"test setup: one storage-partitioned
join:\n$plan")
+ assert(keyedPartitioningsOf(Seq(joins.head)).exists(k =>
+ !k.isGrouped && k.ungroupingOrigin.contains(SPLIT_FOR_JOIN)),
+ s"test setup: the join reports its spread side verbatim:\n$plan")
+ assert(collectAllGroupPartitions(plan).exists(g =>
+ g.child.collectFirst { case _: SortMergeJoinExec => () }.isDefined &&
+ (g.outputPartitioning match {
+ case k: KeyedPartitioning => k.isGrouped &&
k.ungroupingOrigin.isEmpty
+ case _ => false
+ })),
+ s"the aggregate's side is grouped over the join:\n$plan")
+ }
+ }
}
/**
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
index fb9afa9cf3f3..7bab293edf36 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
@@ -21,7 +21,7 @@ import org.apache.spark.{SparkContext, SparkException}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute,
AttributeReference, SortOrder, TransformExpression}
-import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution,
KeyedPartitioning, KeyReducer, Partitioning, PartitioningCollection,
UnknownPartitioning}
+import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution,
KeyedPartitioning, KeyReducer, OrderedDistribution, Partitioning,
PartitioningCollection, REPLICATED_FOR_JOIN, SPLIT_FOR_JOIN,
UnknownPartitioning}
import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
import org.apache.spark.sql.connector.catalog.functions.{BucketFunction,
BucketReducer, DaysFunctionWithToYearsReducerWithLongResult,
DaysToYearsReducerWithLongResult, Reducer,
YearsFunctionWithToYearsReducerWithLongResult}
import org.apache.spark.sql.execution.{DummySparkPlan, LeafExecNode,
SafeForKWayMerge}
@@ -50,7 +50,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
val childKp = KeyedPartitioning(Seq(exprA, exprB), keys)
.withLayout(_.copy(isCollapsed = childCollapsed))
GroupPartitionsExec(DummySparkPlan(outputPartitioning = childKp),
joinKeyPositions,
- expected, distributePartitions = distribute)
+ expected, ungroupingOrigin = Option.when(distribute)(SPLIT_FOR_JOIN))
.outputPartitioning.asInstanceOf[KeyedPartitioning]
}
def keyOf(a: Int): InternalRowComparableWrapper =
@@ -590,7 +590,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2), row(2))))
val gpe = GroupPartitionsExec(child,
expectedPartitionKeys = Some(Seq(keyOf(1) -> 1, keyOf(2) -> 3, keyOf(3)
-> 1)),
- distributePartitions = true)
+ ungroupingOrigin = Some(SPLIT_FOR_JOIN))
gpe.execute()
assert(gpe.metrics("numInputPartitions").value === 3)
@@ -601,6 +601,52 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
assert(!gpe.metrics.contains("numReplicatedPartitionReads"), "distribute
never replicates")
}
+ test("SPARK-59671: the output claim is the node's own, never inherited") {
+ // The child's origin describes the shape its producer built. A node that
groups those
+ // partitions settles them: its output carries only the claim its own
producer stamped it
+ // with, none here, so no stale stamp above turns away a sound pair of
grouped sides.
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val child = ExecutableKeyedLeaf(
+ KeyedPartitioning(Seq(exprA), Seq(row(1), row(1), row(2)))
+ .withLayout(_.copy(ungroupingOrigin = Some(SPLIT_FOR_JOIN))))
+
+ val grouped = GroupPartitionsExec(child)
+ val groupedOut = grouped.outputPartitioning.asInstanceOf[KeyedPartitioning]
+ assert(groupedOut.isGrouped && groupedOut.ungroupingOrigin.isEmpty)
+ assert(groupedOut.partitionKeys.map(_.row) == Seq(row(1), row(2)))
+
+ // A stamped node stamps its own output whatever the child claimed, and
its routing follows
+ // the stamp: a repeating side groups first.
+ val restamped = GroupPartitionsExec(child,
+ expectedPartitionKeys = Some(Seq(keyOf(1) -> 2, keyOf(2) -> 1)),
+ ungroupingOrigin = Some(REPLICATED_FOR_JOIN))
+ assert(!restamped.distributePartitions, "a repeating side groups first")
+ val restampedOut =
restamped.outputPartitioning.asInstanceOf[KeyedPartitioning]
+ assert(!restampedOut.isGrouped)
+ assert(restampedOut.ungroupingOrigin.contains(REPLICATED_FOR_JOIN))
+ }
+
+ test("SPARK-59671: an alignment that settles every key spends its stamp") {
+ // A stamp says why the keys repeat on purpose; a grouping where every key
got one slot
+ // leaves nothing repeating, so the layout reports no claim even though
the node was stamped
+ // for one -- and it keeps the ordering claim a grouped layout has always
had.
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val child = ExecutableKeyedLeaf(KeyedPartitioning(Seq(exprA), Seq(row(1),
row(2))))
+ val gpe = GroupPartitionsExec(child,
+ expectedPartitionKeys = Some(Seq(keyOf(1) -> 1, keyOf(2) -> 1)),
+ ungroupingOrigin = Some(SPLIT_FOR_JOIN))
+ assert(gpe.ungroupingOrigin.contains(SPLIT_FOR_JOIN),
+ "the node keeps the stamp it was handed, and the routing it derives")
+ val out = gpe.outputPartitioning.asInstanceOf[KeyedPartitioning]
+ assert(out.isGrouped && out.ungroupingOrigin.isEmpty)
+ withSQLConf(SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
+ assert(out.satisfies(OrderedDistribution(Seq(SortOrder(exprA,
Ascending)))),
+ "a grouped layout keeps its ordering claim")
+ }
+ }
+
test("SPARK-59310: alignment prunes unmatched keys, pads missing ones") {
// The expected keys carry key 1 and a key-3 slot the child does not hold,
as an inner
// join's intersection combined with the other side's layout would. The 2
splits of key 2
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
index 6b1314fb88f1..6d3a7708a000 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
@@ -33,7 +33,7 @@ import org.apache.spark.sql.connector.catalog.functions._
import org.apache.spark.sql.execution.{BinaryExecNode, DummySparkPlan,
LeafExecNode, SafeForKWayMerge, SortExec, UnaryExecNode}
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
GroupPartitionsExec}
-import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec,
SortMergeJoinExec}
+import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec,
SortMergeAsOfJoinExec, SortMergeJoinExec}
import org.apache.spark.sql.execution.python.FlatMapCoGroupsInPandasExec
import org.apache.spark.sql.execution.window.WindowExec
import org.apache.spark.sql.internal.SQLConf
@@ -2036,6 +2036,34 @@ class EnsureRequirementsSuite extends SharedSparkSession
{
}
}
+ test("SPARK-59671: an ordering spread carries its own reason and serves only
the ordering") {
+ val exprKey = AttributeReference("k", IntegerType)()
+ // The keys arrive out of order for the ascending ordering the operator
reads, and key 1 holds
+ // two splits. The spread that settles the order keeps the splits, so the
layout it reports
+ // stays ungrouped, on purpose and for the ordering only.
+ val child = new DummySparkPlanWithBatchScanChild(
+ outputPartitioning = KeyedPartitioning(Seq(exprKey),
+ Seq(InternalRow(2), InternalRow(1), InternalRow(1))))
+ val ordering = OrderedDistribution(Seq(SortOrder(exprKey, Ascending)))
+
+ withSQLConf(SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
+ val newChild = EnsureRequirements.apply(parentRequiring(child,
ordering)).children.head
+ val groupings = groupPartitionsNodes(newChild)
+ assert(groupings.size == 1,
+ s"the unsorted keys are spread into order, got
${newChild.getClass.getSimpleName} " +
+ s"with ${groupings.size} grouping nodes")
+ assert(groupings.head.distributePartitions)
+ assert(groupings.head.ungroupingOrigin == Some(SPLIT_FOR_ORDERING))
+ val layout = newChild.outputPartitioning.asInstanceOf[KeyedPartitioning]
+ assert(!layout.isGrouped)
+ assert(layout.partitionKeys.map(_.row.getInt(0)) == Seq(1, 1, 2))
+ // The spread serves the ordering it was built for...
+ assert(layout.satisfies(ordering))
+ // ...and its stamp is no pairing role: the same layout must not serve a
clustering.
+ assert(!layout.satisfies(ClusteredDistribution(Seq(exprKey))))
+ }
+ }
+
test("SPARK-58968: a projection whose resulting count matches
requiredNumPartitions still " +
"groups") {
val exprN = AttributeReference("n", IntegerType)()
@@ -2597,6 +2625,71 @@ class EnsureRequirementsSuite extends SharedSparkSession
{
((false, true)))
}
+ test("SPARK-59671: the alignment stamps one repeating side, only where the
join type may") {
+ // The replicate side is the one with fewer pre-alignment partitions (the
dummy plans carry no
+ // `logicalLink`, which forces the fallback). A join type may duplicate
one side, both, or
+ // neither, and the stamps follow: a side that may not repeat groups
instead, and a join type
+ // that may repeat neither side builds no spread pair at all.
+ def stampsOf(smj: SparkPlan): Seq[(Option[UngroupingOrigin], Boolean)] =
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true") {
+ EnsureRequirements.apply(smj).children.map { child =>
+ val gpe = child.collectFirst { case g: GroupPartitionsExec => g }.get
+ (gpe.ungroupingOrigin,
+ gpe.outputPartitioning.asInstanceOf[KeyedPartitioning].isGrouped)
+ }
+ }
+ def smj(joinType: JoinType, leftKeys: Seq[InternalRow],
+ rightKeys: Seq[InternalRow]): SortMergeJoinExec =
+ SortMergeJoinExec(Seq(exprA), Seq(exprB), joinType, None,
+ DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(exprA),
leftKeys)),
+ DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(exprB),
rightKeys)))
+ val threeSplits = Seq(InternalRow(1), InternalRow(1), InternalRow(1))
+ val twoKeys = Seq(InternalRow(1), InternalRow(2))
+
+ // The right side holds fewer partitions and is picked to repeat. An inner
join may duplicate
+ // either side and a left outer join the right one, so both get one spread
and one repeating
+ // side.
+ assert(stampsOf(smj(Inner, threeSplits, twoKeys)) ===
+ Seq((Some(SPLIT_FOR_JOIN), false), (Some(REPLICATED_FOR_JOIN), false)))
+ assert(stampsOf(smj(LeftOuter, threeSplits, twoKeys)) ===
+ Seq((Some(SPLIT_FOR_JOIN), false), (Some(REPLICATED_FOR_JOIN), false)))
+
+ // The mirrored left outer join picks the left side, which it may not
duplicate: the spread
+ // is skipped and both sides group.
+ assert(stampsOf(smj(LeftOuter, Seq(InternalRow(1)), threeSplits)) ===
+ Seq((None, true), (None, true)))
+
+ // A full outer join may duplicate neither side: both group, whichever is
picked.
+ assert(stampsOf(smj(FullOuter, threeSplits, twoKeys)) ===
+ Seq((None, true), (None, true)))
+ }
+
+ test("SPARK-59671: an as-of join gets no keyed alignment, so its sides
shuffle") {
+ // `EnsureRequirements.checkKeyGroupCompatible` plans a keyed alignment
for the sort-merge and
+ // shuffled-hash joins only, and an as-of join is a `ShuffledJoin` outside
that list: its
+ // matches are read in the order within a partition, which neither a
spread nor a repeating
+ // side preserves. The same shapes an inner join stamps and aligns must
fall back to a
+ // shuffle here.
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true") {
+ val asOf = SortMergeAsOfJoinExec(Seq(exprA), Seq(exprB), Seq(exprA),
Seq(exprB),
+ GreaterThan(exprA, exprB), exprA, Inner, None,
+ DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(exprA),
+ Seq(InternalRow(1), InternalRow(1), InternalRow(1)))),
+ DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(exprB),
+ Seq(InternalRow(1), InternalRow(2)))))
+ EnsureRequirements.apply(asOf).children.foreach { child =>
+ assert(groupPartitionsNodes(child).isEmpty,
+ s"an as-of join builds no stamped side:\n$child")
+ assert(child.collectFirst { case s: ShuffleExchangeExec => s
}.isDefined,
+ s"the side shuffles onto a hash layout:\n$child")
+ }
+ }
+ }
+
test("SPARK-58996: a single-child operator over a partially clustered layout
still gets " +
"grouped") {
// A partially clustered `GroupPartitionsExec` reports a non-grouped
partitioning by design,
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/ValidateRequirementsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/ValidateRequirementsSuite.scala
index 4e02a10eb41e..b3002b450406 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/ValidateRequirementsSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/ValidateRequirementsSuite.scala
@@ -17,12 +17,17 @@
package org.apache.spark.sql.execution.exchange
-import org.apache.spark.sql.catalyst.expressions.{Ascending, SortOrder}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Ascending,
AttributeReference, Literal, SortOrder}
+import org.apache.spark.sql.catalyst.optimizer.BuildLeft
import org.apache.spark.sql.catalyst.plans.Inner
-import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning,
SinglePartition}
-import org.apache.spark.sql.execution.SortExec
-import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution,
HashPartitioning, KeyedPartitioning, Partitioning, PartitioningCollection,
REPLICATED_FOR_JOIN, SinglePartition, SPLIT_FOR_JOIN, UngroupingOrigin}
+import org.apache.spark.sql.execution.{CoGroupExec, DummySparkPlan, SortExec,
SparkPlan}
+import org.apache.spark.sql.execution.datasources.v2.GroupPartitionsExec
+import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec,
SortMergeJoinExec}
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{IntegerType, ObjectType}
class ValidateRequirementsSuite extends SharedSparkSession {
@@ -158,4 +163,303 @@ class ValidateRequirementsSuite extends
SharedSparkSession {
testNestedJoin(Seq((2, 2), (1, 1)), Seq((2, 2)), Seq(5, 5, 5), false)
testNestedJoin(Seq((2, 2), (1, 1)), Seq((2, 5)), Seq(5, 5, 5), false)
}
+
+ test("SPARK-59671: a co-partitioning operator judges keyed children by their
pairing") {
+ // The sides of a storage-partitioned join aligned for skew repeat their
spread keys on purpose:
+ // neither satisfies a clustering on its own, and each side's layout says
where that ungrouping
+ // comes from, so the pairing is what the operator reads.
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2))
+ def side(
+ attr: AttributeReference,
+ role: UngroupingOrigin,
+ keys: Seq[InternalRow]): DummySparkPlan =
+ DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(attr), keys).withLayout(_.copy(ungroupingOrigin
= Some(role))))
+ val left = side(a, SPLIT_FOR_JOIN, rows)
+ val right = side(b, REPLICATED_FOR_JOIN, rows)
+ val join = ShuffledHashJoinExec(Seq(a), Seq(b), Inner, BuildLeft, None,
left, right)
+ assert(ValidateRequirements.validate(join),
+ s"a spread side and one that repeats the group pair, which is the
requirement:\n$join")
+
+ // The same key rows in another order do not pair: position by position is
the contract.
+ val off = side(b, REPLICATED_FOR_JOIN, Seq(InternalRow(1), InternalRow(2),
InternalRow(1)))
+ assert(!ValidateRequirements.validate(join.copy(right = off)),
+ "the same keys in a different order are not aligned")
+
+ // Nor does a keyed side pair with one that never pairs: a hashed side
matches neither the keys
+ // nor the layout of a keyed one.
+ val hashed = DummySparkPlan(outputPartitioning = HashPartitioning(Seq(b),
3))
+ assert(!ValidateRequirements.validate(join.copy(right = hashed)),
+ "a keyed side does not pair with a hashed one")
+
+ // A side that does not say why it is ungrouped is one no producer built:
it is waiting for the
+ // node that would settle it, and a finished plan has none left to insert.
+ val unsettled = DummySparkPlan(outputPartitioning =
KeyedPartitioning(Seq(b), rows))
+ assert(!ValidateRequirements.validate(join.copy(right = unsettled)),
+ s"an ungrouped side is admitted only where a producer built one:\n$join")
+
+ // Nor do two sides of the same role pair: an alignment is one of each,
and two spreads leave a
+ // key's rows split between them with nothing repeating the group.
+ assert(!ValidateRequirements.validate(join.copy(right = side(b,
SPLIT_FOR_JOIN, rows))),
+ "two sides that spread the same keys are not an alignment")
+
+ // And a pair each side satisfies on its own is still refused when the
sides do not line up,
+ // which is the operator's requirement: the pairing, not the per-side
answer.
+ val leftKeys = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(2),
InternalRow(3))))
+ val rightKeys = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(b), Seq(InternalRow(4), InternalRow(5),
InternalRow(6))))
+
assert(leftKeys.outputPartitioning.satisfies(ClusteredDistribution(Seq(a))) &&
+ rightKeys.outputPartitioning.satisfies(ClusteredDistribution(Seq(b))),
+ "test setup: each side answers its own clustering")
+ assert(!ValidateRequirements.validate(ShuffledHashJoinExec(
+ Seq(a), Seq(b), Inner, BuildLeft, None, leftKeys, rightKeys)),
+ "the sides do not line up, and nothing else the operator reads says
otherwise")
+ }
+
+ test("SPARK-59671: a keyed pair is judged on the layouts the plan holds") {
+ // Neither side's keys are deduped or re-sorted to make the pair line up:
reading it through
+ // `createShuffleSpec` would do both, so two sides whose layouts disagree
as they stand would
+ // answer as though they agreed.
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2))
+ def side(attr: AttributeReference, role: UngroupingOrigin,
+ keys: Seq[InternalRow]): DummySparkPlan =
+ DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(attr), keys).withLayout(_.copy(ungroupingOrigin
= Some(role))))
+ val left = side(a, SPLIT_FOR_JOIN, rows)
+ // The same keys in another order, which a projection onto distinct sorted
keys would
+ // normalize away.
+ val off = side(b, REPLICATED_FOR_JOIN, Seq(InternalRow(1), InternalRow(2),
InternalRow(1)))
+ assert(!ValidateRequirements.validate(ShuffledHashJoinExec(
+ Seq(a), Seq(b), Inner, BuildLeft, None, left, off)),
+ "the sides report different layouts, and nothing normalizes them")
+
+ // And a count: the other side holds two partitions, and neither side has
the node that would
+ // make the two agree.
+ val twoKeys = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(b), Seq(InternalRow(1), InternalRow(2))))
+ assert(!ValidateRequirements.validate(ShuffledHashJoinExec(
+ Seq(a), Seq(b), Inner, BuildLeft, None, left, twoKeys)),
+ "a child holding three partitions does not pair with one holding two")
+ }
+
+ test("SPARK-59671: a pair is judged on the key the operation clusters on") {
+ // The subset permission applies where the operation's keys are a subset
of the source's
+ // partitioning keys: a side partitioned on `[a, b]` can serve an operator
on `[a]`. Such a side
+ // is judged on its own partitions under `[a]`, which is what its own key
already gives it when
+ // dropping `b` merges no partition, whether or not a node stands over it.
No key is deduped and
+ // none is re-sorted, so the pairs that line up are the ones whose keys
line up as reported.
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val c = AttributeReference("c", IntegerType)()
+ val d = AttributeReference("d", IntegerType)()
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"false") {
+ def joinOnKeys(left: SparkPlan, right: SparkPlan): SparkPlan =
+ ShuffledHashJoinExec(Seq(a), Seq(c), Inner, BuildLeft, None, left,
right)
+
+ // The projection's result, which is what such a plan reports from the
side it grouped.
+ val projected = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(2))))
+ val projectedRight = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(c), Seq(InternalRow(1), InternalRow(2))))
+ assert(ValidateRequirements.validate(joinOnKeys(projected,
projectedRight)),
+ "the layout the grouping node leaves is what the pair is judged on")
+
+ // The source's keys, one step earlier: the second expression is the
operation's to drop, and
+ // what is left is the same pair, so it stands.
+ val source = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 1), InternalRow(2,
2))))
+ val sourceRight = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(c, d), Seq(InternalRow(1, 1), InternalRow(2,
2))))
+ assert(ValidateRequirements.validate(joinOnKeys(source, sourceRight)),
+ s"the key the operation clusters on is what the side offers:\n$source")
+
+ // And it is offered in the order it is reported, so a side whose keys
run the other way is
+ // not aligned: a spec that sorted them would call these two the same
pair.
+ val reversed = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(a, b), Seq(InternalRow(2, 2), InternalRow(1,
1))))
+ assert(!ValidateRequirements.validate(joinOnKeys(reversed, sourceRight)),
+ s"the keys are compared as reported, not as sorted:\n$reversed")
+ }
+ }
+
+ test("SPARK-59671: a single clustered child still owes its own grouping") {
+ // The pairing stands in for the per-child check only where children pair
with each other.
+ // An operator with a single clustered child, an aggregate over a join
output say, is judged
+ // per side: the ungrouped keyed layout does not satisfy it until a
grouping stands under.
+ val a = AttributeReference("a", IntegerType)()
+ val child = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1))))
+ val parent = DummySparkPlan(
+ children = Seq(child),
+ requiredChildDistribution = Seq(ClusteredDistribution(Seq(a))),
+ requiredChildOrdering = Seq(Nil))
+ assert(!ValidateRequirements.validate(parent), "an ungrouped child fails a
lone clustered slot")
+ }
+
+ test("SPARK-59671: a partially clustered pair planned by the rule passes
validation") {
+ // Partially clustered distribution spreads a side, so its replicate side
reports a
+ // non-grouped layout whose keys repeat on purpose, and the validator's
per-side check
+ // refused such a pair. AQE validates a stage's whole candidate plan
before accepting a
+ // shuffle-read change, so the plan's stage took no coalescing either. The
pairing takes
+ // it: the spread is deliberate, and the two sides agree index by index.
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true") {
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val left = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1),
InternalRow(2))))
+ val right = DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(b), Seq(InternalRow(1), InternalRow(2))))
+ val plan = new EnsureRequirements().apply(
+ SortMergeJoinExec(Seq(a), Seq(b), Inner, None, left, right))
+ val join = plan.collectFirst { case j: SortMergeJoinExec => j }
+ .getOrElse(fail(s"expected the join back:\n${plan.treeString}"))
+ def groupingAt(plan: SparkPlan): Option[GroupPartitionsExec] = plan
match {
+ case g: GroupPartitionsExec => Some(g)
+ case s: SortExec if !s.global => groupingAt(s.child)
+ case _ => None
+ }
+ val sideGroupings = join.children.flatMap(groupingAt)
+ assert(sideGroupings.size == 2,
+ s"test setup: both sides are aligned:\n${plan.treeString}")
+ assert(sideGroupings.exists { g =>
+
PartitioningCollection.representativeOf(g.outputPartitioning).exists(!_.isGrouped)
+ }, s"test setup: a side keeps its splits, so its keys
repeat:\n${plan.treeString}")
+ assert(ValidateRequirements.validate(plan),
+ s"a spread pair that agrees on its keys is
accepted:\n${plan.treeString}")
+ }
+ }
+
+ test("SPARK-59671: a collection of keyed members is judged by its pairing") {
+ // A side that reports several keyed alternatives (a projection over a
join keeps one per join
+ // key column) is judged on the members the admission keeps: a side whose
members all carry the
+ // same origin for their ungrouped keys is a pair the producer aligns, and
it holds up when the
+ // member keyed on the join key is the one that lines up with the other
side.
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2))
+ def keyed(attr: AttributeReference, role: UngroupingOrigin):
KeyedPartitioning =
+ KeyedPartitioning(Seq(attr), rows).withLayout(_.copy(ungroupingOrigin =
Some(role)))
+ def side(partitioning: Partitioning): DummySparkPlan =
+ DummySparkPlan(outputPartitioning = partitioning)
+ val bothAlternatives = side(PartitioningCollection.fromPartitionings(
+ Seq(keyed(a, REPLICATED_FOR_JOIN), keyed(b, REPLICATED_FOR_JOIN))))
+ assert(ValidateRequirements.validate(ShuffledHashJoinExec(
+ Seq(a), Seq(a), Inner, BuildLeft, None, bothAlternatives, side(keyed(a,
SPLIT_FOR_JOIN)))),
+ "the member keyed on the join key pairs with the other side")
+
+ // The pairing still has to be there: a side whose members are keyed on
something else offers
+ // nothing to pair with.
+ val wrongKeys = side(PartitioningCollection.fromPartitionings(
+ Seq(keyed(b, REPLICATED_FOR_JOIN), keyed(b, REPLICATED_FOR_JOIN))))
+ assert(!ValidateRequirements.validate(ShuffledHashJoinExec(
+ Seq(a), Seq(a), Inner, BuildLeft, None, wrongKeys, side(keyed(a,
SPLIT_FOR_JOIN)))),
+ "a side offering no member keyed on the join keys does not pair")
+ }
+
+ test("SPARK-59671: a multi-child clustered operator is judged on its
children's pairing") {
+ // Every operator whose children all owe a `ClusteredDistribution` reads
one layout they have to
+ // hold together, whichever operator it is: a cogroup zips corresponding
partitions, so two
+ // sides that each satisfy the distribution on their own are not enough.
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2))
+ def cogroupOf(left: SparkPlan, right: SparkPlan): SparkPlan = CoGroupExec(
+ (key: Any, l: Iterator[Any], r: Iterator[Any]) => Nil,
+ Literal(1), Literal(1), Literal(1), Seq(a), Seq(b), Seq(a), Seq(b), Nil,
Nil,
+ AttributeReference("obj", ObjectType(classOf[AnyRef]))(), left, right)
+ def grouped(attr: AttributeReference, keys: Seq[Int]): DummySparkPlan =
DummySparkPlan(
+ outputOrdering = Seq(SortOrder(attr, Ascending)),
+ outputPartitioning = KeyedPartitioning(Seq(attr),
keys.map(InternalRow(_))))
+
+ // A side that does not say why it is ungrouped is one no producer built,
which is the whole of
+ // what the validator can read: the operator kind is the producer's to
check, and it is checked
+ // where the alignment is built.
+ val unsettled = DummySparkPlan(
+ outputOrdering = Seq(SortOrder(a, Ascending)),
+ outputPartitioning = KeyedPartitioning(Seq(a), rows))
+ assert(!ValidateRequirements.validate(cogroupOf(unsettled, unsettled)),
+ "an ungrouped side no producer built is refused, whoever reads it")
+
+ // What a cogroup shares with a join is the mutual check on the layouts
its sides report.
+ def cogroup(aKeys: Seq[Int], bKeys: Seq[Int]): SparkPlan =
+ cogroupOf(grouped(a, aKeys), grouped(b, bKeys))
+ assert(ValidateRequirements.validate(cogroup(Seq(1, 2), Seq(1, 2))),
+ "two sides holding the same grouped layout are read as they stand")
+ assert(!ValidateRequirements.validate(cogroup(Seq(1, 2), Seq(1, 2, 3))),
+ "a side holding three partitions does not pair with one holding two")
+ assert(!ValidateRequirements.validate(cogroup(Seq(1, 2), Seq(2, 1))),
+ "the same keys in another order are not aligned")
+ }
+
+ test("SPARK-59671: a side is judged on the member that pairs, whichever one
it is") {
+ // A side that reports several keyed alternatives offers all of them: the
one keyed on the
+ // join key is the one that lines up, and it need not be the first the
side reports. Reading a
+ // single member would refuse this pair, which the planner builds whenever
a projection keeps
+ // two key columns.
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2))
+ def keyed(attr: AttributeReference, role: UngroupingOrigin):
KeyedPartitioning =
+ KeyedPartitioning(Seq(attr), rows).withLayout(_.copy(ungroupingOrigin =
Some(role)))
+ val left = DummySparkPlan(outputPartitioning =
+ PartitioningCollection.fromPartitionings(
+ Seq(keyed(b, REPLICATED_FOR_JOIN), keyed(a, REPLICATED_FOR_JOIN))))
+ val right = DummySparkPlan(outputPartitioning = keyed(a, SPLIT_FOR_JOIN))
+ assert(ValidateRequirements.validate(ShuffledHashJoinExec(
+ Seq(a), Seq(a), Inner, BuildLeft, None, left, right)),
+ "the second member is the one keyed on the join key, and it pairs")
+ }
+
+ test("SPARK-59671: the ordering requirement is asked on top of the pairing")
{
+ // A pair can line up and still owe its operator an ordering: the pairing
covers the
+ // distribution clause alone, so a sort-merge join over the aligned sides,
with nothing
+ // ordering them, is refused.
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val rows = Seq(InternalRow(1), InternalRow(1), InternalRow(2))
+ def side(attr: AttributeReference, role: UngroupingOrigin): DummySparkPlan
=
+ DummySparkPlan(outputPartitioning =
+ KeyedPartitioning(Seq(attr), rows).withLayout(_.copy(ungroupingOrigin
= Some(role))))
+ val left = side(a, SPLIT_FOR_JOIN)
+ val right = side(b, REPLICATED_FOR_JOIN)
+ assert(left.outputOrdering.isEmpty && right.outputOrdering.isEmpty,
+ "test setup: nothing orders the sides")
+ assert(!ValidateRequirements.validate(
+ SortMergeJoinExec(Seq(a), Seq(b), Inner, None, left, right)),
+ "the sides pair, and still owe the join its ordering")
+ }
+
+ test("SPARK-59671: a collapsed pair is not admitted on its pairing alone") {
+ // A layout whose keys were collapsed (a partition standing for several of
the source's) serves
+ // a clustering through a grouping node only where that grouping is
permitted, and the
+ // permission is a config. Without it nothing admits such a member, and a
side offering nothing
+ // has nothing to pair with.
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ def collapsed(attr: AttributeReference): DummySparkPlan = DummySparkPlan(
+ outputPartitioning = KeyedPartitioning(Seq(attr),
+ Seq(InternalRow(1), InternalRow(1), InternalRow(2)))
+ .withLayout(_.copy(isCollapsed = true)))
+ def pair: SparkPlan = ShuffledHashJoinExec(Seq(a), Seq(b), Inner,
BuildLeft, None,
+ collapsed(a), collapsed(b))
+
+ // Whether a collapsed layout may be grouped is the producer's permission
to apply, and it
+ // applies it where the stamp goes: the validator reads no configuration,
so the same pair
+ // answers the same way either way.
+ Seq("false", "true").foreach { subset =>
+ withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key
-> subset) {
+ assert(!ValidateRequirements.validate(pair),
+ s"a pair no producer built is refused whatever the permission
says:\n$pair")
+ }
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
