ulysses-you commented on code in PR #58659:
URL: https://github.com/apache/spark/pull/58659#discussion_r4012377678
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -64,70 +72,15 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
*/
case class GroupPartitionsExec(
child: SparkPlan,
- @transient joinKeyPositions: Option[Seq[Int]] = None,
- @transient expectedPartitionKeys:
Option[Seq[(InternalRowComparableWrapper, Int)]] = None,
- @transient reducers: Option[Seq[Option[KeyReducer]]] = None,
- @transient distributePartitions: Boolean = false,
- @transient enableSortedMerge: Boolean = false
+ @transient grouping: PartitionGrouping,
Review Comment:
grouping and outputPartitioning are carried over a child rewrite (:302-303),
and the factory doc (:404-415) states the invariant that nothing in the tree
hands this node a child reporting a different partitioning. I verified it for
the codegen and columnar wrappers (ColumnarToRowExec, RowToColumnarExec,
InputAdapter, WholeStageCodegenExec all report child.outputPartitioning), but
not for AQE, and I could not settle it by reading.
A probe at this head shows what the carried fields do once the invariant is
broken:
- withNewChildren(Seq(childReportingOnePartition)) still reports
numPartitions 2 and groups the two old input partitions, silently, while
child.outputPartitioning reports 1.
- A Java serialization round trip leaves grouping == null (NPE on
groupedPartitions) and an outputPartitioning whose KeyLayout.partitionKeys is
null (KeyLayout:493 is @transient) while numPartitions still reports the old
count, so toString and equals change across the wire.
The base held all three of these as @transient ... lazy val (base :77, :226,
:309), so a rewritten or deserialized node recomputed them from its new child
and could not be stale.
The rewriter I could not rule out is an AQEShuffleReadExec landing between
this node and a keyed shuffle child: CoalesceShufflePartitions.isSupported
accepts a keyed ENSURE_REQUIREMENTS shuffle (CoalesceShufflePartitions:38-44),
and AQEShuffleReadExec reports UnknownPartitioning for a keyed shuffle
(AQEShuffleReadExec:100-105). With the stale keyed claim the join above still
satisfies its distribution, so ValidateRequirements.validate
(AdaptiveSparkPlanExec:205-213) would accept the coalesced plan, where the base
rejects it and reverts the rewrite.
Since I could not build a plan with a GroupPartitionsExec directly over a
keyed ShuffleExchangeExec, I am not calling this blocking. Could we either
re-derive in withNewChildInternal when newChild.outputPartitioning ne
child.outputPartitioning, or name in the doc the rules that guarantee the
invariant? Either way a future child-rewriting rule fails loudly instead of
reading a layout its child no longer has.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -599,7 +700,16 @@ case class EnsureRequirements(
(!compatibleAsIs ||
conf.v2BucketingPartiallyClusteredDistributionEnabled) &&
(conf.v2BucketingPushPartValuesEnabled ||
conf.v2BucketingAllowKeysSubsetOfPartitionKeys)
- if (pushCommonValues) {
+ // Neither route lines the two sides up, so there is no pairing to commit
to. Asked here rather
+ // than after the push branch, which is the only thing below that does any
work.
+ if (!compatibleAsIs && !pushCommonValues) {
+ return None
+ }
+ // What the push branch builds, when it runs. Empty otherwise, and then
each side gets a plain
+ // grouping node instead, which it needs exactly when its source reports
more than one partition
+ // per key. Building both eagerly would derive a grouping the push branch
throws away, and that
+ // is one hash per partition key.
+ val pushed = if (pushCommonValues) {
logInfo("Pushing common partition values for storage-partitioned join")
Review Comment:
This log fires inside `pushed`, but the pairing can still be discarded by
the describesSameKeys gate at :915, which cannot be asked earlier because it
needs the built nodes. So the log can announce a pushdown that never lands. The
bestPair.isEmpty early return at :689 avoids exactly this for the other
discarding case, so this reads as a miss. Moving the log after the gate would
keep the messages truthful, and this line is how a reader explains the extra
shuffle in the SPARK-59050 query.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -788,152 +882,193 @@ case class EnsureRequirements(
}
// Now we need to push-down the common partition information to the
`GroupPartitionsExec`s.
- newLeft = applyGroupPartitions(left, leftSpec.joinKeyPositions,
mergedPartitionKeys,
- leftReducers, distributePartitions = applyPartialClustering &&
!replicateLeftSide)
- newRight = applyGroupPartitions(right, rightSpec.joinKeyPositions,
mergedPartitionKeys,
- rightReducers, distributePartitions = applyPartialClustering &&
!replicateRightSide)
+ Some((
+ GroupPartitionsExec(rawLeft, leftSpec.joinKeyPositions,
+ Some(mergedPartitionKeys), leftReducers,
+ distributePartitions = applyPartialClustering && !replicateLeftSide),
+ GroupPartitionsExec(rawRight, rightSpec.joinKeyPositions,
+ Some(mergedPartitionKeys), rightReducers,
+ distributePartitions = applyPartialClustering &&
!replicateRightSide)))
+ } else {
+ None
}
- if (compatibleAsIs || pushCommonValues) Some(Seq(newLeft, newRight)) else
None
+ // The pairing is only worth committing to if both children still declare
the same aligned key
+ // sequence once they are built. They can fail that. A
`GroupPartitionsExec` gives up its keyed
+ // claim when it turns out to regroup a layout that pins undeclared rows to
+ // `hash(key) % numPartitions` (see
`KeyLayout.mayContainUnknownPartitionKeys`), and
+ // only the node knows the permutation it performs, so that answer arrives
after the pairing
+ // was chosen. Asking before returning is what keeps the join from
skipping both shuffles for a
+ // child that no longer satisfies its distribution, which is a plan
`ValidateRequirements`
+ // rejects and every AQE rule that needs a valid plan then refuses to
touch.
+ //
+ // The check is pairwise, not a per-side `satisfies`. Partially clustered
distribution leaves
+ // both children value-aligned yet not grouped on purpose, so a per-side
gate would refuse that
+ // whole family. What both sides owe each other is the key sequence
`alignToExpectedKeys`
+ // guarantees, each key repeated as many times as the merge expects,
whichever side replicates.
+ // Through `KeyLayout.describesSameKeys`, which carries the reason the key
types are compared
+ // as well as the rows.
+ def declaredLayout(plan: SparkPlan): Option[KeyLayout] =
+
PartitioningCollection.representativeOf(plan.outputPartitioning).map(_.layout)
+ val (newLeft, newRight) =
+ pushed.getOrElse((groupIfNeeded(rawLeft), groupIfNeeded(rawRight)))
+ Option.when(declaredLayout(newLeft).exists { left =>
+ declaredLayout(newRight).exists(left.describesSameKeys)
+ })(Seq(newLeft, newRight))
+ }
+
+ /**
+ * Whether the two children satisfy their distributions and line up with
each other as they
+ * arrive, so that the pairing has nothing to add. This is the question
`ValidateRequirements`
+ * asks of a finished plan, restricted to these two children.
+ *
+ * Both sides have to answer with an unprojected spec. A projected one
describes the layout a
+ * [[GroupPartitionsExec]] would emit rather than the one the child has, so
two sides can agree
+ * through their projections while their partitions do not line up at all.
+ */
+ private def alreadyCoPartitioned(
+ left: SparkPlan,
+ leftRequired: ClusteredDistribution,
+ right: SparkPlan,
+ rightRequired: ClusteredDistribution): Boolean = {
+ def unprojectedSpecs(plan: SparkPlan, required: ClusteredDistribution):
Seq[KeyedShuffleSpec] =
+ if (plan.outputPartitioning.satisfies(required)) {
+ candidatesFor(plan, required).filter(_.joinKeyPositions.isEmpty)
+ } else {
+ Nil
+ }
+ val leftSpecs = unprojectedSpecs(left, leftRequired)
+ val rightSpecs = unprojectedSpecs(right, rightRequired)
+ leftSpecs.exists(l => rightSpecs.exists(l.isCompatibleWith))
}
private def checkShufflePartitionIdPassThroughCompatible(
left: SparkPlan,
+ leftRequired: ClusteredDistribution,
right: SparkPlan,
- requiredChildDistribution: Seq[Distribution]): Boolean = {
+ rightRequired: ClusteredDistribution): Boolean = {
(left.outputPartitioning, right.outputPartitioning) match {
case (p1: ShufflePartitionIdPassThrough, p2:
ShufflePartitionIdPassThrough) =>
- assert(requiredChildDistribution.length == 2)
- val leftSpec = p1.createShuffleSpec(
- requiredChildDistribution.head.asInstanceOf[ClusteredDistribution])
- val rightSpec = p2.createShuffleSpec(
- requiredChildDistribution(1).asInstanceOf[ClusteredDistribution])
- leftSpec.isCompatibleWith(rightSpec)
+
p1.createShuffleSpec(leftRequired).isCompatibleWith(p2.createShuffleSpec(rightRequired))
case _ =>
false
}
}
/**
- * The innermost `GroupPartitionsExec` reachable from `plan` by descending
only through nodes
- * this rule itself inserted above it, together with a function rebuilding
the traversed local
- * sorts over a replacement node. `None` when no `GroupPartitionsExec` is
reachable.
+ * The plan a co-partitioned child's source reports, with every grouping and
local sort this rule
+ * put over it peeled off. `plan` itself when it carries none.
+ *
+ * The grouping peeled here is one an earlier pass left behind, not one this
pass put on. The
+ * per-child step skips co-partitioned children entirely, so nothing of this
pass's is under
+ * there when `checkKeyGroupCompatible` asks. The shuffle step, the other
caller, can see both.
+ * `EnsureRequirements` is re-run on plans it already produced, since
`AdaptiveSparkPlanExec`
+ * builds one instance of this rule, and
`ConvertSortMergeJoinToShuffledHashJoin` and
+ * `OptimizeSkewedJoin` hand the whole tree back to it after rewriting some
other join, all within
+ * one `queryStagePreparationRules` pass. So a join child can arrive as
+ * `GroupPartitionsExec(SortExec(GroupPartitionsExec(scan)))`, and planning
from anything but the
+ * scan would derive the alignment from an already aligned layout and
duplicate rows.
*
* The descent only traverses a `GroupPartitionsExec` and a *local*
`SortExec`. That bound is a
* decision, not an omission: a `GroupPartitionsExec` hidden behind any
other node belongs to a
- * different operator, and reusing it would move that operator's alignment.
Instrumentation of
+ * different operator, and peeling it would undo that operator's alignment.
Instrumentation of
* the descent over `KeyGroupedPartitioningSuite` found these non-`SortExec`
shapes hiding a
* node: `Project > SortMergeJoin > Sort > GroupPartitions` and `Project >
Filter > Window >
* WindowGroupLimit > GroupPartitions`, where refusing to descend is right
every time. A global
* `SortExec` also stops the descent: it requires `OrderedDistribution`,
which a
* `KeyedPartitioning` can satisfy (behind
`spark.sql.sources.v2.bucketing.sorting.enabled`)
- * through a `GroupPartitionsExec` built to emit the partition keys in
sorted order, and
- * reusing that node for a join would destroy the ordering it exists to
provide.
- */
- private def innermostGroupPartition(
- plan: SparkPlan): Option[(GroupPartitionsExec, SparkPlan => SparkPlan)]
= plan match {
- case g: GroupPartitionsExec =>
- // When groupings stack, the outer one is the wrap this invocation's
distribution step
- // just added; the one below is inherited from an earlier pass and owns
the alignment to
- // preserve. Keep the descent below the outer node and drop it.
- innermostGroupPartition(g.child).orElse(Some((g, identity[SparkPlan])))
- case s: SortExec if !s.global =>
- innermostGroupPartition(s.child).map { case (g, rebuild) =>
- (g, (newChild: SparkPlan) => s.withNewChildren(Seq(rebuild(newChild))))
- }
- case _ => None
- }
-
- /**
- * Rewrites the innermost `GroupPartitionsExec` in `plan` with `f` and drops
any redundant
- * grouping stacked above it, per the descent of
[[innermostGroupPartition]]. Returns `None`
- * when `plan` holds no `GroupPartitionsExec`, leaving it to the caller to
create one.
- *
- * This is what makes the rule idempotent for storage-partitioned joins.
`EnsureRequirements`
- * is re-run on plans it already produced: `AdaptiveSparkPlanExec` builds
one instance of this
- * rule, and `ConvertSortMergeJoinToShuffledHashJoin` and
`OptimizeSkewedJoin` hand the whole
- * tree back to it after rewriting some other join, all within one
- * `queryStagePreparationRules` pass. A join child then arrives as
- * `SortExec(GroupPartitionsExec(...))` rather than a bare scan, and the
distribution step adds
- * a plain `GroupPartitionsExec` on top, because a partially clustered
`KeyedPartitioning`
- * reports `isGrouped = false` by design and so is only satisfied "after
grouping". Rewriting
- * that outer node instead of the one below it re-derives the alignment from
an already-aligned
- * layout and duplicates rows; descending to the innermost node and dropping
what sits above it
- * reproduces the plan a single pass would have produced.
+ * through a `GroupPartitionsExec` built to emit the partition keys in
sorted order, and peeling
+ * that node would destroy the ordering it exists to provide.
*
- * Dropping a grouping is safe because only `applyGroupPartitions` calls
this, reached from
- * `checkKeyGroupCompatible`, which runs for joins alone: every
`GroupPartitionsExec` a join
- * child carries is this rule's own. A single-child operator genuinely needs
its non-grouped
- * input grouped and takes the wrap in the children loop instead;
`withJoinKeyPositions`, which
- * other multi-child operators reach, does not reuse at depth.
- */
- private[exchange] def rewriteGroupPartitions(plan: SparkPlan)(
- f: GroupPartitionsExec => GroupPartitionsExec): Option[SparkPlan] =
- innermostGroupPartition(plan).map { case (g, rebuild) =>
- val rewritten = f(g)
- rewritten.copyTagsFrom(g)
- rebuild(rewritten)
- }
-
- /**
- * Unwraps the groupings and local sorts this rule inserted over a child,
down to the
- * pre-alignment plan, per the descent of [[innermostGroupPartition]].
Peeling one level stops
- * at the local sort this rule added, leaving the earlier pass's alignment
in place.
- */
- private def unwrapGroupPartitions(plan: SparkPlan): SparkPlan =
- innermostGroupPartition(plan).map(_._1.child).getOrElse(plan)
-
- /**
- * Applies or updates `GroupPartitionsExec` with the given parameters.
+ * A local sort that is peeled off is re-added by the ordering step at the
end of
+ * `ensureDistributionAndOrdering`, which is what put it there in the first
place.
*
- * Reuses the node this rule inserted over the join child in an earlier
pass, per the descent
- * of [[innermostGroupPartition]], and creates a new one when the child
carries none.
+ * The two callers are `checkKeyGroupCompatible` and the shuffle step, both
on the co-partitioned
+ * path, so every `GroupPartitionsExec` this reaches is one the rule put
there itself. A
+ * single-child operator genuinely needs its non-grouped input grouped and
keeps the wrap
+ * `resolveChild` gave it. `withJoinKeyPositions`, which other multi-child
operators reach, does
+ * not descend at all.
*/
- private def applyGroupPartitions(
- plan: SparkPlan,
- joinKeyPositions: Option[Seq[Int]],
- mergedPartitionKeys: Seq[(InternalRowComparableWrapper, Int)],
- reducers: Option[Seq[Option[KeyReducer]]],
- distributePartitions: Boolean): SparkPlan = {
- rewriteGroupPartitions(plan) { g =>
- g.copy(
- joinKeyPositions = g.joinKeyPositions.orElse(joinKeyPositions),
- expectedPartitionKeys = Some(mergedPartitionKeys),
- // Unlike `joinKeyPositions`, these need no `orElse`. A re-run with
reducers never reaches
- // here. Both sides then report the same reduced keys, so
`compatibleAsIs` holds and the
- // push branch is skipped.
- reducers = reducers,
- distributePartitions = distributePartitions)
- }.getOrElse {
- GroupPartitionsExec(plan, joinKeyPositions, Some(mergedPartitionKeys),
reducers,
- distributePartitions)
+ private[exchange] def peelGroupPartitions(plan: SparkPlan): SparkPlan = {
+ // `None` unless a `GroupPartitionsExec` is actually down there. A local
sort is only this
+ // rule's to drop when it sits over one, otherwise it is the user's
`sortWithinPartitions` and
+ // peeling it would lose an ordering nothing puts back.
+ def peel(p: SparkPlan): Option[SparkPlan] = p match {
+ case g: GroupPartitionsExec => Some(peel(g.child).getOrElse(g.child))
+ case s: SortExec if !s.global => peel(s.child)
+ case _ => None
}
+ peel(plan).getOrElse(plan)
}
/**
* Applies join key positions to a plan by wrapping or updating
GroupPartitionsExec.
*
- * Unlike `applyGroupPartitions`, this does not descend: it serves every
multi-child operator,
+ * Unlike `peelGroupPartitions`, this does not descend. It serves every
multi-child operator,
* not just joins, so a `GroupPartitionsExec` below the top is not known to
be this rule's own.
*/
private[exchange] def withJoinKeyPositions(plan: SparkPlan, positions:
Seq[Int]): SparkPlan = {
plan match {
case g: GroupPartitionsExec =>
- val newGroupPartitions = g.copy(joinKeyPositions = Some(positions))
+ // Rebuilt rather than copied: the positions are an input to the
node's grouping, and a
+ // `copy` would keep the grouping derived from the old ones.
Rebuilding from `g.child`
+ // discards whatever else `g` carried, and the assert says what that
may be.
+ //
+ // `positions` index the layout `g` reports, while `g.child` holds the
raw partition
+ // expressions, so they are composed rather than replaced. The two
index spaces differ
+ // whenever `g` already projects, which happens on two paths.
`resolveChild` projects a
+ // co-partitioned child onto its cluster keys before the pairing
declines, and the two
+ // derivations disagree about which positions those are.
`positionsCoveringClusterKeys`
+ // also keeps an expression that *is* a cluster key, where
`KeyedShuffleSpec.keyPositions`
+ // reads an expression's reference. And a re-run reads the positions
off a report an
+ // earlier pass already projected.
+ assert(g.expectedKeyCount.isEmpty && g.reducers.isEmpty &&
!g.distributePartitions,
+ "expected a grouping this rule inserted for a co-partitioned child")
+ val composed = g.joinKeyPositions.fold(positions)(positions.map(_))
Review Comment:
`positions.map(_)` is `positions.map(oldPositions)` with the Seq used as a
Function1 (SeqOps is a PartialFunction). Typer output: ((x$1: Seq[Int]) =>
positions.map[Int](x$1)). That is the composition the comment intends, but it
reads as positions.map(identity) and would stop compiling if either side became
an Array or a Set. positions.map(g.joinKeyPositions.get) says the same thing
explicitly.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -2576,6 +2552,185 @@ class EnsureRequirementsSuite extends
SharedSparkSession {
s"EnsureRequirements must satisfy the required
distribution:\n${newChild.treeString}")
}
+ test("SPARK-59289: a grouped side is paired on its own key order, not a
sorted one") {
+ // Both sides are grouped, so neither gets a node to reorder anything, and
the left's keys are
+ // not sorted. Planning the pairing against `toGrouped`, which sorts,
would claim an order the
+ // left does not have: the two would look compatible, the left would keep
its own report and
+ // the right would keep the sorted one, and the join would carry two
`KeyedPartitioning`s over
+ // different keys. `toGrouped` is therefore only for a source that is not
grouped.
+ val unsorted = Seq(InternalRow(3), InternalRow(4), InternalRow(1),
InternalRow(2))
+ val sorted = Seq(InternalRow(1), InternalRow(2), InternalRow(3),
InternalRow(4))
+ val left = DummySparkPlan(outputPartitioning =
KeyedPartitioning(Seq(exprA), unsorted))
+ val right = DummySparkPlan(outputPartitioning =
KeyedPartitioning(Seq(exprB), sorted))
+ val smj = SortMergeJoinExec(exprA :: Nil, exprB :: Nil, Inner, None, left,
right)
+
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "true") {
+ val planned = EnsureRequirements.apply(smj)
+ val declared = planned.children.map { child =>
+ child.outputPartitioning.asInstanceOf[Expression]
+ .collectFirst { case k: KeyedPartitioning => k.partitionKeys }
+ }
+ assert(declared.forall(_.isDefined), s"both sides stay
keyed:\n${planned.treeString}")
+ assert(declared.head == declared(1),
+ s"both sides must declare one key sequence:\n${planned.treeString}")
+ }
+ }
+
+ test("SPARK-59289: pushing join key positions into a node re-derives its
grouping") {
+ // The positions are an input to the node's grouping, so
`withJoinKeyPositions` rebuilds through
+ // the factory. A `copy` would keep the grouping the old positions
produced, and nothing about
+ // the node would look wrong.
+ val keys = Seq(InternalRow(1, 1), InternalRow(1, 2))
+ val leaf = DummySparkPlan(
+ outputPartitioning = KeyedPartitioning(Seq(exprA, exprB), keys))
+ val coalescing = GroupPartitionsExec(leaf)
+ assert(coalescing.outputPartitioning.numPartitions == 2, "test setup:
nothing merged yet")
+
+ val projected = EnsureRequirements.withJoinKeyPositions(coalescing, Seq(0))
+ assert(projected.outputPartitioning.numPartitions == 1,
+ "projecting [a, b] onto [a] merges the two partitions that share a = 1")
+ }
+
+ test("SPARK-59289: single-partition children still honour a required
partition count") {
+ // Both sides are already `SinglePartition` and small, so there is no
co-partitioning to
+ // arrange. Each child still has to satisfy its own requirement, and a
stateful operator's
+ // `ClusteredDistribution` carries a partition count that one partition
does not meet.
+ val logicalPlan = StatsTestPlan(Nil, 1L, AttributeMap.empty, Some(1L))
+ val left = DummySparkPlan(outputPartitioning = SinglePartition)
+ left.setLogicalLink(logicalPlan)
+ val right = DummySparkPlan(outputPartitioning = SinglePartition)
+ right.setLogicalLink(logicalPlan)
+ val parent = DummySparkPlan(
+ children = Seq(left, right),
+ requiredChildDistribution = Seq(
+ ClusteredDistribution(Seq(exprA), requiredNumPartitions = Some(5)),
+ ClusteredDistribution(Seq(exprC), requiredNumPartitions = Some(5))),
+ requiredChildOrdering = Seq(Nil, Nil))
+
+ val planned = EnsureRequirements.apply(parent)
+ assert(ValidateRequirements.validate(planned),
+ s"each child must satisfy its own distribution:\n${planned.treeString}")
+ }
+
+ test("SPARK-59289: a local sort with no grouping under it is left alone") {
+ // `peelGroupPartitions` peels the local sorts this rule added over a
grouping. A sort with no
+ // grouping beneath it is the user's `sortWithinPartitions`, and the
ordering step only puts
+ // back the ordering the operator required, so peeling it would lose that
ordering for good.
+ val keys = Seq(InternalRow(1), InternalRow(2))
+ val leaf = DummySparkPlan(outputPartitioning =
KeyedPartitioning(Seq(exprA), keys))
+ val userSort = SortExec(Seq(SortOrder(exprB, Ascending)), global = false,
leaf)
+ assert(EnsureRequirements.peelGroupPartitions(userSort) === userSort)
+
+ // Over a grouping it is this rule's own, and both come off.
+ val ruleSort = SortExec(Seq(SortOrder(exprB, Ascending)), global = false,
+ GroupPartitionsExec(leaf))
+ assert(EnsureRequirements.peelGroupPartitions(ruleSort) === leaf)
+ }
+
+ test("SPARK-59289: a partitioning is never projected onto no position") {
+ // The co-partitioned path takes its projection positions from
`KeyedShuffleSpec.keyPositions`,
+ // which maps an expression's single *reference* onto the clustering.
`keysSatisfy` also
+ // accepts an expression that *is* a cluster key, through
+ // `KeyedPartitioning.positionsCoveringClusterKeys`. This is the shape
where the two disagree:
+ // the join keys name `years(ts)`, so the spec is built, and then no
position covers a cluster
+ // key.
+ //
+ // Projecting onto no position coalesces every partition into one. Before
the guard in
+ // `createShuffleSpec` this planned as a storage-partitioned join over a
single partition per
+ // side, two `GroupPartitionsExec`s and no shuffle, which is correct and
ruinous. The spec now
+ // comes back unprojected, `areKeysCompatible` refuses it, and the join
takes the shuffle.
+ val exprTs = AttributeReference("ts", IntegerType)()
+ val transform = years(exprTs)
+ val keys = Seq(InternalRow(1), InternalRow(2))
+ val left = DummySparkPlan(outputPartitioning =
KeyedPartitioning(Seq(transform), keys))
+ val right = DummySparkPlan(outputPartitioning =
KeyedPartitioning(Seq(transform), keys))
+ val smj = SortMergeJoinExec(transform :: Nil, transform :: Nil, Inner,
None, left, right)
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ENABLED.key -> "true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true") {
+ val planned = EnsureRequirements.apply(smj)
+ assert(groupPartitionsNodes(planned).isEmpty,
+ "a grouping here could only be the one that collapses both sides onto
a single partition")
+ assert(planned.collect { case s: ShuffleExchangeExec => s }.size == 2,
+ "neither side can serve the distribution, so both are shuffled")
+ }
+ }
+
+ test("SPARK-59289: a second pass leaves a pairing this rule already made
alone") {
+ val aL = AttributeReference("aL", IntegerType)()
+ val bL = AttributeReference("bL", IntegerType)()
+ val bR = AttributeReference("bR", IntegerType)()
+ // The left is keyed on (aL, bL) and the join is on bL alone, so the
left's spec projects onto
+ // position 1, and the right, which reports nothing, is shuffled onto that
layout. The first
+ // pass therefore ends with a projecting grouping node over the left and a
keyed shuffle over
+ // the right. That is the shape the second pass used to re-pair: peeling
the node hands the
+ // pairing a raw left that projects again, so it stopped being
co-partitioned as it stands and
+ // both sides were aligned to a merged key set they already held.
+ val leftKeys = Seq(InternalRow(1, 10), InternalRow(2, 20))
+ val left = new DummySparkPlanWithBatchScanChild(
+ outputPartitioning = KeyedPartitioning(Seq(aL, bL), leftKeys))
+ val right = DummySparkPlan(outputPartitioning = UnknownPartitioning(3))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val smj = SortMergeJoinExec(Seq(bL), Seq(bR), Inner, None, left, right)
+ val planned = EnsureRequirements.apply(smj)
+ assert(groupPartitionsNodes(planned).map(_.joinKeyPositions) ===
Seq(Some(Seq(1))),
+ "test setup: the left keeps its own layout with a projection")
+ assert(planned.collect { case s: ShuffleExchangeExec => s }.size == 1,
+ "test setup: the right is shuffled onto it")
+
+ assert(EnsureRequirements.apply(planned) == planned,
Review Comment:
keepArrivedPairing only fires with partially clustered distribution off, and
the PR pins one shape. Since idempotency is now an explicit decision rather
than a property of the reuse-at-depth mechanism, could we get an assertion set
over a handful of SPJ queries, EnsureRequirements.apply(planned) == planned,
including a partially clustered one so the re-plan path runs? The 2713
applications / 0 differs number is a measurement, not a test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -561,34 +352,289 @@ case class GroupPartitionsExec(
val joinKeyStr = joinKeyPositions.map { p =>
s"JoinKeyPositions: ${truncatedString(p, "[", ", ", "]",
joinKeyMaxFields)}"
}.iterator
- val expectedStr = expectedPartitionKeys.map(ks => s"ExpectedPartitionKeys:
${ks.size}")
+ val expectedStr = expectedKeyCount.map(n => s"ExpectedPartitionKeys: $n")
val reducersStr = reducers.map { seq =>
val names = seq.map(_.map(_.reducer.displayName()).getOrElse("identity"))
s"Reducers: ${truncatedString(names, "[", ", ", "]", joinKeyMaxFields)}"
}
val distributeStr = Iterator(s"DistributePartitions:
$distributePartitions")
- // Rendered from the constructor field, as `DistributePartitions` above
is. Not from
- // `usesSortedMerge`, because that forces `grouping`, which can throw, and
this method feeds
- // `simpleString`, which `treeString` calls on error paths.
+ // Rendered from the constructor field, as `DistributePartitions` above
is, and not from
+ // `usesSortedMerge`, which reads the child's ordering. This method feeds
`simpleString`, which
+ // `treeString` calls on error paths, so it stays on what the node was
built with.
val sortedMergeStr = Iterator(s"SortedMerge: $enableSortedMerge")
joinKeyStr ++ expectedStr ++ reducersStr ++ distributeStr ++ sortedMergeStr
}
}
/**
- * What a [[GroupPartitionsExec]] computes once and reports from several
members: which of the
- * child's partitions each of its own is built from, and the layout that
describes them. Every
- * member is given the same `layout` instance, which is what
- * `PartitioningCollection.fromPartitionings` needs to return them untouched.
The last two fields
- * count the alignment's effect on the reads of the child's splits (see
`alignToExpectedKeys`), and
- * are 0 outside the alignment path.
+ * What a [[GroupPartitionsExec]] does to its child's partitions.
+ *
+ * @param partitions the index groups the node emits, each with the key it
stands for
+ * @param layout what the node reports about those groups. Every member of a
+ * [[PartitioningCollection]] is given this same instance, which
is what
+ * `PartitioningCollection.fromPartitionings` needs to return
them untouched.
+ * @param isIdentity whether the grouping leaves the declared keys and every
partition where they
+ * were. Nothing rewrote the keys, output partition i holds
exactly input
+ * partition i, and there is one output per input. That is
the only grouping that
+ * keeps a marked layout's undeclared rows at hash(key) %
numPartitions. A
+ * projection or reduction re-labels the groups into a
different key space, so
+ * even a grouping whose indices line up would pin the claim
to keys it no longer
+ * declares. A rewrite is rejected up front, covering a
narrowing
+ * projection, a reordering one, and any reducer slot. A
reducer slot counts as
+ * key-changing even though a conforming self-reducer cannot
rewrite a reachable
+ * key value, so the give-up there loses at most an
optimization. A grouping that
+ * drops trailing declared keys reads identity for every
group it keeps, but the
+ * partition count shrinks and the hash modulus with it.
+ * @param numPrunedPartitions the alignment's effect on the reads of the
child's splits, 0 outside
+ * the alignment path. See `alignToExpectedKeys`.
+ * @param numReplicatedPartitionReads as above.
*/
-private case class PartitionGrouping(
+private[sql] case class PartitionGrouping(
partitions: Seq[(InternalRowComparableWrapper, Seq[Int])],
layout: KeyLayout,
- keysRewritten: Boolean,
+ isIdentity: Boolean,
numPrunedPartitions: Int,
- numReplicatedPartitionReads: Int)
+ numReplicatedPartitionReads: Int) {
+
+ def isGrouped: Boolean = layout.isGrouped
+}
+
+private[sql] object GroupPartitionsExec {
+
+ /**
+ * Builds a node over `child`, deriving `grouping` and `outputPartitioning`
from the parameters.
+ *
+ * **Both are derived, and neither `copy` nor the generated `apply`
re-derives them**, so a change
+ * to `child`, `joinKeyPositions`, `expectedPartitionKeys`, `reducers` or
`distributePartitions`
Review Comment:
The doc lists expectedPartitionKeys among the parameters that have to come
back through the factory, but that name only exists on apply; the field it
feeds is expectedKeyCount (:78). One word so a reader can find the field.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1252,7 +1336,26 @@ case class PartitioningCollection(partitionings:
Seq[Partitioning])
partitionings.exists(_.satisfies(required))
override def createShuffleSpec(distribution: ClusteredDistribution):
ShuffleSpec = {
- val filtered = partitionings.filter(_.satisfies(distribution))
+ // `maySatisfyAfterGrouping`, not `satisfies`. A spec says what its
partitioning could
+ // co-partition on, and a `KeyedPartitioning` that needs a
`GroupPartitionsExec` first still
+ // can. The strict question would drop a member whose keys are coarser
than the operation's,
+ // which is narrower than what this filter admitted before `satisfies`
became strict.
+ //
+ // `ValidateRequirements` builds a spec from a finished plan through here
too, and this stays
+ // sound for it. A member that is not grouped reports its own ungrouped
keys, duplicates and
+ // all, so it can only pair with a side holding the same sequence, which
is co-partitioned.
+ // `KeyedShuffleSpec.canCreatePartitioning` gates on `isGrouped`
separately, so such a member
+ // is never a shuffle template either.
+ //
+ // Every admitted member has to stay, because `isCompatibleWith` answers
for any of them and
+ // the collection cannot know which one the other side matched. That has a
cost worth knowing:
+ // `KeyedShuffleSpec.canCreatePartitioning` is false without
`v2BucketingShuffleEnabled` and
+ // `ShuffleSpecCollection.canCreatePartitioning` is a `forall`, so a
groupable keyed member
+ // beside a usable non-keyed one would cost the collection its role as a
shuffle template. No
+ // operator is known to report that mixture, since `EnsureRequirements`
groups a keyed child
+ // before it can reach a join's output.
+ val filtered =
Review Comment:
This filter also feeds ValidateRequirements (ValidateRequirements:57-68),
which is the gate that makes AQE drop a rewrite that breaks co-partitioning
(AdaptiveSparkPlanExec:209), and it now admits members whose spec describes a
layout the child does not have.
The soundness argument in the comment holds for two grouped sides. For an
ungrouped side a duplicated key row means the group is split across the
partitions carrying that key, and two such sides are co-partitioned only when
their splits line up (partial clustering makes them line up by construction; an
ordinary ungrouped source does not).
I could not find a producer of a collection holding an ungrouped keyed
member next to a satisfying member: scans report a bare KeyedPartitioning
(DataSourceV2ScanExecBase:120-134), and checkKeyedPartitioningInvariant forces
all keyed members of a collection to share one layout. One producer worth
ruling out explicitly, since it is the natural guess: unions do not report this
mixture. UnionExec's all-keyed arm returns a bare KeyedPartitioning.concat
(basicPhysicalOperators.scala:994, :1006) and its co-located arm keeps only
HashPartitioningLike and SinglePartition members
(basicPhysicalOperators.scala:1014, :1016-1020), so a keyed child beside a hash
child falls back to UnknownPartitioning.
So this is not a finding. But given the filter is now load-bearing for a
safety gate, please either narrow it to what the base admitted, or pin the
shape as unreachable with a test. If narrowing: satisfies || (keyed &&
isGrouped && keysMaySatisfy) restores the base admission set while keeping the
new ShuffleSpecSuite test, which needs a fallback because strict filtering
would leave the collection empty.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
