[SPARK-59671][SQL] Validate a co-partitioned pair by its children's pairing

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala

@@ -564,8 +564,8 @@ case class KeyLayout(
  *
  * == Distribution Satisfaction and Grouping ==
  * Besides the default `satisfies()`, `KeyedPartitioning` answers a family of questions. They differ
- * in what they let happen to the data before the distribution counts as met. Only
- * `keysMaySatisfy()` is asked from outside the class. The rest build it and `satisfies()` up.
+ * in what they let happen to the data before the distribution counts as met. `keysMaySatisfy()` and
+ * `mayServeUngrouped()` are asked from outside the class. The rest build them and `satisfies()` up.
  *
  * - `keysSatisfy()`: do the keys as they stand co-locate every cluster key, with nothing left for
  *   a node to project away? This is the strict question, and `satisfies()` is it plus `isGrouped`.
@@ -578,6 +578,9 @@ case class KeyLayout(
  *   distribution? It is asked of non-grouped partitionings only, where such a node coalesces the
  *   duplicate keys. So the answer is `keysCanSatisfy()`, plus whether that coalescing is allowed.
  *   "Key Collapse" below says when it is not.
+ * - `mayServeUngrouped()`: the same question for a partitioning a finished plan reports ungrouped,
+ *   where nobody is left to insert that node. So the answer is `keysSatisfy()`, plus the same
+ *   coalescing permission, and the projection half of `keysCanSatisfy()` is out.
  * - `keysMaySatisfy()`: the same question for a partitioning that may already be grouped, which is
  *   what `EnsureRequirements` asks. It is `mayGroupToSatisfy()` for a non-grouped one, and
  *   `keysCanSatisfy()` for a grouped one, which has no duplicate keys left for the node to
@@ -937,6 +940,10 @@ case class KeyedPartitioning(
    * side whose keys run the wrong way still answers `true` here. `EnsureRequirements.resolveChild`
    * is what compares the keys against the required ordering and builds the sorting node, and it has
    * to stay: nothing below tells it the keys are already in order.
+   *
+   * `mayServeUngrouped` is the caller that asks this outside the class: a member a finished plan
+   * reports ungrouped has nothing left to group it, so what it holds is what the question is about.
+   * See there for the permission it adds on top.
    */
   private def keysSatisfy(required: Distribution): Boolean = {
     required match {
@@ -990,14 +997,30 @@ case class KeyedPartitioning(
     satisfiesAfterProjection || keysSatisfy(required)
   }
 
+  /**
+   * The `Key Collapse` permission: whether a collapsed layout may be grouped at all. Read wherever
+   * a grouping node or its absence has to be judged. See the class doc's `Key Collapse` section for
+   * what the flag records.
+   */
+  private def collapsedLayoutMayBeGrouped: Boolean =
+    !isCollapsed || SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys
+
+  /**
+   * Whether this partitioning, reported ungrouped, may still serve `required` as it stands: the
+   * keys it holds co-locate every cluster key, and the collapse a grouping node would settle is
+   * permitted. `mayGroupToSatisfy` is the planner's question, which also admits a member such a
+   * node would serve; a finished plan has nobody left to insert one, and this is what it leaves.
+   */
+  private[sql] def mayServeUngrouped(required: ClusteredDistribution): Boolean =
+    collapsedLayoutMayBeGrouped && keysSatisfy(required)
+
   /**
    * Ask this only of a partitioning that is not grouped, since a grouped one has nothing to
    * coalesce and `keysMaySatisfy` covers both. See the class doc.
    */
   private[sql] def mayGroupToSatisfy(required: Distribution): Boolean = {
     val mayCoalesce = required match {
-      case _: ClusteredDistribution =>
-        !isCollapsed || SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys
+      case _: ClusteredDistribution => collapsedLayoutMayBeGrouped
       case _ => true
     }
     // The permission is the cheap half, so it is asked first.
@@ -1371,8 +1394,9 @@ case class PartitioningCollection(partitionings: Seq[Partitioning])
     // operation's still can, through the projection a `GroupPartitionsExec` performs. That is the
     // admission set this filter had before `satisfies` became strict, up to one shape it now also
     // keeps, a partition expression that *is* a cluster key, which `areKeysCompatible` turns away
-    // anyway. The set matters because `ValidateRequirements` builds a spec from a finished plan
-    // through here.
+    // anyway. The set matters because the planner's `shuffleToCoPartition` reads it through here to
+    // pick the layout the other children are laid out on, and a member this filter drops is one the
+    // collection cannot offer as that layout.
     //
     // Every admitted member stays, because `isCompatibleWith` answers for any of them and the
     // collection cannot know which one the other side matched. The cost is that
@@ -1420,10 +1444,13 @@ object PartitioningCollection {
    * stands, and only a keyed partitioning answers the two differently.
    *
    * A partitioning that is not grouped is not admitted, even though a node would also group it.
-   * This is the admission set `satisfies` gave the one caller before it became strict, and the
-   * caller feeds `ValidateRequirements` as well as the planner, so it does not widen what a
-   * finished plan is checked against. `EnsureRequirements.createKeyedShuffleSpecs` is where the
-   * planner asks the wider question, for a child it is about to group itself.
+   * This is the admission set `satisfies` gave the caller before it became strict, and it is what
+   * `PartitioningCollection.createShuffleSpec` answers with, which the planner's
+   * `shuffleToCoPartition` reads to lay the other children out on what this one reports.
+   * `ValidateRequirements` does not read it: `specsForPairing` is its admission, and it admits an
+   * ungrouped member only where partially clustered distribution reports one.
+   * `EnsureRequirements.createKeyedShuffleSpecs` is where the planner asks the wider question, for
+   * a child it is about to group itself.
    *
    * The partition count clause is here for consistency with the strict question. A node changes the
    * count, so the pre-grouping one is no prediction of it.
@@ -1437,6 +1464,74 @@ object PartitioningCollection {
     case other => other.satisfies(required)
   }
 
+  /**
+   * The specs `p` offers for `distribution`, one per member that may serve it, each of them the
+   * layout that member reports (`reportedSpecOf`): the partitions it has, in the order it reports
+   * them, with the partition expressions that carry no cluster key left out. Leaving those out is
+   * what lets two members pair when the operation clusters on part of what they are partitioned
+   * on, and it invents nothing: a kept key is that partition's own, the count is the member's own,
+   * and the order is the member's own.
+   *
+   * 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. It is admitted only
+   * where the caller says so, through `mayBeUngrouped`, which carries both halves of that
+   * admission: the configuration that builds such a shape, and the operator whose producer spreads
+   * it. `mayServeUngrouped` 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 cover only a subset of the operation's keys is a sound pairing.
+   *
+   * @param mayBeUngrouped whether an ungrouped keyed member of `p` may serve `distribution` here.
+   *                      Only the caller knows both halves of that; see above.
+   */
+  private[sql] def specsForPairing(
+      p: Partitioning,
+      distribution: ClusteredDistribution,
+      mayBeUngrouped: Boolean): Seq[ShuffleSpec] =
+    flatten(p).flatMap {
+      case k: KeyedPartitioning =>
+        val pairsAsIs = k.satisfies(distribution) ||
+          (mayBeUngrouped &&
+            distribution.requiredNumPartitions.forall(_ == k.numPartitions) &&
+            k.mayServeUngrouped(distribution))
+        Option.when(pairsAsIs)(reportedSpecOf(k, distribution))
+      case other =>
+        Option.when(other.satisfies(distribution))(other.createShuffleSpec(distribution))
+    }
+
+  /**
+   * The spec for a member a finished plan reports: its own layout with the partition expressions
+   * that carry no cluster key left out. No key is deduped and no key is re-sorted, so a side offers
+   * a view of the partitions the plan has rather than the layout
+   * `KeyedPartitioning.createShuffleSpec` builds for a node a planner would insert: that one dedups
+   * and sorts, and its count is the grouped one, which the plan does not hold.
+   *
+   * A projection here is a relabelling of the member's partitions, and this leans on the caller for
+   * that: only `keysSatisfy`'s projecting half admits such a member, and it holds the projection to
+   * one that merges no partition, handing the positions it checked to this method.
+   *
+   * A marked layout is not projected: its claim routes the rows of an undeclared key by a hash over
+   * these keys in this order, and dropping one breaks the claim. The clause cannot fire as the
+   * admission stands, and it is kept because the causality is not local: an admitted marked member
+   * has its keys covering the clustering structurally, since the projecting half of `keysSatisfy`
+   * excludes a marked one, so it leaves here unprojected either way.
+   */
+  private def reportedSpecOf(
+      k: KeyedPartitioning,
+      distribution: ClusteredDistribution): KeyedShuffleSpec = {
+    val positions = k.positionsCoveringClusterKeys(distribution).toSeq
+    if (positions.size == k.expressions.length || k.mayContainUnknownPartitionKeys) {
+      KeyedShuffleSpec(k, distribution)
+    } else {
+      KeyedShuffleSpec(k.project(positions), distribution, Some(positions))
+    }
+  }
+
   /**
    * The uniform `mayContainUnknownPartitionKeys` marker of `p`'s keyed members, read from its
    * first keyed member, the same one `checkKeyedPartitioningInvariant` compares against.
@@ -1639,12 +1734,14 @@ case object SinglePartitionShuffleSpec extends LeafShuffleSpec {
     // disagree when the subset config projects them onto different key sets.
     //
     // `EnsureRequirements` never reaches this: a spec whose `canCreatePartitioning` is false is
-    // never the best one. The one production caller that can put a collection on the `other` side
-    // is `ValidateRequirements`' `specs.tail.forall(_.isCompatibleWith(specs.head))`, and there the
-    // stricter answer is the safer one. It does leave this direction stricter than the collection's
-    // own `exists`, against the symmetry this trait's doc assumes, but nothing observable follows:
-    // only `KeyedShuffleSpec` can make members disagree on `numPartitions`, and it has no
-    // `SinglePartitionShuffleSpec` case, so that direction is already false.
+    // never the best one. Nor does any other production caller put a collection on the `other`
+    // side: `pickCoPartitionTarget` flattens before it pairs a member, and `ValidateRequirements`
+    // offers leaf specs and pairs those, where it used to compare every child against the first
+    // one's whole spec. So what is left is the answer `ShuffleSpecSuite` pins. It stays `forall`
+    // rather than the `exists` a collection answers with, against the symmetry the trait doc
+    // assumes, but nothing observable follows: only `KeyedShuffleSpec` can make members disagree on
+    // `numPartitions`, and it has no `SinglePartitionShuffleSpec` case, so that direction is
+    // already false.
     case ShuffleSpecCollection(specs) => specs.forall(isCompatibleWith)
   }
 

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala

@@ -93,9 +93,8 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
   test("SPARK-59289: createShuffleSpec drops a keyed member that is not grouped") {
     val a = AttributeReference("a", IntegerType)()
     val clustered = ClusteredDistribution(Seq(a))
-    // A node would group this one too, but the admission set here is what a finished plan is
-    // checked against by `ValidateRequirements`, so it stays what the strict question admitted
-    // before a projection was allowed to answer it.
+    // A node would group this one too, but a collection is a layout the plan holds, so it stays
+    // what the strict question admitted before a projection was allowed to answer it.
     val ungrouped = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1), InternalRow(2)))
     val grouped = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(2), InternalRow(3)))
     assert(!ungrouped.isGrouped && grouped.isGrouped, "test setup")
@@ -114,6 +113,75 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
     assert(specs.head.isInstanceOf[HashShuffleSpec], s"and it is the hash one, got $specs")
   }
 
+  test("SPARK-59671: the specs a side offers are the layouts it reports") {
+    val a = AttributeReference("a", IntegerType)()
+    val b = AttributeReference("b", IntegerType)()
+    val cd = ClusteredDistribution(Seq(a))
+    // An ungrouped keyed member is the layout partially clustered distribution spreads, and it is
+    // offered as the layout it holds: nobody is left to group it. Only a caller that knows both
+    // halves of that admission, the configuration and the operator whose producer spreads one, says
+    // so; the helper takes the answer rather than reading a configuration for itself.
+    val spread = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1), InternalRow(2)))
+    val specs = PartitioningCollection.specsForPairing(spread, cd, mayBeUngrouped = true)
+    assert(specs.size == 1, s"one member, one spec, got $specs")
+    val spec = specs.head.asInstanceOf[KeyedShuffleSpec]
+    assert((spec.partitioning eq spread) && spec.joinKeyPositions.isEmpty,
+      s"its own layout, no projection to make, got $specs")
+
+    // Everywhere else a member that only a grouping node would serve is not offered at all.
+    assert(PartitioningCollection.specsForPairing(spread, cd, mayBeUngrouped = false).isEmpty,
+      "an ungrouped member is a plan only where something builds one and produces it")
+
+    // The permission for a collapse is the planner's, and it survives into this admission: a
+    // collapsed member is not offered ungrouped where the planner would refuse to group it either.
+    val collapsed = spread.withLayout(_.copy(isCollapsed = true))
+    withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") {
+      assert(PartitioningCollection
+        .specsForPairing(collapsed, cd, mayBeUngrouped = true).isEmpty,
+        "a collapsed member the planner would not group is not offered ungrouped")
+    }
+
+    // The subset permission applies where the operation's keys are a subset of the member's
+    // partitioning keys, so a member may carry an expression the operation does not cluster on. It
+    // is offered as its own partitions under the key the operation clusters on: that expression is
+    // left out, and no key is deduped or re-sorted for it.
+    withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+      val source = KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 1), InternalRow(2, 2)))
+      val offered = PartitioningCollection.specsForPairing(source, cd, mayBeUngrouped = false)
+        .head.asInstanceOf[KeyedShuffleSpec]
+      assert(offered.partitioning.expressions === Seq(a) &&
+        offered.partitioning.numPartitions === source.numPartitions &&
+        offered.joinKeyPositions === Some(Seq(0)),
+        s"the member's own partitions under the operation's key, got $offered")
+
+      // A marked layout is never offered through a projection: the projecting half of `keysSatisfy`
+      // excludes a marked member, so the keys an offered marked one carries are the keys its claim
+      // is over.
+      val marked = source.withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+      assert(PartitioningCollection.specsForPairing(marked, cd, mayBeUngrouped = false).isEmpty,
+        "a marked member is not offered through a projection")
+    }
+
+    // A member whose keys do not cover the clustering is not offered, and a member that is not
+    // keyed is asked for its own spec.
+    val elsewhere = KeyedPartitioning(Seq(b), Seq(InternalRow(1), InternalRow(2)))
+    assert(PartitioningCollection.specsForPairing(elsewhere, cd, mayBeUngrouped = false).isEmpty)
+    assert(PartitioningCollection
+      .specsForPairing(HashPartitioning(Seq(a), 2), cd, mayBeUngrouped = false).head
+      .isInstanceOf[HashShuffleSpec])
+
+    // A count the operation pinned is asked as it stands: a member of another size is not offered
+    // even where its keys cover the clustering.
+    val four = KeyedPartitioning(Seq(a), (1 to 4).map(InternalRow(_)))
+    def pairingFour(pinned: Int): Seq[ShuffleSpec] =
+      PartitioningCollection.specsForPairing(
+        four, ClusteredDistribution(Seq(a), requiredNumPartitions = Some(pinned)),
+        mayBeUngrouped = false)
+    assert(pairingFour(4).size == 1,
+      "the size the operation asks for is the one the member reports")
+    assert(pairingFour(3).isEmpty, "a member of another size does not serve the distribution")
+  }
+
   protected def checkCompatible(
       left: ShuffleSpec,
       right: ShuffleSpec,

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala

@@ -591,13 +591,8 @@ case class EnsureRequirements(
       leftRequired: ClusteredDistribution,
       right: SparkPlan,
       rightRequired: ClusteredDistribution): Option[Seq[SparkPlan]] = {
-    parent match {
-      case smj: SortMergeJoinExec =>
-        checkKeyGroupCompatible(left, leftRequired, right, rightRequired, smj.joinType)
-      case sj: ShuffledHashJoinExec =>
-        checkKeyGroupCompatible(left, leftRequired, right, rightRequired, sj.joinType)
-      case _ =>
-        None
+    ShuffledJoin.partiallyClusteredJoinType(parent).flatMap { joinType =>
+      checkKeyGroupCompatible(left, leftRequired, right, rightRequired, joinType)
     }
   }
 
@@ -904,6 +899,15 @@ case class EnsureRequirements(
       }
 
       // Now we need to push-down the common partition information to the `GroupPartitionsExec`s.
+      //
+      // The two arguments below say which side does what, and exactly one of them is true:
+      // `replicateRightSide` is the negation of `replicateLeftSide`, and the branch above is taken
+      // only when the side it picked may replicate, so `applyPartialClustering` holds with one flag
+      // set. That split is what makes an aligned pair of these layouts sound, and it is the whole
+      // of it: for a key, a partition holding part of it on the spread side holds all of it on the
+      // side that repeats the group, so pairing the two index by index loses no match.
+      // `ValidateRequirements` reads such a pair as aligned and cannot tell it from two sides that
+      // split the key between them, so the decision has to be right here.
       (
         GroupPartitionsExec(rawLeft, leftSpec.joinKeyPositions,
           Some(mergedPartitionKeys), leftReducers,