peter-toth commented on code in PR #58262:
URL: https://github.com/apache/spark/pull/58262#discussion_r3872256694
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -867,42 +867,234 @@ 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.
+ *
+ * Only call this for a `kp` that can satisfy `distribution`. For a `kp`
that cannot, no position
+ * would be covered, and an empty result means something else here.
+ *
+ * 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 =>
+ val positions = 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)
+ // A member that can satisfy covers at least one position, unless it
declares no partition
+ // expression at all. `supportsExpressions` rules out a reference-free
expression, so a
+ // non-empty expression list has non-empty `references`, and every
`groupedSatisfies` branch
+ // then needs a cluster key among them. A connector can declare
`KeyGroupedPartitioning` with
+ // no keys though, and nothing rejects it: there the empty result is
harmless, because it
+ // makes the projection a no-op and the node only coalesces, exactly as
before this change.
+ assert(positions.nonEmpty || kp.expressions.isEmpty,
+ s"no partition expression of ${kp.expressions} covers an operation
key")
+ positions
+ case _ => kp.expressions.indices.to(BitSet)
+ }
+
+ /**
+ * Splits a partitioning into three categories with respect to
`distribution`:
+ * 1. Every non-KeyedPartitioning (HashPartitioning, RangePartitioning, etc.)
+ * 2. The KeyedPartitioning that needs no [[GroupPartitionsExec]] at all: it
satisfies the
+ * distribution and such a node would leave nothing changed
+ * 3. Or the one that satisfies it only after a [[GroupPartitionsExec]],
paired with the partition
+ * expression positions that node has to project to (`None` when it only
has to coalesce
+ * duplicate partition keys)
+ *
+ * Categories 2 and 3 hold at most one partitioning each and never both,
because the caller acts
+ * on a single one of them: whichever it takes, the child then satisfies the
distribution and the
+ * rest of the child's partitioning is irrelevant. A partitioning that
satisfies the distribution
+ * can still land in 3, because `satisfies` over-claims under
+ * `v2BucketingAllowKeysSubsetOfPartitionKeys`.
+ *
+ * Categories 2 and 3 classify by what still has to happen to the data, not
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 category leaves
the caller a single
+ * `ClusteredDistribution` arm that inserts the node, and one place that
decides the projection.
+ *
+ * KeyedPartitionings that cannot satisfy the distribution at all are
dropped.
*
* @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`)
+ * @return A tuple of (other, satisfying, needsGrouping)
*/
- private def splitKeyedPartitionings(partitioning: Partitioning) = {
+ private def splitKeyedPartitionings(
+ partitioning: Partitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean) = {
val otherPartitionings = ArrayBuffer.empty[Partitioning]
- val groupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
- val nonGroupedKeyedPartitionings = ArrayBuffer.empty[KeyedPartitioning]
+ // 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.
+ //
+ // 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
Review Comment:
Agreed, and thanks for tracing it across the files. I have added a sentence
at the `candidates` declaration naming `checkKeyedPartitioningInvariant` and
the interning in `fromPartitionings` as the actual guarantee, and saying that
relaxing it means changing `GroupPartitionsExec`'s `collectFirst` at the same
time rather than this side alone.
Worth adding that the assumption is already only partly guaranteed: the
`partitionKeys` reference and the arity are enforced, the per-position
`expressionDataTypes` are not. Two members of one collection really can declare
different types over the same keys - `pushPartValues` plus
`allowCompatibleTransforms`, an `identity(ts)`-partitioned table joined to a
`years(ts)`-partitioned one - which is why the projected-count memo in this
method is keyed on `(BitSet, Seq[DataType])` rather than on the position set
alone.
I would rather not add a `require` on the reference identity here:
`PartitioningCollection` already asserts it on construction, so a second check
would go stale the day that one moves.
--
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]
