cloud-fan commented on code in PR #58942:
URL: https://github.com/apache/spark/pull/58942#discussion_r4072240006
##########
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:
**Non-blocking (P2):** This waiver is broader than its producer invariant:
`EnsureRequirements.checkKeyGroupCompatible` builds spread/replicate layouts
only for `SortMergeJoinExec` and `ShuffledHashJoinExec`, while
`SortMergeAsOfJoinExec` also matches `ShuffledJoin`. With partial clustering
enabled, an AQE rule candidate over aligned repeated AS-OF children can
therefore pass without either side being known to repeat the whole key group,
and the partition-local AS-OF scan can miss or duplicate matches. Please
restrict the waiver to the producer-supported join kinds or carry equivalent
explicit provenance.
See **Shared repair plan 1** in the review body.
##########
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:
**Blocking (P1):** For a non-`ShuffledJoin`, this branch now returns `true`
after per-child satisfaction without the merge-target validator's mutual
shuffle-spec compatibility check. `CoGroupExec` and
`FlatMapCoGroupsInBatchExec` zip corresponding partitions, so individually
satisfying children with different counts can fail at execution, while
equal-count but misaligned layouts can emit partial or missing groups. Please
restore generic compatibility validation for multi-child clustered consumers
and keep only the ungrouped waiver join-specific.
**Recommended change:** Restore a generic mutual compatibility check for
every multi-child operator whose requirements are clustered, while retaining
the producer-specific ungrouped exception only on the eligible shuffled-join
path; add grouped cogroup coverage for incompatible counts and layouts.
**Why this works:** After individual distribution and ordering validation,
compare compatible shuffle specs for all multi-child clustered consumers. Use
the specialized as-held pairing only where an eligible join needs the ungrouped
waiver; non-join consumers must still prove a mutually compatible grouped
layout.
**Scope:** Reestablish cross-child compatibility as a generic validator
responsibility without broadening the ungrouped join waiver.
**Compatibility:** The join-specific partially clustered waiver remains
separate from ordinary grouped compatibility validation.
**Risks:** A generic check must not apply the join-only ungrouped semantics
to operators that emit one result per key group.
**Constraints:** Single-child operators continue to owe only their own
distribution and ordering. Valid grouped multi-child plans remain accepted.
**Success:** Two individually grouped cogroup children with different
partition counts are rejected before execution. Equal-count but mutually
misaligned grouped cogroup layouts are rejected. Compatible grouped non-join
children and valid shuffled joins continue to validate.
##########
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:
**Non-blocking (P2):** With partial clustering enabled, this fallback asks
`keysSatisfy`, which does not apply the `isCollapsed` gate.
`EnsureRequirements` asks `keysMaySatisfy`; for an ungrouped collapsed layout,
`mayGroupToSatisfy` rejects it while `allowKeysSubsetOfPartitionKeys` is off.
Two matching collapsed sides can therefore validate even though the planner
refuses that shape. The added negative leaves partial clustering disabled and
exits before this branch, so please preserve the planner's collapsed-layout
permission here and cover the two flags together.
See **Shared repair plan 1** in the review body.
##########
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
Review Comment:
**Nit (P3):** Disabling `requireAllClusterKeysForCoPartition` skips
`allClusterKeysCovered`, so it can admit partitioning on `[a]` for an operation
on `[a, b]`: the partitioning keys cover only a subset of the operation keys. A
partition-key superset is the separate projection case controlled by
`allowKeysSubsetOfPartitionKeys`. Please reverse this example/rationale so it
describes the configuration used here.
##########
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:
**Nit (P3):** This states the regression as present behavior, but the
assertion below requires the unrelated `AQEShuffleReadExec` to have a coalesced
partition. Please make this historical (for example, the join *used to* keep
unrelated shuffles uncoalesced) so the explanation agrees with the behavior
this test pins.
--
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]
