You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "peter-toth (via GitHub)" <gi...@apache.org> on 2026/09/01 10:59:35 UTC
Re: [PR] [SPARK-58968][SQL] Fix SPJ allowKeysSubsetOfPartitionKeys correctness for non-join operators [spark]
peter-toth commented on code in PR #58262:
URL: https://github.com/apache/spark/pull/58262#discussion_r3903414328
##########
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:
Taken, and it shipped separately as #58420 (SPARK-59120), now on master, 4.2, 4.3 and 4.x. `KeyedShuffleSpec.canCreatePartitioning` ends with `partitioning.expressionsDescribeKeyShape`, which is the clause you asked for. It compares shapes rather than plain types, because `createPartitioning` puts the other child's expressions over these keys, so a struct field name can legitimately differ with no reducer involved.
The whole `keyDataTypes` half left this PR with it, so the comment you flagged is not in this diff any more.
What the clause still does not catch is recorded next to it. Matching shapes are only a proxy: `bucket(12)` and `bucket(8)` reducing onto `bucket(4)` keep the type, pass the gate, and misroute rows anyway. SPARK-59121 replaces the proxy with the real test.
##########
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:
Still deferred, and now tracked rather than only mentioned in a description. The arm is one of three readers that reduced keys break in a way no choice of types can fix, and the `keyDataTypes` scaladoc names all three: this arm, the `UnionExec` key merge, and `KeyedShuffleSpec.reducers`. SPARK-59121 covers them together.
I did not take the interim gate. It would refuse the differently-typed reductions and let the same-typed ones through, so SPARK-59121 has to solve that class anyway and the gate would come straight back out. The crash is identical on `master` and needs `v2BucketingAllowSorting`, which is off by default.
##########
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:
Agreed, and left as it is, for the reason you gave. The comment at the memo now says what it keys on and why, so the trade is on the record rather than implied.
The default shape does not pay it. A single surviving candidate takes the `maximal.size == 1` fast path and no count is computed at all.
##########
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:
Fixed in https://github.com/apache/spark/pull/58262/commits/3509a8873e2ef6e68e4edbde1ace9b494a3e7abe.
`satisfiedAsIs` is two passes now instead of one loop over a `var`. The first `find` asks every admitted member whether its positions cover all of its expressions, which touches no partition key. Only the `orElse` pass computes a count. So a narrowing member no longer pays a projection that a later full-coverage member discards, and the scaladoc at the two `find`s says that is why the order is what it is.
##########
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:
Half of it, by construction. `satisfies` is asked first and `||` short-circuits, so every member that is admitted matches once. A member that fails `satisfies` still matches twice, which is the case you named.
I left that. Composing the two out of locals means putting `satisfies0`'s structure in the caller, and `nonGroupedSatisfies` and `keysSatisfy` are both private now, which is what `keysMaySatisfy` exists to keep. Widening them to save a constant factor at planning time looks like the wrong trade. The comment above `admitted` records why the strict question is asked first.
##########
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:
Kept, and there is a test for it now: `a KeyedPartitioning with no partition expressions is kept as it is`.
Dropping it does not make such a member shuffle uniformly like the reference-free one. It reaches the shuffle branch under `UnspecifiedDistribution`, whose `createPartitioning` throws outright, so the guard is what stops a planning failure rather than a bad plan. The test's comment says that, and it also says what you did, that no in-tree producer builds one.
##########
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:
Fixed in https://github.com/apache/spark/pull/58262/commits/3509a8873e2ef6e68e4edbde1ace9b494a3e7abe.
`splitKeyedPartitionings` starts with `PartitioningCollection.flatten(partitioning)` and then splits by type. No local recursion, no mutable buffers.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org