[SPARK-59289][SQL] Internal review

Author: peter-toth

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

@@ -582,6 +582,8 @@ case class KeyLayout(
  *
  * - `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`.
+ *   Strict for a `ClusteredDistribution`; its `OrderedDistribution` arm is a gate on what may claim
+ *   a global ordering, and the order itself is still the caller's to check.
  * - `keysCanSatisfy()`: `keysSatisfy()`, or the keys co-locate them after a node has projected
  *   away the expressions that carry none. Only
  *   `spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys` admits the second half.
@@ -818,22 +820,18 @@ case class KeyedPartitioning(
   /**
    * The counts `numPartitionsProjectedOn` has already worked out, by position set.
    *
-   * Memoized here rather than at a caller because there is no one caller. `satisfies` asks this
-   * question now, and `EnsureRequirements`, `ValidateRequirements` and every `AQEShuffleReadRule`
-   * application all ask `satisfies` of the same partitioning, so a memo scoped to one rule pass
-   * would miss most of the repeats. The partitioning is immutable and the answer depends only on
-   * the keys and the position set, so the cache cannot go stale.
+   * Kept on the partitioning rather than at a caller because there is no one caller: `satisfies`
+   * asks this question, and the planner, `ValidateRequirements` and every `AQEShuffleReadRule`
+   * application all ask `satisfies` of the same partitioning. The partitioning is immutable and the
+   * answer depends only on the keys and the position set, so the cache cannot go stale.
    *
    * Not a constructor field, so it stays out of the product and leaves equality, `hashCode` and
-   * canonicalization alone. The map is built on first use, and only for a partitioning something
-   * actually asks a narrowing question of. On the default configuration nothing does, since
-   * `mayProjectToClusterKeys` requires
+   * canonicalization alone. The map is built on first use, and on the default configuration nothing
+   * asks a narrowing question at all, since `mayProjectToClusterKeys` requires
    * `spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys`.
    *
-   * [[TransientBestEffortLazyVal]] rather than a `lazy val`, for the reason that class exists: a
-   * Scala 2 `lazy val` locks the instance to initialize it, and this is read from whichever thread
-   * is planning or validating. Two threads racing here each build an empty map and one is dropped,
-   * which is all the best-effort part costs.
+   * [[TransientBestEffortLazyVal]] because a Scala 2 `lazy val` locks the instance to initialize
+   * it, and this is read from whichever thread is planning or validating.
    */
   private val projectedPartitionCounts =
     new TransientBestEffortLazyVal[ConcurrentHashMap[BitSet, java.lang.Integer]](
@@ -935,8 +933,8 @@ case class KeyedPartitioning(
       expressions.forall(_.references.size == 1)
 
   /**
-   * The strict question of the family the class doc lists. `true` only when the
-   * keys as they stand co-locate every cluster key, with nothing left for a
+   * The strict question of the family the class doc lists, for a [[ClusteredDistribution]]. `true`
+   * only when the keys as they stand co-locate every cluster key, with nothing left for a
    * `GroupPartitionsExec` to do about it.
    *
    * Two ways to be true. Every partition expression is a function of cluster keys alone, so
@@ -945,6 +943,12 @@ case class KeyedPartitioning(
    * second is the only branch that reads a partition key, and only
    * `spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys` admits it, so the ordinary
    * configuration answers structurally.
+   *
+   * The strictness claim stops at [[ClusteredDistribution]]. The [[OrderedDistribution]] arm is a
+   * local gate on what may claim a global ordering at all, and it does not read the key order, so a
+   * 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.
    */
   private def keysSatisfy(required: Distribution): Boolean = {
     required match {
@@ -1383,24 +1387,18 @@ case class PartitioningCollection(partitionings: Seq[Partitioning])
   override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = {
     // `maySatisfyAfterProjection`, not `satisfies`. A spec says what its partitioning could
     // co-partition on, and a grouped `KeyedPartitioning` whose keys are coarser than the
-    // operation's still can, through the projection a `GroupPartitionsExec` performs. The strict
-    // question would drop it, which is narrower than what this filter admitted before `satisfies`
-    // became strict.
-    //
-    // It is not wider either. `ValidateRequirements` builds a spec from a finished plan through
-    // here, and is the gate that makes AQE drop a rewrite that breaks co-partitioning, so this
-    // admits exactly what it admitted before and no more. `checkKeyedPartitioningInvariant` forces
-    // one `KeyLayout` on all keyed members, so `isGrouped` is uniform across them and the filter
-    // either keeps every keyed member or none: the latter only for a partitioning that does not
-    // serve the distribution through a keyed member at all, which neither caller passes.
+    // 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.
     //
-    // Every admitted member has to stay, because `isCompatibleWith` answers for any of them and
-    // the collection cannot know which one the other side matched. That has a cost worth knowing:
-    // `KeyedShuffleSpec.canCreatePartitioning` is false without `v2BucketingShuffleEnabled` and
-    // `ShuffleSpecCollection.canCreatePartitioning` is a `forall`, so a groupable keyed member
-    // beside a usable non-keyed one would cost the collection its role as a shuffle template. No
-    // operator is known to report that mixture, since `EnsureRequirements` groups a keyed child
-    // before it can reach a join's output.
+    // 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
+    // `ShuffleSpecCollection.canCreatePartitioning` is a `forall` over members whose keyed half is
+    // false without `v2BucketingShuffleEnabled`, so a groupable keyed member beside a usable
+    // non-keyed one would cost the collection its role as a shuffle template. No operator is known
+    // to report that mixture.
     val filtered =
       partitionings.filter(PartitioningCollection.maySatisfyAfterProjection(_, distribution))
     ShuffleSpecCollection(filtered.map(_.createShuffleSpec(distribution)))
@@ -1446,18 +1444,15 @@ object PartitioningCollection {
    * finished plan is checked against. `EnsureRequirements.createKeyedShuffleSpecs` is where the
    * planner asks the wider question, for a child it is about to group itself.
    *
-   * A required partition count still has to match, which is what `satisfies` applies to every other
-   * partitioning. A node changes the count, so the pre-grouping one is not a bound on what it will
-   * be. This clause is here for consistency with the strict question, not as a prediction.
+   * 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.
    */
   private[sql] def maySatisfyAfterProjection(
       p: Partitioning,
-      required: ClusteredDistribution): Boolean = p match {
+      required: ClusteredDistribution): Boolean = flatten(p).exists {
     case k: KeyedPartitioning =>
       k.isGrouped && required.requiredNumPartitions.forall(_ == k.numPartitions) &&
         k.keysMaySatisfy(required)
-    case pc: PartitioningCollection =>
-      pc.partitionings.exists(maySatisfyAfterProjection(_, required))
     case other => other.satisfies(required)
   }
 

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

@@ -63,20 +63,28 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
     }
   }
 
-  test("SPARK-59289: createShuffleSpec drops a member that is not grouped") {
+  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. Asked of the predicate rather than of a mixed
-    // collection, since `checkKeyedPartitioningInvariant` forces one `KeyLayout` on all keyed
-    // members and `isGrouped` is therefore uniform across them.
+    // 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")
 
     assert(!PartitioningCollection.maySatisfyAfterProjection(ungrouped, clustered))
     assert(PartitioningCollection.maySatisfyAfterProjection(grouped, clustered))
+
+    // And through the filter itself, which is the predicate's one caller. The keyed member goes and
+    // the hash member stays, so the collection keeps a spec and its `require` is not tripped. The
+    // pairing is asked of a mixed collection rather than of two keyed members, because
+    // `checkKeyedPartitioningInvariant` forces one `KeyLayout` on all keyed members and `isGrouped`
+    // is therefore uniform across them: a collection cannot hold one of each.
+    val mixed = PartitioningCollection(Seq(ungrouped, HashPartitioning(Seq(a), 3)))
+    val specs = mixed.createShuffleSpec(clustered).asInstanceOf[ShuffleSpecCollection].specs
+    assert(specs.size == 1, s"only the hash member survives the filter, got $specs")
+    assert(specs.head.isInstanceOf[HashShuffleSpec], s"and it is the hash one, got $specs")
   }
 
   protected def checkCompatible(

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala

@@ -53,10 +53,9 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
  *                            the child still reports what it was planned over, see
  *                            `outputPartitioning`.
  * @param childPartitioning The child's partitioning this node was decided for. Kept so that a node
- *                          handed a different child can say so. A `Partitioning` that is an
- *                          `Expression` is in the case class product on purpose, so that
- *                          `QueryPlan.doCanonicalize` normalizes its exprIds along with
- *                          `plannedPartitioning`'s and the comparison stays meaningful.
+ *                          handed a different child can say so, and in the case class product so
+ *                          that `QueryPlan.doCanonicalize` normalizes its exprIds along with
+ *                          `plannedPartitioning`'s, which is what keeps the comparison meaningful.
  * @param joinKeyPositions Optional projection selecting a subset of the partitioning key positions,
  *                         so that partitions sharing the projected key are coalesced. Used whenever
  *                         the cluster keys are a subset of the partition keys, either the join
@@ -103,17 +102,18 @@ case class GroupPartitionsExec(
    * node a child reporting something else invalidates both, so the keyed claim is dropped and only
    * the physical output count is reported. `ValidateRequirements` then finds the parent's
    * distribution unsatisfied and `AdaptiveSparkPlanExec` reverts the rewrite, which is what an
-   * `AQEShuffleReadExec` landing over a keyed shuffle stage relies on.
+   * `AQEShuffleReadExec` landing over a keyed shuffle stage relies on. That revert is also what
+   * keeps `doExecute` off the changed child, since it still coalesces on the planning-time
+   * `grouping.partitions`. The base relied on the same revert, having re-derived to get there.
    *
-   * Re-deriving from the new child instead would be worse. The recipe was chosen for a pairing, and
-   * applying it to a different child yields a claim nothing validated against the other side. It
-   * would also buy nothing today: no rule in the tree is known to report a *different*
-   * `KeyedPartitioning`, so the two reachable outcomes are an equal one, where this check passes,
-   * and `UnknownPartitioning`, where re-deriving reports exactly what this reports.
+   * Re-deriving instead would apply a recipe chosen for one pairing to a different child, yielding
+   * a claim nothing validated against the other side. It would also buy nothing today: no rule is
+   * known to report a *different* `KeyedPartitioning`, so the reachable outcomes are an equal one,
+   * where this check passes, and `UnknownPartitioning`, where re-deriving reports the same thing.
    *
-   * Asked on the read rather than in `withNewChildInternal` because plan canonicalization rebuilds
-   * this node over a canonicalized child and never reads its partitioning, which it cannot: a
-   * canonicalized `BatchScanExec` throws from `reportedKeyedPartitioning`.
+   * Asked on the read rather than in `withNewChildInternal` because canonicalization rebuilds this
+   * node over a canonicalized child whose partitioning cannot be read at all: a canonicalized
+   * `BatchScanExec` throws from `reportedKeyedPartitioning`.
    */
   @transient override lazy val outputPartitioning: Partitioning =
     if (child.outputPartitioning == childPartitioning) {
@@ -417,15 +417,11 @@ case class GroupPartitionsExec(
  * @param isIdentity whether the grouping leaves the declared keys and every partition where they
  *                   were. Nothing rewrote the keys, output partition i holds exactly input
  *                   partition i, and there is one output per input. That is the only grouping that
- *                   keeps a marked layout's undeclared rows at hash(key) % numPartitions. A
- *                   projection or reduction re-labels the groups into a different key space, so
- *                   even a grouping whose indices line up would pin the claim to keys it no longer
- *                   declares. A rewrite is rejected up front, covering a narrowing
- *                   projection, a reordering one, and any reducer slot. A reducer slot counts as
- *                   key-changing even though a conforming self-reducer cannot rewrite a reachable
- *                   key value, so the give-up there loses at most an optimization. A grouping that
- *                   drops trailing declared keys reads identity for every group it keeps, but the
- *                   partition count shrinks and the hash modulus with it.
+ *                   keeps a marked layout's undeclared rows at hash(key) % numPartitions, so a
+ *                   projection, a reduction or a dropped trailing key all fail it: each moves the
+ *                   claim into a different key space or changes the modulus. A reducer slot counts
+ *                   as key-changing even where a conforming self-reducer could not rewrite a
+ *                   reachable key value, so the give-up there loses at most an optimization.
  * @param numPrunedPartitions the alignment's effect on the reads of the child's splits, 0 outside
  *                            the alignment path. See `alignToExpectedKeys`.
  * @param numReplicatedPartitionReads as above.
@@ -501,15 +497,16 @@ private[sql] object GroupPartitionsExec {
     // describes the keys, so the reduce marks it (see `KeyedShuffleSpec.reducersBothWays`).
     //
     // A marked claim pins undeclared rows to hash(key) % numPartitions (see [[KeyLayout]]'s
-    // `mayContainUnknownPartitionKeys`). Only an identity grouping keeps that relationship: a
-    // reorder, a coalesce, or a resize moves those rows, and a projection or reduction rewrites
-    // the keys the claim speaks for (see `PartitionGrouping.isIdentity`). Clearing only the marker
-    // would misreport the undeclared rows that remain, so give up the keyed partitioning at the
-    // physical output count (one per group, padding included) that a parent's
-    // `PartitioningCollection` requires for uniformity. The give-up deliberately under-reports.
-    // The node still physically groups the partitions but no longer claims a keyed layout, so a
-    // join planned over it does not see its required distribution satisfied. `EnsureRequirements`
-    // asks this before it commits to a pairing and falls back to a shuffle when it happens.
+    // `mayContainUnknownPartitionKeys`), and only an identity grouping keeps that relationship (see
+    // `PartitionGrouping.isIdentity`). Clearing the marker alone would misreport the undeclared
+    // rows that remain, so the keyed partitioning goes and only the physical output count is
+    // reported, one per group with padding included, which is the count a parent's
+    // `PartitioningCollection` needs for uniformity. `EnsureRequirements` asks this before it
+    // commits to a pairing and falls back to a shuffle when it happens.
+    //
+    // The marker read off `p` is the same bit the `else` branch publishes through
+    // `grouping.layout`, since `computeGrouping` copies the child's layout without touching it. A
+    // change to that copy has to move this guard with it.
     if (PartitioningCollection.keyedMarkerOf(p).contains(true) && !grouping.isIdentity) {
       UnknownPartitioning(grouping.partitions.size)
     } else {