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]

Reply via email to