dongjoon-hyun commented on code in PR #58262:
URL: https://github.com/apache/spark/pull/58262#discussion_r3881911759
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala:
##########
@@ -415,6 +415,12 @@ object ShuffleExchangeExec {
// .createPartitioning` substitutes this side's expressions but
retains the other side's
// `partitionKeys`, whose wrappers may carry a schema that differs in
struct field names,
// which wrapper equality would reject.
+ //
+ // This is the one path that declares the keys at
`expressionDataTypes` rather than at
+ // `keyDataTypes`: the lookup keys are evaluated from the expressions
per row, so both sides
+ // of the comparison have to be the expressions' types. A partitioning
whose keys a reducer
+ // rewrote cannot be shuffled onto at all -- its stored keys live in
the reduced key space
+ // while the evaluated ones do not, and no choice of declared types
brings the two together.
val wrapperFactory = InternalRowComparableWrapper
Review Comment:
**[correctness]** This comment asserts that a reducer-rewritten partitioning
"cannot be shuffled onto at all", but nothing enforces it:
`KeyedShuffleSpec.canCreatePartitioning` checks only `isGrouped` and the
expression shapes, never `keyDataTypes == expressionDataTypes`.
Moreover, this PR removes the accidental planning-time fail-fast that used
to stop this path: pre-PR, `createShuffleSpec`'s subset path read the keys at
`expressionDataTypes` and threw CCE at planning for a reduced partitioning;
post-PR it succeeds at `keyDataTypes`, so under `allowCompatibleTransforms` +
`v2BucketingShuffleEnabled` + `allowKeysSubsetOfPartitionKeys`, an SPJ output
(`TimestampType`-declared expressions over reduced `IntegerType` year keys) can
become `bestSpec`, the other side gets shuffled onto it, and the `valueMap`
below wraps the Integer keys at `expressionDataTypes` — CCE at execution, or
silent misrouting for same-width reductions. The failure moves from planning
time to execution time.
A `keyDataTypes == expressionDataTypes` clause in `canCreatePartitioning`
would make this comment true and restore the planning-time fallback to a
shuffle.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -62,89 +72,91 @@ case class EnsureRequirements(
shuffleOrigin: ShuffleOrigin): Seq[SparkPlan] = {
assert(requiredChildDistributions.length == originalChildren.length)
assert(requiredChildOrderings.length == originalChildren.length)
+ // Get the indexes of children which have specified distribution
requirements and need to be
+ // co-partitioned.
+ val childrenIndexes = requiredChildDistributions.zipWithIndex.filter {
+ case (_: ClusteredDistribution, _) => true
+ case _ => false
+ }.map(_._2)
+ // The multi-child co-partitioning block below is where the projected-key
`GroupPartitionsExec`
+ // belongs for an operator that co-partitions two children -- a
storage-partitioned join gets it
+ // from `checkKeyGroupCompatible`, anything else from
`withJoinKeyPositions`. Projecting inline
+ // here as well would leave that block deriving positions from an already
projected partitioning
+ // and applying them to the unprojected keys.
+ val isCoPartitioned = childrenIndexes.length > 1
// Ensure that the operator's children satisfy their output distribution
requirements.
var children = originalChildren.zip(requiredChildDistributions).map {
case (child, distribution) =>
- // Split child's partitioning into categories
- val (other, grouped, nonGrouped) =
splitKeyedPartitionings(child.outputPartitioning)
+ // Ask what the child's partitioning still needs to satisfy the
distribution
+ val (otherSatisfies, keyed) =
+ splitKeyedPartitionings(child.outputPartitioning, distribution,
isCoPartitioned)
- // If non-KeyedPartitioning already satisfies, no changes needed
- if (other.exists(_.satisfies(distribution))) {
+ // If a non-KeyedPartitioning already satisfies, no changes needed
+ if (otherSatisfies) {
child
} else {
- // Check KeyedPartitioning satisfaction conditions
- val groupedSatisfies = grouped.find(_.satisfies(distribution))
- val nonGroupedSatisfiesAsIs =
nonGrouped.exists(_.nonGroupedSatisfies(distribution))
- val nonGroupedSatisfiesWhenGrouped =
nonGrouped.find(_.groupedSatisfies(distribution))
-
- // Check if any KeyedPartitioning satisfies the distribution
- if (groupedSatisfies.isDefined || nonGroupedSatisfiesAsIs
- || nonGroupedSatisfiesWhenGrouped.isDefined) {
- distribution match {
- case o: OrderedDistribution =>
- // OrderedDistribution requires grouped KeyedPartitioning with
sorted keys
- // according to the distribution's ordering.
- // Find any KeyedPartitioning that satisfies via
groupedSatisfies.
- val satisfyingKeyedPartitioning =
- groupedSatisfies.orElse(nonGroupedSatisfiesWhenGrouped).get
- // 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)
- if
(satisfyingKeyedPartitioning.partitionKeys.sliding(2).forall {
- case Seq(k1, k2) => keyOrdering.lteq(k1, k2)
- }) {
+ keyed match {
+ case Some(resolution) =>
+ distribution 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)
Review Comment:
**[correctness / acknowledged-deferred]** This arm still binds the ordering
to the expressions' declared attribute types and evaluates it over
`partitionKeys` rows (`keyOrdering.lteq` below, and the `sortBy` in the else
branch) — the one key-reading site in this rule not migrated to `keyDataTypes`.
With `v2BucketingAllowSorting` + `allowCompatibleTransforms`, a global `ORDER
BY` over a reduced join reaches this and throws CCE at planning (verified
identical on `master`, and the PR description already defers it).
The deferral rationale is sound — reduced keys live in a different key
space, so a type fix alone would give silently wrong ordering. One thought: an
interim planning-time gate refusing `keyDataTypes != expressionDataTypes` here
would trade part of the crash surface for a shuffle until the follow-up lands,
though it would not cover same-typed reductions.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -867,42 +879,258 @@ case class EnsureRequirements(
}
/**
- * Splits a partitioning into three categories:
- * 1. Non-KeyedPartitioning (HashPartitioning, RangePartitioning, etc.)
- * 2. Grouped KeyedPartitioning (isGrouped = true)
- * 3. Non-grouped KeyedPartitioning (isGrouped = false)
+ * The positions of `kp`'s partition expressions that are operation keys of
`distribution`, and so
+ * have to survive a projection. All of them when nothing needs projecting.
+ *
+ * Under `v2BucketingAllowKeysSubsetOfPartitionKeys` a [[KeyedPartitioning]]
may be grouped on
+ * more keys than the operation requires, in which case partitions sharing
an operation key are
+ * still separate. A partition expression is an operation key in two ways:
one of its *references*
+ * is a cluster key - the form `groupedSatisfies` and
`KeyedShuffleSpec.keyPositions` both use,
+ * where a `bucket(4, a)` transform covers the cluster key `a` - or the
expression *itself* is a
+ * cluster key. The second is never decisive in practice, because
`IdentityTransform` resolves to
+ * the attribute itself and then the reference-level test matches the same
position anyway; it is
+ * kept so that a partition expression which is a cluster key can never be
projected away.
+ *
+ * Returns every position for a co-partitioned operator: there the
multi-child block owns the
+ * projection, and doing it here as well would leave that block deriving
positions from an already
+ * projected partitioning and applying them to the unprojected partition
expressions.
+ *
+ * An empty result means no partition expression covers an operation key, so
there is nothing to
+ * project onto. That is the answer for a member that cannot satisfy
`distribution`, which is why
+ * the caller only asks for members that can. It also happens for a member
that can: one whose
+ * expressions have no references at all makes every `groupedSatisfies`
branch vacuously true, and
+ * nothing at `KeyedPartitioning` construction rejects that. The caller
skips such a member rather
+ * than project it to no position, which would collapse every partition into
one.
+ *
+ * Keeping a position is only sound because `groupedSatisfies`' subset
branch also requires
+ * `expressions.forall(_.references.size == 1)`: a kept expression is then a
function of a single
+ * cluster key, so coalescing on the projected keys cannot put rows that
share an operation key on
+ * different partitions.
+ */
+ private def clusterKeyPositions(
+ kp: KeyedPartitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): BitSet = distribution match {
+ case c: ClusteredDistribution if !isCoPartitioned =>
+ kp.expressions.indices.filter { i =>
+ val e = kp.expressions(i)
+ c.clustering.exists(_.semanticEquals(e)) ||
+ e.references.exists(ref =>
c.clustering.exists(_.semanticEquals(ref)))
+ }.to(BitSet)
+ case _ => kp.expressions.indices.to(BitSet)
+ }
+
+ /**
+ * Splits a partitioning into the two questions the caller acts on, in this
order:
+ * 1. does one of its non-[[KeyedPartitioning]] members (HashPartitioning,
RangePartitioning,
+ * etc.) already satisfy `distribution`, in which case the child needs
nothing
+ * 2. and if not, how can a [[KeyedPartitioning]] member satisfy it: as it
is (`Left`), or after a
+ * [[GroupPartitionsExec]] projecting to the given partition expression
positions (`Right`,
+ * with `None` positions when the node only has to coalesce duplicate
partition keys), or not
+ * at all (`None`)
+ *
+ * The order matters for more than tidiness: the first question touches no
partition key, the
+ * second projects them. And a `Left` is not the same answer as a satisfying
non-keyed member --
+ * the `OrderedDistribution` arm has to look at the keys of the partitioning
it gets.
+ *
+ * At most one `KeyedPartitioning` comes back, because the caller acts on a
single one: whichever
+ * it takes, the child then satisfies the distribution and the rest of its
partitioning is
+ * irrelevant. A partitioning that satisfies the distribution can still come
back as a `Right`,
+ * because `satisfies` over-claims under
`v2BucketingAllowKeysSubsetOfPartitionKeys`.
+ *
+ * That is the point of classifying by what still has to happen to the data
rather than by how the
+ * partitioning was built. An already grouped `KeyedPartitioning` can still
need a
+ * `GroupPartitionsExec`, because
`v2BucketingAllowKeysSubsetOfPartitionKeys` lets it be grouped
+ * on more keys than the operation requires -- `isGrouped` only tells
whether the *full*
+ * partition keys are unique. Keeping both reasons in one answer leaves the
caller a single
+ * `ClusteredDistribution` arm that inserts the node, and one place that
decides the projection.
*
* @param partitioning The partitioning to split
- * @return A tuple of (other, grouped, nonGrouped) where:
- * - other: Option containing non-KeyedPartitioning(s)
- * - grouped: Seq of grouped KeyedPartitionings
- * - nonGrouped: Seq of non-grouped KeyedPartitionings
+ * @param distribution The distribution to satisfy
+ * @param isCoPartitioned Whether the parent operator co-partitions more
than one child, in which
+ * case the projection is not done here (see
`clusterKeyPositions`)
*/
- private def splitKeyedPartitionings(partitioning: Partitioning) = {
+ private def splitKeyedPartitionings(
+ partitioning: Partitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): (Boolean, Option[KeyedResolution]) = {
val otherPartitionings = ArrayBuffer.empty[Partitioning]
- val groupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
- val nonGroupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
+ val keyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
def split(p: Partitioning): Unit = p match {
case c: PartitioningCollection => c.partitionings.foreach(split)
- case k: KeyedPartitioning =>
- if (k.isGrouped) {
- groupedKeyedPartitionings += k
- } else {
- nonGroupedKeyedPartitionings += k
- }
+ case k: KeyedPartitioning => keyedPartitionings += k
case o => otherPartitionings += o
}
split(partitioning)
- val other = otherPartitionings.length match {
- case 0 => None
- case 1 => Some(otherPartitionings.head)
- case _ => Some(PartitioningCollection(otherPartitionings.toSeq))
+ if (otherPartitionings.exists(_.satisfies(distribution))) {
+ (true, None)
+ } else {
+ (false, resolveKeyedPartitioning(keyedPartitionings.toSeq, distribution,
isCoPartitioned))
}
+ }
- (other, groupedKeyedPartitionings.toSeq,
nonGroupedKeyedPartitionings.toSeq)
+ /**
+ * How one of `keyedPartitionings` can satisfy `distribution`, or `None`
when none of them can.
+ * See `splitKeyedPartitionings`, which is the only caller.
+ */
+ private def resolveKeyedPartitioning(
+ keyedPartitionings: Seq[KeyedPartitioning],
+ distribution: Distribution,
+ isCoPartitioned: Boolean): Option[KeyedResolution] = {
+ // A member that needs no node at all settles the whole child, so it is
kept apart from the
+ // candidates that would need one.
+ var satisfiedAsIs: Option[KeyedPartitioning] = None
+ // The candidates that would need a node, keyed by the positions the node
would project them to.
+ // One entry per distinct position set is enough, and the first member
wins: the same set
+ // projects to the same keys whichever member applies it, because
`PartitioningCollection`
+ // guarantees its members share the `partitionKeys` reference and their
arity, so position `i`
+ // addresses the same key column in all of them.
+ //
+ // `GroupPartitionsExec` re-derives the member independently, with a
`collectFirst` over its
+ // child's partitioning, so the member recorded here and the one used at
execution agree only
+ // because of that same guarantee --
`PartitioningCollection.checkKeyedPartitioningInvariant`,
+ // and the value-equality interning in `fromPartitionings` behind it.
Relaxing the invariant
+ // means changing both places together, not just this one. What the
members may still differ in
+ // is their `expressionDataTypes`, which nothing enforces; that does not
reach the keys, because
+ // both sides read them at `KeyedPartitioning.keyDataTypes` instead.
+ //
+ // Insertion-ordered so that when two sets leave the same number of
partitions, the one from the
+ // member the child reports first wins. That tie is the only thing the
order decides, and either
+ // winner satisfies the distribution -- but the two project to different
keys, so the choice is
+ // visible in the plan.
+ val candidates = mutable.LinkedHashMap.empty[BitSet, KeyedPartitioning]
+
+ // The number of partitions a `GroupPartitionsExec` projecting to
`positions` would leave.
+ //
+ // A projection that keeps every position is the identity on the key
values, so its count needs
+ // no projected rows at all -- which is the common shape on the default
config, where nothing
+ // narrows the positions and the node is inserted only to coalesce
duplicate keys. The rest
+ // allocate a row per input partition and hash it with an uncached
`hashCode`, which is the
+ // expensive step here, so the answer is memoized.
+ //
+ // The position set is the whole memo key. `projectKeys` reads each key
value at
+ // `KeyedPartitioning.keyDataTypes`, the types the keys were built with,
and every member of a
+ // child's partitioning shares the same keys, so the same position set
projects to the same
+ // count whichever member is asked. Reading the values at the
*expressions'* types would not
+ // have that property, and would not even be sound: a reducer can rewrite
the keys onto another
+ // key space while a member keeps reporting the expressions it was built
from.
+ val projectedNumPartitions = mutable.Map.empty[BitSet, Int]
Review Comment:
**[efficiency, minor]** The memo keeps only the count and discards the
projected keys, so whenever a count *was* computed (satisfying-narrowing
member, `requiredNumPartitions` filter, or the `maxBy` ranking), the inserted
`GroupPartitionsExec` re-runs the identical `projectKeys` + grouping in the
same planning pass as soon as its `outputPartitioning` is consulted (and
`tryEnableSortedMerge`'s `copy` recomputes once more). The default-config
single-candidate shape is unaffected thanks to the `ranked.size == 1` fast path.
No drop-in fix — `expectedPartitionKeys` has different semantics and a
cached-keys parameter would fight the documented independent re-derivation — so
this is just a noted trade-off / possible follow-up.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -867,42 +879,258 @@ case class EnsureRequirements(
}
/**
- * Splits a partitioning into three categories:
- * 1. Non-KeyedPartitioning (HashPartitioning, RangePartitioning, etc.)
- * 2. Grouped KeyedPartitioning (isGrouped = true)
- * 3. Non-grouped KeyedPartitioning (isGrouped = false)
+ * The positions of `kp`'s partition expressions that are operation keys of
`distribution`, and so
+ * have to survive a projection. All of them when nothing needs projecting.
+ *
+ * Under `v2BucketingAllowKeysSubsetOfPartitionKeys` a [[KeyedPartitioning]]
may be grouped on
+ * more keys than the operation requires, in which case partitions sharing
an operation key are
+ * still separate. A partition expression is an operation key in two ways:
one of its *references*
+ * is a cluster key - the form `groupedSatisfies` and
`KeyedShuffleSpec.keyPositions` both use,
+ * where a `bucket(4, a)` transform covers the cluster key `a` - or the
expression *itself* is a
+ * cluster key. The second is never decisive in practice, because
`IdentityTransform` resolves to
+ * the attribute itself and then the reference-level test matches the same
position anyway; it is
+ * kept so that a partition expression which is a cluster key can never be
projected away.
+ *
+ * Returns every position for a co-partitioned operator: there the
multi-child block owns the
+ * projection, and doing it here as well would leave that block deriving
positions from an already
+ * projected partitioning and applying them to the unprojected partition
expressions.
+ *
+ * An empty result means no partition expression covers an operation key, so
there is nothing to
+ * project onto. That is the answer for a member that cannot satisfy
`distribution`, which is why
+ * the caller only asks for members that can. It also happens for a member
that can: one whose
+ * expressions have no references at all makes every `groupedSatisfies`
branch vacuously true, and
+ * nothing at `KeyedPartitioning` construction rejects that. The caller
skips such a member rather
+ * than project it to no position, which would collapse every partition into
one.
+ *
+ * Keeping a position is only sound because `groupedSatisfies`' subset
branch also requires
+ * `expressions.forall(_.references.size == 1)`: a kept expression is then a
function of a single
+ * cluster key, so coalescing on the projected keys cannot put rows that
share an operation key on
+ * different partitions.
+ */
+ private def clusterKeyPositions(
+ kp: KeyedPartitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): BitSet = distribution match {
+ case c: ClusteredDistribution if !isCoPartitioned =>
+ kp.expressions.indices.filter { i =>
+ val e = kp.expressions(i)
+ c.clustering.exists(_.semanticEquals(e)) ||
+ e.references.exists(ref =>
c.clustering.exists(_.semanticEquals(ref)))
+ }.to(BitSet)
+ case _ => kp.expressions.indices.to(BitSet)
+ }
+
+ /**
+ * Splits a partitioning into the two questions the caller acts on, in this
order:
+ * 1. does one of its non-[[KeyedPartitioning]] members (HashPartitioning,
RangePartitioning,
+ * etc.) already satisfy `distribution`, in which case the child needs
nothing
+ * 2. and if not, how can a [[KeyedPartitioning]] member satisfy it: as it
is (`Left`), or after a
+ * [[GroupPartitionsExec]] projecting to the given partition expression
positions (`Right`,
+ * with `None` positions when the node only has to coalesce duplicate
partition keys), or not
+ * at all (`None`)
+ *
+ * The order matters for more than tidiness: the first question touches no
partition key, the
+ * second projects them. And a `Left` is not the same answer as a satisfying
non-keyed member --
+ * the `OrderedDistribution` arm has to look at the keys of the partitioning
it gets.
+ *
+ * At most one `KeyedPartitioning` comes back, because the caller acts on a
single one: whichever
+ * it takes, the child then satisfies the distribution and the rest of its
partitioning is
+ * irrelevant. A partitioning that satisfies the distribution can still come
back as a `Right`,
+ * because `satisfies` over-claims under
`v2BucketingAllowKeysSubsetOfPartitionKeys`.
+ *
+ * That is the point of classifying by what still has to happen to the data
rather than by how the
+ * partitioning was built. An already grouped `KeyedPartitioning` can still
need a
+ * `GroupPartitionsExec`, because
`v2BucketingAllowKeysSubsetOfPartitionKeys` lets it be grouped
+ * on more keys than the operation requires -- `isGrouped` only tells
whether the *full*
+ * partition keys are unique. Keeping both reasons in one answer leaves the
caller a single
+ * `ClusteredDistribution` arm that inserts the node, and one place that
decides the projection.
*
* @param partitioning The partitioning to split
- * @return A tuple of (other, grouped, nonGrouped) where:
- * - other: Option containing non-KeyedPartitioning(s)
- * - grouped: Seq of grouped KeyedPartitionings
- * - nonGrouped: Seq of non-grouped KeyedPartitionings
+ * @param distribution The distribution to satisfy
+ * @param isCoPartitioned Whether the parent operator co-partitions more
than one child, in which
+ * case the projection is not done here (see
`clusterKeyPositions`)
*/
- private def splitKeyedPartitionings(partitioning: Partitioning) = {
+ private def splitKeyedPartitionings(
+ partitioning: Partitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): (Boolean, Option[KeyedResolution]) = {
val otherPartitionings = ArrayBuffer.empty[Partitioning]
- val groupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
- val nonGroupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
+ val keyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
def split(p: Partitioning): Unit = p match {
case c: PartitioningCollection => c.partitionings.foreach(split)
- case k: KeyedPartitioning =>
- if (k.isGrouped) {
- groupedKeyedPartitionings += k
- } else {
- nonGroupedKeyedPartitionings += k
- }
+ case k: KeyedPartitioning => keyedPartitionings += k
case o => otherPartitionings += o
}
split(partitioning)
- val other = otherPartitionings.length match {
- case 0 => None
- case 1 => Some(otherPartitionings.head)
- case _ => Some(PartitioningCollection(otherPartitionings.toSeq))
+ if (otherPartitionings.exists(_.satisfies(distribution))) {
+ (true, None)
+ } else {
+ (false, resolveKeyedPartitioning(keyedPartitionings.toSeq, distribution,
isCoPartitioned))
}
+ }
- (other, groupedKeyedPartitionings.toSeq,
nonGroupedKeyedPartitionings.toSeq)
+ /**
+ * How one of `keyedPartitionings` can satisfy `distribution`, or `None`
when none of them can.
+ * See `splitKeyedPartitionings`, which is the only caller.
+ */
+ private def resolveKeyedPartitioning(
+ keyedPartitionings: Seq[KeyedPartitioning],
+ distribution: Distribution,
+ isCoPartitioned: Boolean): Option[KeyedResolution] = {
+ // A member that needs no node at all settles the whole child, so it is
kept apart from the
+ // candidates that would need one.
+ var satisfiedAsIs: Option[KeyedPartitioning] = None
+ // The candidates that would need a node, keyed by the positions the node
would project them to.
+ // One entry per distinct position set is enough, and the first member
wins: the same set
+ // projects to the same keys whichever member applies it, because
`PartitioningCollection`
+ // guarantees its members share the `partitionKeys` reference and their
arity, so position `i`
+ // addresses the same key column in all of them.
+ //
+ // `GroupPartitionsExec` re-derives the member independently, with a
`collectFirst` over its
+ // child's partitioning, so the member recorded here and the one used at
execution agree only
+ // because of that same guarantee --
`PartitioningCollection.checkKeyedPartitioningInvariant`,
+ // and the value-equality interning in `fromPartitionings` behind it.
Relaxing the invariant
+ // means changing both places together, not just this one. What the
members may still differ in
+ // is their `expressionDataTypes`, which nothing enforces; that does not
reach the keys, because
+ // both sides read them at `KeyedPartitioning.keyDataTypes` instead.
+ //
+ // Insertion-ordered so that when two sets leave the same number of
partitions, the one from the
+ // member the child reports first wins. That tie is the only thing the
order decides, and either
+ // winner satisfies the distribution -- but the two project to different
keys, so the choice is
+ // visible in the plan.
+ val candidates = mutable.LinkedHashMap.empty[BitSet, KeyedPartitioning]
+
+ // The number of partitions a `GroupPartitionsExec` projecting to
`positions` would leave.
+ //
+ // A projection that keeps every position is the identity on the key
values, so its count needs
+ // no projected rows at all -- which is the common shape on the default
config, where nothing
+ // narrows the positions and the node is inserted only to coalesce
duplicate keys. The rest
+ // allocate a row per input partition and hash it with an uncached
`hashCode`, which is the
+ // expensive step here, so the answer is memoized.
+ //
+ // The position set is the whole memo key. `projectKeys` reads each key
value at
+ // `KeyedPartitioning.keyDataTypes`, the types the keys were built with,
and every member of a
+ // child's partitioning shares the same keys, so the same position set
projects to the same
+ // count whichever member is asked. Reading the values at the
*expressions'* types would not
+ // have that property, and would not even be sound: a reducer can rewrite
the keys onto another
+ // key space while a member keeps reporting the expressions it was built
from.
+ val projectedNumPartitions = mutable.Map.empty[BitSet, Int]
+ def numPartitionsAfter(kp: KeyedPartitioning, positions: BitSet): Int =
+ projectedNumPartitions.getOrElseUpdate(positions, {
+ if (positions.size < kp.expressions.length) {
+ kp.projectKeys(positions.toSeq)._2.distinct.size
+ } else if (kp.isGrouped) {
+ kp.numPartitions
+ } else {
+ kp.partitionKeys.distinct.size
+ }
+ })
+
+ // Which members can satisfy the distribution at all, which of their
partition expression
+ // positions are operation keys, and whether any of them needs no node.
The positions are only
+ // computed for a member that can satisfy -- for one that cannot, no
position would be covered
+ // and an empty set means something else there (see `clusterKeyPositions`).
+ keyedPartitionings.foreach { k =>
+ // Once a member needs no node the child is settled, so the rest are
skipped.
+ if (satisfiedAsIs.isEmpty) {
+ // `satisfies` is the strict question: it also enforces
`requiredNumPartitions`. A
+ // non-grouped partitioning never satisfies a `ClusteredDistribution`
as it is, because
+ // `satisfies0` gates that on `isGrouped`; it still needs a node to
coalesce duplicate
+ // keys.
+ val satisfies = k.satisfies(distribution)
+ if (satisfies || k.groupedSatisfies(distribution)) {
+ val positions = clusterKeyPositions(k, distribution, isCoPartitioned)
+ // With no position covered there is nothing to project onto, so the
member is skipped
+ // rather than projected to no position at all, which would collapse
every partition into
+ // one. Only a partitioning whose expressions have no references
gets here, and only
+ // because nothing rejects one -- see `clusterKeyPositions`.
+ if (positions.nonEmpty || k.expressions.isEmpty) {
+ // A node is pointless when `k` satisfies and nothing is left for
the node to do. That
+ // holds in two ways. Either the projection drops no position, so
`satisfies` is not the
+ // over-claim this fix is about -- also the only shape
`UnspecifiedDistribution` and
+ // `AllTuples` ever reach, the two `nonGroupedSatisfies` covers,
because
+ // `clusterKeyPositions` returns every position for them. Or a
position is dropped but
+ // the projection merges nothing, so every operation key already
lives on a single
+ // partition. Keeping `k` is then better than projecting: both
describe the same number
+ // of partitions, and only `k` still names the dropped keys, which
lets a downstream
+ // operator co-partition on them too.
+ //
+ // `numPartitions` is the count to compare against in the second
case, because there the
+ // distribution is a `ClusteredDistribution` -- so `satisfies`
went through
+ // `isGrouped && groupedSatisfies` -- and a grouped partitioning
has distinct keys,
+ // leaving the node nothing to coalesce.
+ if (satisfies && (positions.size == k.expressions.length ||
+ numPartitionsAfter(k, positions) == k.numPartitions)) {
Review Comment:
**[efficiency, minor]** This eagerly pays an O(#partitions) projection for
an early satisfying-but-narrowing member even when a *later* member of the same
collection turns out to satisfy with full positions and zero key work (the
earlier projection is then discarded). Narrow shape (subset config + asymmetric
coverage across collection members) and planning-time only — worth a two-pass
restructure only if it stays simple against the candidate-ordering invariants
documented above.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -867,42 +879,258 @@ case class EnsureRequirements(
}
/**
- * Splits a partitioning into three categories:
- * 1. Non-KeyedPartitioning (HashPartitioning, RangePartitioning, etc.)
- * 2. Grouped KeyedPartitioning (isGrouped = true)
- * 3. Non-grouped KeyedPartitioning (isGrouped = false)
+ * The positions of `kp`'s partition expressions that are operation keys of
`distribution`, and so
+ * have to survive a projection. All of them when nothing needs projecting.
+ *
+ * Under `v2BucketingAllowKeysSubsetOfPartitionKeys` a [[KeyedPartitioning]]
may be grouped on
+ * more keys than the operation requires, in which case partitions sharing
an operation key are
+ * still separate. A partition expression is an operation key in two ways:
one of its *references*
+ * is a cluster key - the form `groupedSatisfies` and
`KeyedShuffleSpec.keyPositions` both use,
+ * where a `bucket(4, a)` transform covers the cluster key `a` - or the
expression *itself* is a
+ * cluster key. The second is never decisive in practice, because
`IdentityTransform` resolves to
+ * the attribute itself and then the reference-level test matches the same
position anyway; it is
+ * kept so that a partition expression which is a cluster key can never be
projected away.
+ *
+ * Returns every position for a co-partitioned operator: there the
multi-child block owns the
+ * projection, and doing it here as well would leave that block deriving
positions from an already
+ * projected partitioning and applying them to the unprojected partition
expressions.
+ *
+ * An empty result means no partition expression covers an operation key, so
there is nothing to
+ * project onto. That is the answer for a member that cannot satisfy
`distribution`, which is why
+ * the caller only asks for members that can. It also happens for a member
that can: one whose
+ * expressions have no references at all makes every `groupedSatisfies`
branch vacuously true, and
+ * nothing at `KeyedPartitioning` construction rejects that. The caller
skips such a member rather
+ * than project it to no position, which would collapse every partition into
one.
+ *
+ * Keeping a position is only sound because `groupedSatisfies`' subset
branch also requires
+ * `expressions.forall(_.references.size == 1)`: a kept expression is then a
function of a single
+ * cluster key, so coalescing on the projected keys cannot put rows that
share an operation key on
+ * different partitions.
+ */
+ private def clusterKeyPositions(
+ kp: KeyedPartitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): BitSet = distribution match {
+ case c: ClusteredDistribution if !isCoPartitioned =>
+ kp.expressions.indices.filter { i =>
+ val e = kp.expressions(i)
+ c.clustering.exists(_.semanticEquals(e)) ||
+ e.references.exists(ref =>
c.clustering.exists(_.semanticEquals(ref)))
+ }.to(BitSet)
+ case _ => kp.expressions.indices.to(BitSet)
+ }
+
+ /**
+ * Splits a partitioning into the two questions the caller acts on, in this
order:
+ * 1. does one of its non-[[KeyedPartitioning]] members (HashPartitioning,
RangePartitioning,
+ * etc.) already satisfy `distribution`, in which case the child needs
nothing
+ * 2. and if not, how can a [[KeyedPartitioning]] member satisfy it: as it
is (`Left`), or after a
+ * [[GroupPartitionsExec]] projecting to the given partition expression
positions (`Right`,
+ * with `None` positions when the node only has to coalesce duplicate
partition keys), or not
+ * at all (`None`)
+ *
+ * The order matters for more than tidiness: the first question touches no
partition key, the
+ * second projects them. And a `Left` is not the same answer as a satisfying
non-keyed member --
+ * the `OrderedDistribution` arm has to look at the keys of the partitioning
it gets.
+ *
+ * At most one `KeyedPartitioning` comes back, because the caller acts on a
single one: whichever
+ * it takes, the child then satisfies the distribution and the rest of its
partitioning is
+ * irrelevant. A partitioning that satisfies the distribution can still come
back as a `Right`,
+ * because `satisfies` over-claims under
`v2BucketingAllowKeysSubsetOfPartitionKeys`.
+ *
+ * That is the point of classifying by what still has to happen to the data
rather than by how the
+ * partitioning was built. An already grouped `KeyedPartitioning` can still
need a
+ * `GroupPartitionsExec`, because
`v2BucketingAllowKeysSubsetOfPartitionKeys` lets it be grouped
+ * on more keys than the operation requires -- `isGrouped` only tells
whether the *full*
+ * partition keys are unique. Keeping both reasons in one answer leaves the
caller a single
+ * `ClusteredDistribution` arm that inserts the node, and one place that
decides the projection.
*
* @param partitioning The partitioning to split
- * @return A tuple of (other, grouped, nonGrouped) where:
- * - other: Option containing non-KeyedPartitioning(s)
- * - grouped: Seq of grouped KeyedPartitionings
- * - nonGrouped: Seq of non-grouped KeyedPartitionings
+ * @param distribution The distribution to satisfy
+ * @param isCoPartitioned Whether the parent operator co-partitions more
than one child, in which
+ * case the projection is not done here (see
`clusterKeyPositions`)
*/
- private def splitKeyedPartitionings(partitioning: Partitioning) = {
+ private def splitKeyedPartitionings(
+ partitioning: Partitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): (Boolean, Option[KeyedResolution]) = {
val otherPartitionings = ArrayBuffer.empty[Partitioning]
- val groupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
- val nonGroupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
+ val keyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
def split(p: Partitioning): Unit = p match {
case c: PartitioningCollection => c.partitionings.foreach(split)
- case k: KeyedPartitioning =>
- if (k.isGrouped) {
- groupedKeyedPartitionings += k
- } else {
- nonGroupedKeyedPartitionings += k
- }
+ case k: KeyedPartitioning => keyedPartitionings += k
case o => otherPartitionings += o
}
split(partitioning)
- val other = otherPartitionings.length match {
- case 0 => None
- case 1 => Some(otherPartitionings.head)
- case _ => Some(PartitioningCollection(otherPartitionings.toSeq))
+ if (otherPartitionings.exists(_.satisfies(distribution))) {
+ (true, None)
+ } else {
+ (false, resolveKeyedPartitioning(keyedPartitionings.toSeq, distribution,
isCoPartitioned))
}
+ }
- (other, groupedKeyedPartitionings.toSeq,
nonGroupedKeyedPartitionings.toSeq)
+ /**
+ * How one of `keyedPartitionings` can satisfy `distribution`, or `None`
when none of them can.
+ * See `splitKeyedPartitionings`, which is the only caller.
+ */
+ private def resolveKeyedPartitioning(
+ keyedPartitionings: Seq[KeyedPartitioning],
+ distribution: Distribution,
+ isCoPartitioned: Boolean): Option[KeyedResolution] = {
+ // A member that needs no node at all settles the whole child, so it is
kept apart from the
+ // candidates that would need one.
+ var satisfiedAsIs: Option[KeyedPartitioning] = None
+ // The candidates that would need a node, keyed by the positions the node
would project them to.
+ // One entry per distinct position set is enough, and the first member
wins: the same set
+ // projects to the same keys whichever member applies it, because
`PartitioningCollection`
+ // guarantees its members share the `partitionKeys` reference and their
arity, so position `i`
+ // addresses the same key column in all of them.
+ //
+ // `GroupPartitionsExec` re-derives the member independently, with a
`collectFirst` over its
+ // child's partitioning, so the member recorded here and the one used at
execution agree only
+ // because of that same guarantee --
`PartitioningCollection.checkKeyedPartitioningInvariant`,
+ // and the value-equality interning in `fromPartitionings` behind it.
Relaxing the invariant
+ // means changing both places together, not just this one. What the
members may still differ in
+ // is their `expressionDataTypes`, which nothing enforces; that does not
reach the keys, because
+ // both sides read them at `KeyedPartitioning.keyDataTypes` instead.
+ //
+ // Insertion-ordered so that when two sets leave the same number of
partitions, the one from the
+ // member the child reports first wins. That tie is the only thing the
order decides, and either
+ // winner satisfies the distribution -- but the two project to different
keys, so the choice is
+ // visible in the plan.
+ val candidates = mutable.LinkedHashMap.empty[BitSet, KeyedPartitioning]
+
+ // The number of partitions a `GroupPartitionsExec` projecting to
`positions` would leave.
+ //
+ // A projection that keeps every position is the identity on the key
values, so its count needs
+ // no projected rows at all -- which is the common shape on the default
config, where nothing
+ // narrows the positions and the node is inserted only to coalesce
duplicate keys. The rest
+ // allocate a row per input partition and hash it with an uncached
`hashCode`, which is the
+ // expensive step here, so the answer is memoized.
+ //
+ // The position set is the whole memo key. `projectKeys` reads each key
value at
+ // `KeyedPartitioning.keyDataTypes`, the types the keys were built with,
and every member of a
+ // child's partitioning shares the same keys, so the same position set
projects to the same
+ // count whichever member is asked. Reading the values at the
*expressions'* types would not
+ // have that property, and would not even be sound: a reducer can rewrite
the keys onto another
+ // key space while a member keeps reporting the expressions it was built
from.
+ val projectedNumPartitions = mutable.Map.empty[BitSet, Int]
+ def numPartitionsAfter(kp: KeyedPartitioning, positions: BitSet): Int =
+ projectedNumPartitions.getOrElseUpdate(positions, {
+ if (positions.size < kp.expressions.length) {
+ kp.projectKeys(positions.toSeq)._2.distinct.size
+ } else if (kp.isGrouped) {
+ kp.numPartitions
+ } else {
+ kp.partitionKeys.distinct.size
+ }
+ })
+
+ // Which members can satisfy the distribution at all, which of their
partition expression
+ // positions are operation keys, and whether any of them needs no node.
The positions are only
+ // computed for a member that can satisfy -- for one that cannot, no
position would be covered
+ // and an empty set means something else there (see `clusterKeyPositions`).
+ keyedPartitionings.foreach { k =>
+ // Once a member needs no node the child is settled, so the rest are
skipped.
+ if (satisfiedAsIs.isEmpty) {
+ // `satisfies` is the strict question: it also enforces
`requiredNumPartitions`. A
+ // non-grouped partitioning never satisfies a `ClusteredDistribution`
as it is, because
+ // `satisfies0` gates that on `isGrouped`; it still needs a node to
coalesce duplicate
+ // keys.
+ val satisfies = k.satisfies(distribution)
+ if (satisfies || k.groupedSatisfies(distribution)) {
Review Comment:
**[efficiency, nit]** `satisfies0` is `nonGroupedSatisfies || (isGrouped &&
groupedSatisfies)` and `Partitioning.satisfies` has no caching, so for a
grouped member that fails `satisfies`, `groupedSatisfies` runs twice (SQLConf
lookup, `AttributeSet` build, `semanticEquals` scans). Evaluating
`nonGroupedSatisfies`/`groupedSatisfies` once as locals and deriving
`satisfies` from them (plus the count gate) would avoid it. Constant-factor
only.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -867,42 +879,258 @@ case class EnsureRequirements(
}
/**
- * Splits a partitioning into three categories:
- * 1. Non-KeyedPartitioning (HashPartitioning, RangePartitioning, etc.)
- * 2. Grouped KeyedPartitioning (isGrouped = true)
- * 3. Non-grouped KeyedPartitioning (isGrouped = false)
+ * The positions of `kp`'s partition expressions that are operation keys of
`distribution`, and so
+ * have to survive a projection. All of them when nothing needs projecting.
+ *
+ * Under `v2BucketingAllowKeysSubsetOfPartitionKeys` a [[KeyedPartitioning]]
may be grouped on
+ * more keys than the operation requires, in which case partitions sharing
an operation key are
+ * still separate. A partition expression is an operation key in two ways:
one of its *references*
+ * is a cluster key - the form `groupedSatisfies` and
`KeyedShuffleSpec.keyPositions` both use,
+ * where a `bucket(4, a)` transform covers the cluster key `a` - or the
expression *itself* is a
+ * cluster key. The second is never decisive in practice, because
`IdentityTransform` resolves to
+ * the attribute itself and then the reference-level test matches the same
position anyway; it is
+ * kept so that a partition expression which is a cluster key can never be
projected away.
+ *
+ * Returns every position for a co-partitioned operator: there the
multi-child block owns the
+ * projection, and doing it here as well would leave that block deriving
positions from an already
+ * projected partitioning and applying them to the unprojected partition
expressions.
+ *
+ * An empty result means no partition expression covers an operation key, so
there is nothing to
+ * project onto. That is the answer for a member that cannot satisfy
`distribution`, which is why
+ * the caller only asks for members that can. It also happens for a member
that can: one whose
+ * expressions have no references at all makes every `groupedSatisfies`
branch vacuously true, and
+ * nothing at `KeyedPartitioning` construction rejects that. The caller
skips such a member rather
+ * than project it to no position, which would collapse every partition into
one.
+ *
+ * Keeping a position is only sound because `groupedSatisfies`' subset
branch also requires
+ * `expressions.forall(_.references.size == 1)`: a kept expression is then a
function of a single
+ * cluster key, so coalescing on the projected keys cannot put rows that
share an operation key on
+ * different partitions.
+ */
+ private def clusterKeyPositions(
Review Comment:
**[altitude, follow-up]** This is now a third derivation of "operation-key
positions", with deliberately different matching rules from
`KeyedShuffleSpec.keyPositions` and `createShuffleSpec`'s `joinKeyPositions`
(the expression-level `semanticEquals` branch, tolerance of reference-free
expressions) — and the co-partitioned path still derives positions the
`keyPositions`-only way, so the expression-level case the new test pins is
honoured on the single-child path only. Latent today (analyzed queries don't
put a `TransformExpression` in a `ClusteredDistribution`), but consolidating
this next to `keyPositions` on `KeyedPartitioning`/`KeyedShuffleSpec` in the
follow-up would make the divergence visible in one place.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -867,42 +879,258 @@ case class EnsureRequirements(
}
/**
- * Splits a partitioning into three categories:
- * 1. Non-KeyedPartitioning (HashPartitioning, RangePartitioning, etc.)
- * 2. Grouped KeyedPartitioning (isGrouped = true)
- * 3. Non-grouped KeyedPartitioning (isGrouped = false)
+ * The positions of `kp`'s partition expressions that are operation keys of
`distribution`, and so
+ * have to survive a projection. All of them when nothing needs projecting.
+ *
+ * Under `v2BucketingAllowKeysSubsetOfPartitionKeys` a [[KeyedPartitioning]]
may be grouped on
+ * more keys than the operation requires, in which case partitions sharing
an operation key are
+ * still separate. A partition expression is an operation key in two ways:
one of its *references*
+ * is a cluster key - the form `groupedSatisfies` and
`KeyedShuffleSpec.keyPositions` both use,
+ * where a `bucket(4, a)` transform covers the cluster key `a` - or the
expression *itself* is a
+ * cluster key. The second is never decisive in practice, because
`IdentityTransform` resolves to
+ * the attribute itself and then the reference-level test matches the same
position anyway; it is
+ * kept so that a partition expression which is a cluster key can never be
projected away.
+ *
+ * Returns every position for a co-partitioned operator: there the
multi-child block owns the
+ * projection, and doing it here as well would leave that block deriving
positions from an already
+ * projected partitioning and applying them to the unprojected partition
expressions.
+ *
+ * An empty result means no partition expression covers an operation key, so
there is nothing to
+ * project onto. That is the answer for a member that cannot satisfy
`distribution`, which is why
+ * the caller only asks for members that can. It also happens for a member
that can: one whose
+ * expressions have no references at all makes every `groupedSatisfies`
branch vacuously true, and
+ * nothing at `KeyedPartitioning` construction rejects that. The caller
skips such a member rather
+ * than project it to no position, which would collapse every partition into
one.
+ *
+ * Keeping a position is only sound because `groupedSatisfies`' subset
branch also requires
+ * `expressions.forall(_.references.size == 1)`: a kept expression is then a
function of a single
+ * cluster key, so coalescing on the projected keys cannot put rows that
share an operation key on
+ * different partitions.
+ */
+ private def clusterKeyPositions(
+ kp: KeyedPartitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): BitSet = distribution match {
+ case c: ClusteredDistribution if !isCoPartitioned =>
+ kp.expressions.indices.filter { i =>
+ val e = kp.expressions(i)
+ c.clustering.exists(_.semanticEquals(e)) ||
+ e.references.exists(ref =>
c.clustering.exists(_.semanticEquals(ref)))
+ }.to(BitSet)
+ case _ => kp.expressions.indices.to(BitSet)
+ }
+
+ /**
+ * Splits a partitioning into the two questions the caller acts on, in this
order:
+ * 1. does one of its non-[[KeyedPartitioning]] members (HashPartitioning,
RangePartitioning,
+ * etc.) already satisfy `distribution`, in which case the child needs
nothing
+ * 2. and if not, how can a [[KeyedPartitioning]] member satisfy it: as it
is (`Left`), or after a
+ * [[GroupPartitionsExec]] projecting to the given partition expression
positions (`Right`,
+ * with `None` positions when the node only has to coalesce duplicate
partition keys), or not
+ * at all (`None`)
+ *
+ * The order matters for more than tidiness: the first question touches no
partition key, the
+ * second projects them. And a `Left` is not the same answer as a satisfying
non-keyed member --
+ * the `OrderedDistribution` arm has to look at the keys of the partitioning
it gets.
+ *
+ * At most one `KeyedPartitioning` comes back, because the caller acts on a
single one: whichever
+ * it takes, the child then satisfies the distribution and the rest of its
partitioning is
+ * irrelevant. A partitioning that satisfies the distribution can still come
back as a `Right`,
+ * because `satisfies` over-claims under
`v2BucketingAllowKeysSubsetOfPartitionKeys`.
+ *
+ * That is the point of classifying by what still has to happen to the data
rather than by how the
+ * partitioning was built. An already grouped `KeyedPartitioning` can still
need a
+ * `GroupPartitionsExec`, because
`v2BucketingAllowKeysSubsetOfPartitionKeys` lets it be grouped
+ * on more keys than the operation requires -- `isGrouped` only tells
whether the *full*
+ * partition keys are unique. Keeping both reasons in one answer leaves the
caller a single
+ * `ClusteredDistribution` arm that inserts the node, and one place that
decides the projection.
*
* @param partitioning The partitioning to split
- * @return A tuple of (other, grouped, nonGrouped) where:
- * - other: Option containing non-KeyedPartitioning(s)
- * - grouped: Seq of grouped KeyedPartitionings
- * - nonGrouped: Seq of non-grouped KeyedPartitionings
+ * @param distribution The distribution to satisfy
+ * @param isCoPartitioned Whether the parent operator co-partitions more
than one child, in which
+ * case the projection is not done here (see
`clusterKeyPositions`)
*/
- private def splitKeyedPartitionings(partitioning: Partitioning) = {
+ private def splitKeyedPartitionings(
+ partitioning: Partitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): (Boolean, Option[KeyedResolution]) = {
val otherPartitionings = ArrayBuffer.empty[Partitioning]
- val groupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
- val nonGroupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
+ val keyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
def split(p: Partitioning): Unit = p match {
case c: PartitioningCollection => c.partitionings.foreach(split)
- case k: KeyedPartitioning =>
- if (k.isGrouped) {
- groupedKeyedPartitionings += k
- } else {
- nonGroupedKeyedPartitionings += k
- }
+ case k: KeyedPartitioning => keyedPartitionings += k
case o => otherPartitionings += o
}
split(partitioning)
- val other = otherPartitionings.length match {
- case 0 => None
- case 1 => Some(otherPartitionings.head)
- case _ => Some(PartitioningCollection(otherPartitionings.toSeq))
+ if (otherPartitionings.exists(_.satisfies(distribution))) {
+ (true, None)
+ } else {
+ (false, resolveKeyedPartitioning(keyedPartitionings.toSeq, distribution,
isCoPartitioned))
}
+ }
- (other, groupedKeyedPartitionings.toSeq,
nonGroupedKeyedPartitionings.toSeq)
+ /**
+ * How one of `keyedPartitionings` can satisfy `distribution`, or `None`
when none of them can.
+ * See `splitKeyedPartitionings`, which is the only caller.
+ */
+ private def resolveKeyedPartitioning(
+ keyedPartitionings: Seq[KeyedPartitioning],
+ distribution: Distribution,
+ isCoPartitioned: Boolean): Option[KeyedResolution] = {
+ // A member that needs no node at all settles the whole child, so it is
kept apart from the
+ // candidates that would need one.
+ var satisfiedAsIs: Option[KeyedPartitioning] = None
+ // The candidates that would need a node, keyed by the positions the node
would project them to.
+ // One entry per distinct position set is enough, and the first member
wins: the same set
+ // projects to the same keys whichever member applies it, because
`PartitioningCollection`
+ // guarantees its members share the `partitionKeys` reference and their
arity, so position `i`
+ // addresses the same key column in all of them.
+ //
+ // `GroupPartitionsExec` re-derives the member independently, with a
`collectFirst` over its
+ // child's partitioning, so the member recorded here and the one used at
execution agree only
+ // because of that same guarantee --
`PartitioningCollection.checkKeyedPartitioningInvariant`,
+ // and the value-equality interning in `fromPartitionings` behind it.
Relaxing the invariant
+ // means changing both places together, not just this one. What the
members may still differ in
+ // is their `expressionDataTypes`, which nothing enforces; that does not
reach the keys, because
+ // both sides read them at `KeyedPartitioning.keyDataTypes` instead.
+ //
+ // Insertion-ordered so that when two sets leave the same number of
partitions, the one from the
+ // member the child reports first wins. That tie is the only thing the
order decides, and either
+ // winner satisfies the distribution -- but the two project to different
keys, so the choice is
+ // visible in the plan.
+ val candidates = mutable.LinkedHashMap.empty[BitSet, KeyedPartitioning]
+
+ // The number of partitions a `GroupPartitionsExec` projecting to
`positions` would leave.
+ //
+ // A projection that keeps every position is the identity on the key
values, so its count needs
+ // no projected rows at all -- which is the common shape on the default
config, where nothing
+ // narrows the positions and the node is inserted only to coalesce
duplicate keys. The rest
+ // allocate a row per input partition and hash it with an uncached
`hashCode`, which is the
+ // expensive step here, so the answer is memoized.
+ //
+ // The position set is the whole memo key. `projectKeys` reads each key
value at
+ // `KeyedPartitioning.keyDataTypes`, the types the keys were built with,
and every member of a
+ // child's partitioning shares the same keys, so the same position set
projects to the same
+ // count whichever member is asked. Reading the values at the
*expressions'* types would not
+ // have that property, and would not even be sound: a reducer can rewrite
the keys onto another
+ // key space while a member keeps reporting the expressions it was built
from.
+ val projectedNumPartitions = mutable.Map.empty[BitSet, Int]
+ def numPartitionsAfter(kp: KeyedPartitioning, positions: BitSet): Int =
+ projectedNumPartitions.getOrElseUpdate(positions, {
+ if (positions.size < kp.expressions.length) {
+ kp.projectKeys(positions.toSeq)._2.distinct.size
+ } else if (kp.isGrouped) {
+ kp.numPartitions
+ } else {
+ kp.partitionKeys.distinct.size
+ }
+ })
+
+ // Which members can satisfy the distribution at all, which of their
partition expression
+ // positions are operation keys, and whether any of them needs no node.
The positions are only
+ // computed for a member that can satisfy -- for one that cannot, no
position would be covered
+ // and an empty set means something else there (see `clusterKeyPositions`).
+ keyedPartitionings.foreach { k =>
+ // Once a member needs no node the child is settled, so the rest are
skipped.
+ if (satisfiedAsIs.isEmpty) {
+ // `satisfies` is the strict question: it also enforces
`requiredNumPartitions`. A
+ // non-grouped partitioning never satisfies a `ClusteredDistribution`
as it is, because
+ // `satisfies0` gates that on `isGrouped`; it still needs a node to
coalesce duplicate
+ // keys.
+ val satisfies = k.satisfies(distribution)
+ if (satisfies || k.groupedSatisfies(distribution)) {
+ val positions = clusterKeyPositions(k, distribution, isCoPartitioned)
+ // With no position covered there is nothing to project onto, so the
member is skipped
+ // rather than projected to no position at all, which would collapse
every partition into
+ // one. Only a partitioning whose expressions have no references
gets here, and only
+ // because nothing rejects one -- see `clusterKeyPositions`.
+ if (positions.nonEmpty || k.expressions.isEmpty) {
Review Comment:
**[simplification, nit]** The `|| k.expressions.isEmpty` disjunct keeps
alive a zero-expression `KeyedPartitioning` that no in-tree producer can
construct (the scan, `AliasAwareOutputExpression`, and `GroupPartitionsExec`
all guarantee at least one expression), and the branch is untested — even the
reference-free test uses `Seq(Literal(1))` and exercises the *skip* path.
Dropping it would let a zero-expression member shuffle uniformly like the
reference-free case and simplify the trickiest guard in this function.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -867,42 +879,258 @@ case class EnsureRequirements(
}
/**
- * Splits a partitioning into three categories:
- * 1. Non-KeyedPartitioning (HashPartitioning, RangePartitioning, etc.)
- * 2. Grouped KeyedPartitioning (isGrouped = true)
- * 3. Non-grouped KeyedPartitioning (isGrouped = false)
+ * The positions of `kp`'s partition expressions that are operation keys of
`distribution`, and so
+ * have to survive a projection. All of them when nothing needs projecting.
+ *
+ * Under `v2BucketingAllowKeysSubsetOfPartitionKeys` a [[KeyedPartitioning]]
may be grouped on
+ * more keys than the operation requires, in which case partitions sharing
an operation key are
+ * still separate. A partition expression is an operation key in two ways:
one of its *references*
+ * is a cluster key - the form `groupedSatisfies` and
`KeyedShuffleSpec.keyPositions` both use,
+ * where a `bucket(4, a)` transform covers the cluster key `a` - or the
expression *itself* is a
+ * cluster key. The second is never decisive in practice, because
`IdentityTransform` resolves to
+ * the attribute itself and then the reference-level test matches the same
position anyway; it is
+ * kept so that a partition expression which is a cluster key can never be
projected away.
+ *
+ * Returns every position for a co-partitioned operator: there the
multi-child block owns the
+ * projection, and doing it here as well would leave that block deriving
positions from an already
+ * projected partitioning and applying them to the unprojected partition
expressions.
+ *
+ * An empty result means no partition expression covers an operation key, so
there is nothing to
+ * project onto. That is the answer for a member that cannot satisfy
`distribution`, which is why
+ * the caller only asks for members that can. It also happens for a member
that can: one whose
+ * expressions have no references at all makes every `groupedSatisfies`
branch vacuously true, and
+ * nothing at `KeyedPartitioning` construction rejects that. The caller
skips such a member rather
+ * than project it to no position, which would collapse every partition into
one.
+ *
+ * Keeping a position is only sound because `groupedSatisfies`' subset
branch also requires
+ * `expressions.forall(_.references.size == 1)`: a kept expression is then a
function of a single
+ * cluster key, so coalescing on the projected keys cannot put rows that
share an operation key on
+ * different partitions.
+ */
+ private def clusterKeyPositions(
+ kp: KeyedPartitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): BitSet = distribution match {
+ case c: ClusteredDistribution if !isCoPartitioned =>
+ kp.expressions.indices.filter { i =>
+ val e = kp.expressions(i)
+ c.clustering.exists(_.semanticEquals(e)) ||
+ e.references.exists(ref =>
c.clustering.exists(_.semanticEquals(ref)))
+ }.to(BitSet)
+ case _ => kp.expressions.indices.to(BitSet)
+ }
+
+ /**
+ * Splits a partitioning into the two questions the caller acts on, in this
order:
+ * 1. does one of its non-[[KeyedPartitioning]] members (HashPartitioning,
RangePartitioning,
+ * etc.) already satisfy `distribution`, in which case the child needs
nothing
+ * 2. and if not, how can a [[KeyedPartitioning]] member satisfy it: as it
is (`Left`), or after a
+ * [[GroupPartitionsExec]] projecting to the given partition expression
positions (`Right`,
+ * with `None` positions when the node only has to coalesce duplicate
partition keys), or not
+ * at all (`None`)
+ *
+ * The order matters for more than tidiness: the first question touches no
partition key, the
+ * second projects them. And a `Left` is not the same answer as a satisfying
non-keyed member --
+ * the `OrderedDistribution` arm has to look at the keys of the partitioning
it gets.
+ *
+ * At most one `KeyedPartitioning` comes back, because the caller acts on a
single one: whichever
+ * it takes, the child then satisfies the distribution and the rest of its
partitioning is
+ * irrelevant. A partitioning that satisfies the distribution can still come
back as a `Right`,
+ * because `satisfies` over-claims under
`v2BucketingAllowKeysSubsetOfPartitionKeys`.
+ *
+ * That is the point of classifying by what still has to happen to the data
rather than by how the
+ * partitioning was built. An already grouped `KeyedPartitioning` can still
need a
+ * `GroupPartitionsExec`, because
`v2BucketingAllowKeysSubsetOfPartitionKeys` lets it be grouped
+ * on more keys than the operation requires -- `isGrouped` only tells
whether the *full*
+ * partition keys are unique. Keeping both reasons in one answer leaves the
caller a single
+ * `ClusteredDistribution` arm that inserts the node, and one place that
decides the projection.
*
* @param partitioning The partitioning to split
- * @return A tuple of (other, grouped, nonGrouped) where:
- * - other: Option containing non-KeyedPartitioning(s)
- * - grouped: Seq of grouped KeyedPartitionings
- * - nonGrouped: Seq of non-grouped KeyedPartitionings
+ * @param distribution The distribution to satisfy
+ * @param isCoPartitioned Whether the parent operator co-partitions more
than one child, in which
+ * case the projection is not done here (see
`clusterKeyPositions`)
*/
- private def splitKeyedPartitionings(partitioning: Partitioning) = {
+ private def splitKeyedPartitionings(
+ partitioning: Partitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): (Boolean, Option[KeyedResolution]) = {
val otherPartitionings = ArrayBuffer.empty[Partitioning]
- val groupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
- val nonGroupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
+ val keyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
def split(p: Partitioning): Unit = p match {
Review Comment:
**[reuse, nit]** This hand-rolled recursive traversal duplicates the
existing `PartitioningCollection.flatten`; `flatten(partitioning)` followed by
a partition on `KeyedPartitioning` yields the same two sequences without local
recursion or mutable buffers. The recursion predates this PR, but since the
function is rewritten wholesale anyway, the cleanup is in scope.
--
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]
