ulysses-you commented on code in PR #58942:
URL: https://github.com/apache/spark/pull/58942#discussion_r4081500867
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala:
##########
@@ -45,29 +46,73 @@ object ValidateRequirements extends Logging {
assert(requiredChildDistributions.length == children.length)
assert(requiredChildOrderings.length == children.length)
+ // A `ClusteredDistribution` is the one distribution an operator can owe
its children together
+ // rather than one by one, so a join with two of them is judged by their
pairing below, and
+ // every other child, an operator with a single clustered child included,
answers for itself.
+ // Only a shuffling join gets that reading: it matches within a partition
and unions the
+ // matches, so an alignment that spreads one side against a repeating
other is sound, while an
+ // operator that emits a per-key result from one partition, a cogroup for
instance, would emit a
+ // partial one for every spread key.
`EnsureRequirements.checkKeyGroupCompatible` is the join
+ // path that plans such a pair, and the only producer of one.
+ val coPartitioning = children.length > 1 &&
plan.isInstanceOf[ShuffledJoin] &&
Review Comment:
Thanks @cloud-fan! Restricted in f95e44a9f42, and not by a second copy of
the list: the operator kinds now live in one predicate,
`ShuffledJoin.partiallyClusteredJoinType`, which
`EnsureRequirements.checkKeyGroupCompatible` dispatches on and
`ValidateRequirements` reads through `.isDefined`, so the two cannot drift
apart. A `SortMergeAsOfJoinExec` is excluded with them, and the pair a
shuffled-hash join accepts is pinned as refused for it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala:
##########
@@ -45,29 +46,73 @@ object ValidateRequirements extends Logging {
assert(requiredChildDistributions.length == children.length)
assert(requiredChildOrderings.length == children.length)
+ // A `ClusteredDistribution` is the one distribution an operator can owe
its children together
+ // rather than one by one, so a join with two of them is judged by their
pairing below, and
+ // every other child, an operator with a single clustered child included,
answers for itself.
+ // Only a shuffling join gets that reading: it matches within a partition
and unions the
+ // matches, so an alignment that spreads one side against a repeating
other is sound, while an
+ // operator that emits a per-key result from one partition, a cogroup for
instance, would emit a
+ // partial one for every spread key.
`EnsureRequirements.checkKeyGroupCompatible` is the join
+ // path that plans such a pair, and the only producer of one.
+ val coPartitioning = children.length > 1 &&
plan.isInstanceOf[ShuffledJoin] &&
+ requiredChildDistributions.forall(_.isInstanceOf[ClusteredDistribution])
+
val satisfied =
children.zip(requiredChildDistributions.zip(requiredChildOrderings)).forall {
case (child, (distribution, ordering))
- if !child.outputPartitioning.satisfies(distribution)
+ if (!child.outputPartitioning.satisfies(distribution) &&
+ !(coPartitioning &&
+
PartitioningCollection.representativeOf(child.outputPartitioning).exists(k =>
+ k.isGrouped || PartitioningCollection.mayUngroupedMember)))
|| !SortOrder.orderingSatisfies(child.outputOrdering, ordering) =>
logDebug(s"ValidateRequirements failed: $distribution,
$ordering\n$plan")
false
case _ => true
}
- if (satisfied && children.length > 1 &&
-
requiredChildDistributions.forall(_.isInstanceOf[ClusteredDistribution])) {
- // Check the co-partitioning requirement.
- val specs =
children.map(_.outputPartitioning).zip(requiredChildDistributions).map {
- case (p, d) =>
p.createShuffleSpec(d.asInstanceOf[ClusteredDistribution])
- }
- if (specs.tail.forall(_.isCompatibleWith(specs.head))) {
- true
- } else {
+ // What a co-partitioning operator reads is the pairing: a pair aligned
without grouping, which
+ // partially clustered distribution builds on purpose, is one the sides
agree on while neither
+ // is grouped. The pairing cannot tell how the two sides hold a key's
rows, since a spread side
+ // and one that repeats the whole group report the same keys as two sides
that split the key,
+ // so that rests on the producer: only a join is admitted above, and
+ // `checkKeyGroupCompatible` is where it plans such a pair.
+ if (!satisfied) {
+ false
+ } else if (coPartitioning) {
+ val paired = satisfiesForPairing(children, requiredChildDistributions)
+ if (!paired) {
logDebug(s"ValidateRequirements failed: children not co-partitioned
in\n$plan")
- false
}
+ paired
} else {
- satisfied
+ true
Review Comment:
Thanks @cloud-fan! Restored in f95e44a9f42. The pairing is asked of every
multi-child clustered operator again, which is what the base did, and only the
ungrouped waiver stays producer-scoped: the waiver's condition is the
multi-child cluster, the partially clustered configuration and the producer's
kinds, computed once and passed into the pairing. The cogroup test now covers
the mutual check from three sides, two sides holding the same grouped layout
passing, a differing partition count refused, and a differing key order refused.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -9228,6 +9229,100 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-59671: a partially clustered join leaves AQE's shuffle
coalescing alone") {
+ // AQE validates a stage's whole candidate plan before accepting a
shuffle-read change, so a
Review Comment:
Thanks @cloud-fan! Reworded in f95e44a9f42: the failure is stated against
the base now, "on the base a storage-partitioned join whose sides are aligned
but not grouped kept every shuffle in its stage uncoalesced", and the closing
sentence says the assertions below pin the read the join's presence must leave
alone.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1437,6 +1445,54 @@ object PartitioningCollection {
case other => other.satisfies(required)
}
+ /**
+ * Whether a finished plan may report an ungrouped keyed member, which only
partially clustered
+ * distribution builds. Read by `specsForPairing` and by
`ValidateRequirements`' per-side
+ * admission, so the two cannot drift.
+ */
+ private[sql] def mayUngroupedMember: Boolean =
+ SQLConf.get.v2BucketingPartiallyClusteredDistributionEnabled
+
+ /**
+ * The specs `p` offers for `distribution`, one per member that may serve
it, each of them the
+ * layout that member reports. A keyed member is admitted on `satisfies`,
the as-it-stands
+ * question, count included; a member that is not keyed is asked on
`satisfies` as well and has no
+ * projection to make.
+ *
+ * The one member whose layout a finished plan may report without satisfying
the distribution is
+ * an ungrouped keyed one: that is the shape partially clustered
distribution spreads, and
+ * `mayUngroupedMember` confines it to the configuration that builds one.
`keysSatisfy` is the
+ * as-it-stands question for it, since `satisfies` adds `isGrouped` on top.
+ *
+ * This is the planner's admission of a member
(`EnsureRequirements.createKeyedShuffleSpecs`) less
+ * the coverage of every operation key it requires there
+ * (`spark.sql.requireAllClusterKeysForCoPartition`), which is a skew
heuristic: a member whose
+ * partitioning keys are a superset of the operation's keys is a sound
pairing.
+ *
+ * Two things the planner does are deliberately not done here, because they
build a layout the
+ * plan does not hold, which is the one thing the validator must not read:
+ *
+ * - no projection onto the operation keys
(`KeyedPartitioning.createShuffleSpec` makes one under
+ * `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys`). A
plan in which the
+ * permission applies holds the projection already, in the grouping node
+ * `EnsureRequirements` inserted, so the member reports it and no
projection is wanted.
+ * - no grouping of an ungrouped member (`tryCreate` builds from
`toGrouped`), and no re-sorting
+ * of a grouped one: `toGrouped` dedups and sorts, both of which move the
member's key rows.
+ */
+ private[sql] def specsForPairing(
+ p: Partitioning,
+ distribution: ClusteredDistribution): Seq[ShuffleSpec] =
+ flatten(p).flatMap {
+ case k: KeyedPartitioning =>
+ val pairsAsIs = k.satisfies(distribution) ||
+ (mayUngroupedMember &&
Review Comment:
Thanks @cloud-fan! The permission is back in f95e44a9f42, read once and
shared: `collapsedLayoutMayBeGrouped` is what `mayGroupToSatisfy` and the
ungrouped admission both ask, so a member a finished plan reports ungrouped is
admitted only where the planner would agree to group it. The negative covers
the two flags together, partially clustered distribution on with the subset
permission off, and a unit test pins that the member is not offered in that
configuration.
--
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]