[SPARK-59272][SQL][FOLLOWUP] Split the two refusal cases in the comments, and pin the new one

Three comment corrections from a review round, and the unit test the changed contract was
missing.

`createShuffleSpec` refuses to project any marked layout, but what the unprojected spec is then
good for differs between the two cases. Where every partition expression covers a clustering key
the spec is usable, and a consumer is laid out on this side's own key order. Where the layout
narrows, one `keyPositions` entry is empty and `canCreatePartitioning` and `areKeysCompatible`
both turn the spec away, which is what the sentence this commit brings back used to say. The
`@param joinKeyPositions` doc and the comment at the guard now split the two.

`ShuffleSpecSuite` gains the covering case beside the narrowing one it already had, so the
contract this PR changed is pinned at the unit that changed. It fails with the guard back at its
narrowing-only form.

The gate comment in `EnsureRequirements` loses a sentence fragment and the cost argument against
re-asking on the `compatibleAsIs` path. `describesSameKeys` walks the same keys on the push path,
so the reason that holds is the other one already given: there the children are the ones the
pairing read.

The `SPARK-59050` test comment numbered the shuffles wrong. The final join's own shuffle cannot
have existed while that join storage-partitioned, so before SPARK-59050 the plan held the first
and third of today's four. The fourth is that join's one-side shuffle, not a re-shuffle of a
regrouped side, since a declined pairing discards the node it tried.

Raised by dongjoon-hyun.

Author: peter-toth

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

@@ -917,8 +917,13 @@ case class KeyedPartitioning(
       // either would also need a `GroupPartitionsExec` on this side, and that node gives the keyed
       // claim up rather than perform a non-identity regrouping of a marked layout.
       //
-      // The unprojected spec that comes back is still usable, and is what a consumer is offered:
-      // the other child is laid out on the keys in the order this one reports them.
+      // What the unprojected spec is good for depends on the layout. Where every partition
+      // expression covers a clustering key, so the projection would only have re-sorted, the spec
+      // is usable and is what a consumer is offered: the other child is laid out on the keys in the
+      // order this one reports them. Where the layout narrows, one `keyPositions` entry is empty,
+      // and `canCreatePartitioning` and `areKeysCompatible` both turn the spec away. Nothing is
+      // lost there, because `keysSatisfy` does not call such a child grouped on the operation keys
+      // either.
       if (mayContainUnknownPartitionKeys) {
         return result
       }
@@ -1770,12 +1775,14 @@ case class IdentityReducer(transform: TransformExpression) extends Reducer[Any,
  *                         projection changed it. `None` therefore means `partitioning` is the one
  *                         the child reports, and a consumer needs no `GroupPartitionsExec` to
  *                         produce it. Only `v2BucketingAllowKeysSubsetOfPartitionKeys` projects at
- *                         all, and it reaches `None` two ways: an identity projection over already
- *                         grouped and sorted keys rebuilds the same partitioning, and a marked
- *                         claim is refused outright, since neither narrowing nor sorting its keys
- *                         leaves the routing it promises for its undeclared rows. Both say the
- *                         child is grouped on the operation keys as it stands, and both leave a
- *                         spec a consumer can be laid out on. See
+ *                         all, and it reaches `None` two ways. An identity projection over already
+ *                         grouped and sorted keys rebuilds the same partitioning. A marked claim is
+ *                         refused outright, since neither narrowing nor sorting its keys leaves the
+ *                         routing it promises for its undeclared rows. The first always leaves a
+ *                         spec a consumer can be laid out on, and so does the second where every
+ *                         partition expression covers a clustering key. A marked claim that narrows
+ *                         does not: one `keyPositions` entry is empty there, which
+ *                         `canCreatePartitioning` and `areKeysCompatible` both turn away. See
  *                         `KeyedPartitioning.createShuffleSpec`.
  */
 case class KeyedShuffleSpec(

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

@@ -650,6 +650,29 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
     }
   }
 
+  test("createShuffleSpec: a marked covering projection yields a usable unprojected spec") {
+    val a = $"a".int
+    // Grouped but not sorted, which is what a one-side shuffle onto a union's key order leaves
+    // behind. Every partition expression covers a clustering key here, so the projection that is
+    // refused would only have re-sorted the keys.
+    val marked = KeyedPartitioning(Seq(a),
+      Seq(InternalRow(3), InternalRow(4), InternalRow(1), InternalRow(2)))
+      .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+    assert(marked.isGrouped, "test setup: the keys are distinct, so the layout is grouped")
+    withSQLConf(
+        SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+        SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+      val spec = marked.createShuffleSpec(ClusteredDistribution(Seq(a)))
+        .asInstanceOf[KeyedShuffleSpec]
+      // The refusal leaves the child's own layout, and unlike the narrowing case above that spec
+      // is usable: a consumer is laid out on the keys in the order this side reports them. Sorting
+      // them would send the consumer's undeclared rows where this side does not hold them.
+      assert(spec.joinKeyPositions.isEmpty)
+      assert(spec.canCreatePartitioning)
+      assert(spec.partitioning.partitionKeys == marked.partitionKeys)
+    }
+  }
+
   test("areKeysCompatible: unknown keys require the same function, not just a compatible one") {
     // A bucket-like reducible function: like the built-in `bucket`, a pair of the same function
     // with coarser/finer bucket counts is compatible (a reducer exists) but not the same

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

@@ -794,27 +794,25 @@ case class EnsureRequirements(
     }
 
     // The pairing is only worth committing to if both children still declare the same aligned key
-    // sequence once the grouping has been pushed into them. They can fail that. A
-    // `GroupPartitionsExec` gives up its keyed claim when it turns out to regroup a layout that
-    // pins undeclared rows to `hash(key) % numPartitions` (see
-    // `KeyLayout.mayContainUnknownPartitionKeys`), and only the node knows the permutation it
-    // performs, so that answer arrives after the pairing was chosen. Asking before returning is
-    // what keeps the join from skipping both shuffles for a child that no longer satisfies its
-    // distribution, which is a plan `ValidateRequirements` rejects and every AQE rule that needs a
-    // valid plan then refuses to touch.
+    // sequence once the grouping has been pushed into them. A `GroupPartitionsExec` gives up its
+    // keyed claim when it turns out to regroup a layout that pins undeclared rows to
+    // `hash(key) % numPartitions` (see `KeyLayout.mayContainUnknownPartitionKeys`), and only the
+    // node knows the permutation it performs, so that answer arrives after the pairing was chosen.
+    // Asking before returning is what keeps the join from skipping both shuffles for a child that
+    // no longer satisfies its distribution, which is a plan `ValidateRequirements` rejects and
+    // every AQE rule that needs a valid plan then refuses to touch.
     //
     // The check is pairwise, not a per-side `satisfies`. Partially clustered distribution leaves
     // both children value-aligned yet not grouped on purpose, so a per-side gate would refuse that
     // whole family. What both sides owe each other is the key sequence `alignToExpectedKeys`
     // guarantees, each key repeated as many times as the merge expects, whichever side replicates.
-    // Through `KeyLayout.describesSameKeys`, which carries the reason the key types are compared as
-    // well as the rows.
+    // `KeyLayout.describesSameKeys` asks exactly that, and it compares the key types as well as
+    // the rows.
     //
     // Only the push branch rebuilds the children, so only it has to be asked. Where it did not run,
     // the children are the ones the pairing read: `KeyedShuffleSpec.isCompatibleWith` already ends
     // in `describesSameKeys`, and all keyed members of a `PartitioningCollection` share one
     // `KeyLayout`, so whichever member the spec matched on declares what the representative does.
-    // Asking again there would walk both sides' partition keys for an answer that cannot be no.
     def declaredLayout(plan: SparkPlan): Option[KeyLayout] =
       PartitioningCollection.representativeOf(plan.outputPartitioning).map(_.layout)
     def sidesDeclareSameKeys: Boolean =

sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala

@@ -8249,9 +8249,9 @@ class KeyGroupedPartitioningSuite
         checkAnswer(df, expected)
         // The regrouped marked side can no longer claim the hash-routing contract. Four one-side
         // shuffles, all keyed with unknown partition keys: rt onto the union, rs onto the marked
-        // side once the second join declines, rt2 onto ra2, and the final join's re-shuffle of the
-        // side that was regrouped. Before SPARK-59050 the final join storage-partitioned (only the
-        // first, third and fourth shuffles) and silently lost the id=5 row.
+        // side once the second join declines, rt2 onto ra2, and the final join's one-side shuffle.
+        // Before SPARK-59050 the final join storage-partitioned (only the first and third
+        // shuffles) and silently lost the id=5 row.
         //
         // The second join's shuffle is this fix's. The give-up happens inside the node, so it used
         // to arrive after that join had already committed to the pairing and skipped both