cloud-fan commented on code in PR #58659:
URL: https://github.com/apache/spark/pull/58659#discussion_r4038851069
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -561,34 +415,286 @@ 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, so a
+ * projection, a reduction or a dropped trailing key all
fail it: each moves the
+ * claim into a different key space or changes the modulus.
A reducer slot counts
+ * as key-changing even where a conforming self-reducer
could not rewrite a
+ * reachable key value, so the give-up there loses at most
an optimization.
+ * @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 `plannedPartitioning`
from the parameters.
+ *
+ * **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
+ * an input to either, which is why `tryEnableSortedMerge` may `copy` it.
+ *
+ * Two other `copy` calls in this file are deliberate.
`withNewChildInternal` carries both fields
+ * over a child rewrite, and a child that turns out to report something else
is what
+ * `outputPartitioning` answers for, so the carried pair is never reported
as if it still held.
+ * And `doCanonicalize` rewrites `reducers` without re-deriving, which holds
because neither field
Review Comment:
**Nit (P3):** This factory description does not match the values it builds.
`plannedPartitioning` contains `AttributeReference` exprIds and is normalized
during canonicalization. `computeGrouping` sorts when expected keys are absent;
when they are present it distributes or replicates according to
`distributePartitions`, and partial clustering can use either branch. Could we
describe those actual conditions instead of the no-exprId and mode-split claims?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -327,8 +114,307 @@ case class EnsureRequirements(
.getOrElse(SortExec(requiredOrdering, global = false, child = child))
}
}
+ }
+
+ /**
+ * What `child` needs to satisfy `distribution` on its own: nothing, a
[[GroupPartitionsExec]], a
+ * broadcast, or a shuffle.
+ */
+ private def resolveChild(
+ child: SparkPlan,
+ distribution: Distribution,
+ shuffleOrigin: ShuffleOrigin): SparkPlan = {
+ // Ask what the child's partitioning still needs to satisfy the
distribution
+ val (otherSatisfies, keyed) =
+ splitKeyedPartitionings(child.outputPartitioning, distribution)
+
+ // If a non-KeyedPartitioning already satisfies, no changes needed
+ if (otherSatisfies) {
+ child
+ } else {
+ keyed match {
+ case Some(resolution) =>
+ (distribution, resolution) match {
+ case (o: OrderedDistribution, _) =>
+ // OrderedDistribution requires grouped KeyedPartitioning with
sorted keys
+ // according to the distribution's ordering.
+ val satisfyingKeyedPartitioning = resolution.fold(identity, _._1)
+ // The single-column invariant in
KeyedPartitioning.supportsExpressions guarantees
+ // one attribute per partition expression.
+ val attrs =
satisfyingKeyedPartitioning.expressions.flatMap(_.references)
+ val keyRowOrdering = RowOrdering.create(o.ordering, attrs)
+ 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) }) {
+ child
+ } else {
+ // Use distributePartitions to spread splits across expected
partitions
+ val sortedGroupedKeys = keys
+ .groupBy(identity).view.mapValues(_.size)
+ .toSeq.sortBy(_._1)(keyOrdering)
+ GroupPartitionsExec(child,
+ expectedPartitionKeys = Some(sortedGroupedKeys),
+ distributePartitions = true
+ )
+ }
+
+ // A KeyedPartitioning satisfies the distribution and a node would
change nothing
+ case (_, scala.Left(_)) =>
+ child
+
+ // A KeyedPartitioning satisfies the distribution only after a
GroupPartitionsExec:
+ // to coalesce duplicate partition keys, to project the partition
keys down to the
+ // cluster keys, or both. The positions to project to come from
whichever member
+ // of the child's partitioning leaves the most partitions.
+ case (_, scala.Right((_, positions))) =>
+ GroupPartitionsExec(child, joinKeyPositions = positions)
+ }
+
+ case None =>
+ // No partitioning satisfies - need broadcast or shuffle
+ val numPartitions = distribution.requiredNumPartitions
+ .getOrElse(conf.numShufflePartitions)
+ distribution match {
+ case BroadcastDistribution(mode) =>
+ BroadcastExchangeExec(mode, child)
+ case _: StatefulOpClusteredDistribution =>
+ ShuffleExchangeExec(
+ distribution.createPartitioning(numPartitions), child,
+ REQUIRED_BY_STATEFUL_OPERATOR)
+ case _ =>
+ ShuffleExchangeExec(
+ distribution.createPartitioning(numPartitions), child,
shuffleOrigin)
+ }
+ }
+ }
+ }
+
+ /**
+ * Resolves each co-partitioned child on its own, which is what is left once
nothing is arranged
+ * between them.
+ */
+ private def resolveEachChild(
+ children: Seq[SparkPlan],
+ coPartitioned: Seq[Option[ClusteredDistribution]],
+ shuffleOrigin: ShuffleOrigin): Seq[SparkPlan] =
children.zip(coPartitioned).map {
+ case (child, required) => required.fold(child)(resolveChild(child, _,
shuffleOrigin))
+ }
+
+ /**
+ * Decides the children an operator co-partitions, together, because neither
side's answer stands
+ * on its own.
+ *
+ * Three shapes need nothing arranged between the children, and every one of
them still resolves
+ * each child on its own: a child can owe its distribution a
[[GroupPartitionsExec]] or a
+ * partition count whether or not it lines up with its siblings.
+ * 1. Every side is a single small partition, so they line up whatever
they hold.
+ * 2. The sources report layouts a storage-partitioned join can align,
which
+ * `checkKeyGroupCompatible` decides and builds end to end.
+ * 3. Both sides pass shuffle partition ids through, on compatible specs.
+ *
+ * Failing all three, `shuffleToCoPartition` shuffles whichever children are
not aligned onto one
+ * that is.
+ */
+ private def coPartitionChildren(
+ parent: Option[SparkPlan],
+ children: Seq[SparkPlan],
+ coPartitioned: Seq[Option[ClusteredDistribution]],
+ shuffleOrigin: ShuffleOrigin): Seq[SparkPlan] = {
+ // Every way of lining the children up as they are still needs this, and
only the shuffle below
+ // reads more than it, hence `lazy`: the storage-partitioned join plans
from the children as
+ // their sources report them, so when it succeeds nothing here is asked
for.
+ lazy val resolved = resolveEachChild(children, coPartitioned,
shuffleOrigin)
+
+ // Special case: if all sides of the join are single partition and it's
physical size less than
+ // or equal spark.sql.maxSinglePartitionBytes.
+ val preferSinglePartition = children.zip(coPartitioned).forall {
+ case (child, Some(_)) =>
+ child.outputPartitioning == SinglePartition &&
+ child.logicalLink
+ .forall(_.stats.sizeInBytes <=
conf.getConf(SQLConf.MAX_SINGLE_PARTITION_BYTES))
+ case _ => true
+ }
+
+ if (preferSinglePartition) {
+ // Nothing to arrange between them, but each still has to satisfy its
own requirement. A
+ // `SinglePartition` child does unless the distribution asks for a
partition count, which a
+ // stateful operator's does.
+ resolved
+ } else {
+ // The two checks below decide on a pair: the operator, and what each of
its children has to
+ // satisfy. Spark doesn't support multi-way join at the moment, so a
parent that
+ // co-partitions has exactly two children and both of them are in the
decision.
+ val linedUp = (parent, coPartitioned) match {
+ case (Some(operator), Seq(Some(leftRequired), Some(rightRequired))) =>
+ // key group compatibility check
+ checkKeyGroupCompatible(
+ operator, children.head, leftRequired, children(1), rightRequired)
+ // If key group check fails, check ShufflePartitionIdPassThrough
compatibility. That
+ // one reads the children once each has satisfied its distribution
on its own.
+ .orElse(Option.when(checkShufflePartitionIdPassThroughCompatible(
+ resolved.head, leftRequired, resolved(1),
rightRequired))(resolved))
+ case _ => None
+ }
+ linedUp.getOrElse(shuffleToCoPartition(resolved, coPartitioned))
+ }
+ }
+
+ /**
+ * The layout the co-partitioned children are lined up on, and the member
each of them pairs with
+ * it. A child with no pair is one that has to be shuffled onto `member`.
+ */
+ private case class CoPartitionTarget(
+ member: LeafShuffleSpec,
+ pairedMembers: Seq[Option[LeafShuffleSpec]])
- children
+ /**
+ * Lines the co-partitioned children up by shuffle, which is what is left
once no way of taking
+ * them as they are has worked. One child's layout is picked and the others
are shuffled onto it.
+ *
+ * The children arrive having each satisfied its distribution on its own.
+ */
+ private def shuffleToCoPartition(
+ children: Seq[SparkPlan],
+ coPartitioned: Seq[Option[ClusteredDistribution]]): Seq[SparkPlan] = {
+ // One spec per co-partitioned child, in child order, and nothing for a
child that is not in
+ // the decision.
+ val specs = children.zip(coPartitioned).map { case (child, required) =>
+ required.map(child.outputPartitioning.createShuffleSpec)
+ }
+ val target = pickCoPartitionTarget(children, specs)
+ // Nothing pairs when no child can serve as the layout, and then each of
them is shuffled below
+ // on its own required distribution.
+ val pairedMembers =
target.map(_.pairedMembers).getOrElse(Seq.fill(children.length)(None))
+
+ children.lazyZip(coPartitioned).lazyZip(pairedMembers).map {
+ case (child, None, _) =>
+ child
+
+ // The positions come from this child's own paired member, since they
index into its own
+ // partition expressions -- the layout picked only says which member of
it the two sides
+ // agreed on.
+ //
+ // The storage-partitioned join above declined, but one can still be had
by shuffling the
+ // other side onto this one's keys (see the last case below). So the
partitioned side's scan
+ // has to end up grouped by those keys, which is what pushing the
positions in does. They
+ // index into the child's own report, which already carries whatever
projection
+ // `resolveChild` gave it, so `withJoinKeyPositions` composes the two
rather than replacing.
+ case (child, Some(_), Some(KeyedShuffleSpec(_, _,
Some(joinKeyPositions)))) =>
+ withJoinKeyPositions(child, joinKeyPositions)
+
+ // Paired as it stands, so there is nothing to shuffle and nothing to
push in.
+ case (child, Some(_), Some(_)) =>
+ child
+
+ case (child, Some(required), None) =>
+ val newPartitioning = target.map { picked =>
+ // Use the picked layout to create a new partitioning to re-shuffle
this child
+ picked.member.createPartitioning(required.clustering)
+ }.getOrElse {
+ // No layout was picked, so we create default partitioning from the
required distribution
+ val numPartitions = required.requiredNumPartitions
+ .getOrElse(conf.numShufflePartitions)
+ required.createPartitioning(numPartitions)
+ }
+
+ child match {
+ case s: ShuffleExchangeExec =>
+ s.copy(outputPartitioning = newPartitioning)
+ case gpe: GroupPartitionsExec =>
+ // Strip every grouping this rule inserted (they can stack on a
re-run): a
+ // replicating one repeats every row, so none of them may feed the
shuffle.
+ ShuffleExchangeExec(newPartitioning, peelGroupPartitions(gpe))
+ case _ => ShuffleExchangeExec(newPartitioning, child)
+ }
+ }
+ }
+
+ /**
+ * Picks the layout to line the co-partitioned children up on. It is one of
their own, so that at
+ * least the child offering it keeps its partitioning and only the others
are shuffled. `None`
+ * when no child can serve, and then they all take a shuffle.
+ *
+ * Find out the shuffle spec that gives better parallelism. Currently this
is done by
Review Comment:
**Nit (P3):** Several contracts in this new explanatory block are broader
than the code: target selection ranks avoiding an exchange before partition
count; the co-partitioned dispatch is binary only because it handles concrete
join cases; peeled groupings are from an earlier pass; only one tested shape
needs `keepArrivedPairing`; grouped sources can still receive
projection/alignment nodes; `satisfies` can read partition keys; and the
no-node `AdmittedMember` interpretation does not hold generically for
`OrderedDistribution`. Please narrow these statements to the implemented
branches so later edits do not rely on false planner invariants.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -497,6 +350,7 @@ case class GroupPartitionsExec(
override def supportsColumnar: Boolean = child.supportsColumnar &&
!usesSortedMerge
override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ checkChildStillMatches()
Review Comment:
**Nit (P3):** The stale-child check is duplicated in `doExecute` and
`doExecuteColumnar`, but the new execution assertion calls only `execute()`.
Removing this columnar call would leave the test green. Could the owning suite
drive a genuinely columnar child through `executeColumnar()` and assert the
same refusal, so either execution path losing the guard has a regression signal?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -834,41 +868,100 @@ case class KeyedPartitioning(
KeyedPartitioning.reduceKeys(partitionKeys, keyDataTypes, reducers)
override def satisfies0(required: Distribution): Boolean = {
- nonGroupedSatisfies(required) || (isGrouped && keysSatisfy(required))
+ super.satisfies0(required) || (isGrouped && keysSatisfy(required))
}
- /** The first of the four questions the class doc lists. */
- private def nonGroupedSatisfies(required: Distribution): Boolean =
super.satisfies0(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:
+ *
+ * - one of its *references* is a cluster key, as a `bucket(4, a)` covers
`a`;
+ * - the expression *itself* is a cluster key, as a `years(ts)` under a
clustering that names
+ * `years(ts)` rather than `ts`.
+ *
+ * `EnsureRequirements.resolveKeyedPartitioning` and `keysSatisfy` both read
this, so the
+ * predicate deciding whether a projection is needed and the one choosing
the positions to
+ * project onto cannot drift apart.
+ *
+ * Not the same question as `KeyedShuffleSpec.keyPositions`, and the two
must not be merged. That
+ * one answers, per expression, which cluster key the expression is a
function *of*. That is
+ * what `KeyedShuffleSpec.createPartitioning` needs, since it rebuilds the
expression over the
+ * other side's clustering with `te.copy(children = ...
clustering(positionSet.head))`. Handing it
+ * the second form above would rebuild a `years(ts)` clustered on
`years(ts)` as
+ * `years(years(ts))`.
+ * The two coincide exactly where every expression has a single reference
and no expression is
+ * itself a cluster key. Where they diverge, `createShuffleSpec` refuses to
project rather than
+ * projecting onto nothing.
+ */
+ def positionsCoveringClusterKeys(required: ClusteredDistribution): BitSet =
+ expressions.zipWithIndex.collect {
+ case (e, i) if required.isClusterKey(e) ||
e.references.exists(required.isClusterKey) => i
+ }.to(BitSet)
+
+ /**
+ * Whether every partition expression is a function of cluster keys alone,
so that rows sharing
+ * a cluster key share a partition and nothing needs projecting.
+ *
+ * The sibling of `positionsCoveringClusterKeys`, and the two differ in one
word: this quantifies
+ * over an expression's references with `forall`, that one with `exists`.
Both are right for what
+ * they are asked. A `bucket(4, b, c)` over a clustering naming `b` alone
*carries* a cluster
+ * key, so its position survives a projection, yet two rows sharing `b` land
in different buckets
+ * when their `c` differs, so it is not a function of the cluster keys.
Deliberately not
+ * `positionsCoveringClusterKeys(required).size == expressions.length`.
+ */
+ private def isFunctionOfClusterKeys(required: ClusteredDistribution):
Boolean =
+ expressions.forall { e =>
+ required.isClusterKey(e) || e.references.forall(required.isClusterKey)
+ }
+
+ /**
+ * Whether a [[GroupPartitionsExec]] may project these keys down to the
cluster keys at all. Only
+ * `spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys`
permits it, and a marked
+ * layout cannot survive it, since the projection coarsens the declared set
the out-of-set routing
+ * speaks for (see the `@param`). Whole keys co-locate there and subsets do
not, so a window or
+ * aggregate keyed on a strict subset of the partition columns must still
shuffle. The
+ * single-reference requirement is what makes keeping a position sound at
all. A kept expression
+ * is then a function of one cluster key, so coalescing on the projected
keys cannot put rows that
+ * share a cluster key on different partitions.
+ *
+ * Read by `keysSatisfy` and `keysCanSatisfy`, which is the point. They ask
different things once
+ * the projection is permitted, so the permission has to be one rule rather
than two spellings.
+ */
+ private def mayProjectToClusterKeys(required: ClusteredDistribution):
Boolean =
+ !required.requireAllClusterKeys && !mayContainUnknownPartitionKeys &&
+ SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys &&
+ expressions.forall(_.references.size == 1)
- /** The second of the four questions the class doc lists. */
+ /**
+ * The strict question of the family the class doc lists, for a
[[ClusteredDistribution]]. `true`
Review Comment:
**Nit (P3):** These predicate and mapping claims need qualification.
`keysSatisfy` does not inspect `isGrouped`, so an admitted layout can still
need grouping; with clustering `[years(ts), ts]`, `years(ts)` is itself a
cluster key while its `ts` reference maps to the second position. Also,
Catalyst Scaladoc cannot resolve `[[GroupPartitionsExec]]`, which exists only
in downstream sql/core. Please state the admitted cases precisely and use
module-resolvable wording for that node.
--
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]
