[SPARK-59057][SQL] Make KeyedPartitioning.isNarrowed mean actual key collapse, and rename it to isCollapsed

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -2515,9 +2515,11 @@ object SQLConf {
         "tables. At planning time, Spark will group the partitions by only those keys that are " +
         "in the operation's keys. That is currently enabled only if " +
         s"${REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key} is false. This config also gates " +
-        "grouping a partitioning that was narrowed to a subset of its keys and whose keys are no " +
-        "longer distinct, which carries the same risk of skew; that applies regardless of " +
-        s"${REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key}."
+        "grouping a partitioning whose keys collapsed, that is, where a projection or a " +
+        "reduction mapped keys that were distinct in the source onto the same key, so that " +
+        "grouping them would produce a partition larger than any the source declared. That " +
+        s"applies regardless of ${REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key}. It does not " +
+        "apply to duplicate keys the source itself reported, which are grouped without this config."
       )
       .version("4.0.0")
       .withBindingPolicy(ConfigBindingPolicy.SESSION)

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

@@ -466,4 +466,22 @@ class DistributionSuite extends SparkFunSuite {
     }
     assert(arityMismatch.getMessage.contains("matching expression arity"))
   }
+
+  test("SPARK-59057: toGrouped and KeyedShuffleSpec.createPartitioning keep isCollapsed sticky") {
+    val x = AttributeReference("x", IntegerType)()
+    val y = AttributeReference("y", IntegerType)()
+
+    val collapsedKP = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(1), InternalRow(2)))
+      .copy(isCollapsed = true)
+    // The protection this flag provides: a coarsened, non-grouped KP must not satisfy
+    // ClusteredDistribution via grouping without allowKeysSubsetOfPartitionKeys.
+    assert(!collapsedKP.groupedSatisfies(ClusteredDistribution(Seq(x))))
+
+    assert(collapsedKP.toGrouped.isCollapsed, "toGrouped must keep isCollapsed sticky")
+
+    val spec = KeyedShuffleSpec(collapsedKP, ClusteredDistribution(Seq(x)))
+    val created = spec.createPartitioning(Seq(y)).asInstanceOf[KeyedPartitioning]
+    assert(created.isCollapsed, "createPartitioning must keep isCollapsed sticky: the shuffled " +
+      "side is laid out on the collapsed side's partition keys and inherits its skew risk")
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala

@@ -137,12 +137,24 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
       if (projectablePositions.length == numPositions) keySource.partitionKeys
       else keySource.projectKeys(projectablePositions)._2
 
-    val isGrouped = sharedKeys.distinct.size == sharedKeys.size
-    // A KP is narrowed if this node drops positions, or if the input KPs were already narrowed
-    // (i.e. came from a finer-grained partitioning). The flag must be sticky: a subsequent
-    // PartitioningPreservingUnaryExecNode that passes all positions through would otherwise
-    // recompute isNarrowed=false, silently dropping the protection.
-    val isNarrowed = projectablePositions.length < numPositions || keySource.isNarrowed
+    val distinctSharedKeys = sharedKeys.distinct
+    val isGrouped = distinctSharedKeys.size == sharedKeys.size
+    // This projection collapses keys when it maps keys that were distinct in the input
+    // onto the same projected key -- dropping positions is not enough on its own, since the
+    // projected keys can stay just as distinct as the originals. The flag is sticky: a subsequent
+    // PartitioningPreservingUnaryExecNode that passes all positions through must not recompute it
+    // as false and drop the protection, and no projection can make a partitioning finer again.
+    //
+    // Both cheap terms come first: an inherited flag or a projection that drops no position
+    // settles the question without counting distinct keys. A pass-through projection cannot
+    // collapse anything, since it keeps the input's keys as they are.
+    //
+    // The inherited flag is read from all inputs rather than from the key source alone. A
+    // `PartitioningCollection` normalizes it across its members, so the two agree today; reading
+    // all of them keeps this producer correct without depending on that.
+    val isCollapsed = kps.exists(_.isCollapsed) ||
+      (projectablePositions.length < numPositions &&
+        keySource.collapsesOnProjection(distinctSharedKeys.size))
 
     // Cross-product the per-position alternatives to produce all concrete KPs.
     // Note: generateCartesianProduct expects thunks () => Seq[T], but wrapping LazyLists in thunks
@@ -152,7 +164,7 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
     MultiTransform.generateCartesianProduct(
       projectablePositions.map(i => () => alternativesPerPosition(i)))
       .map(projectedExprs =>
-        new KeyedPartitioning(projectedExprs, sharedKeys, isGrouped, isNarrowed))
+        new KeyedPartitioning(projectedExprs, sharedKeys, isGrouped, isCollapsed))
   }
 }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala

@@ -989,8 +989,12 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup
         val mergedKeys = kps.flatMap(_.partitionKeys)
         val mergedExpressions = headKp.expressions
         val isGrouped = mergedKeys.distinct.size == mergedKeys.size
-        val isNarrowed = kps.exists(_.isNarrowed)
-        return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isNarrowed)
+        // A collapse is a property of each child's own keys: a child key that stands for several
+        // finer-grained keys keeps standing for them in the concatenation. Keys repeating *across*
+        // children is not a collapse -- grouping merges partitions that share a key, which is what
+        // it does for any partitioning that never went through one, and needs no opt-in.
+        val isCollapsed = kps.exists(_.isCollapsed)
+        return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isCollapsed)
       } else {
         return super.outputPartitioning
       }

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

@@ -68,13 +68,20 @@ case class GroupPartitionsExec(
     child.outputPartitioning match {
       case p: Partitioning with Expression =>
         // There can be multiple `KeyedPartitioning`s in an output partitioning of a join, but they
-        // can only differ in `expressions`; their `partitionKeys` reference is shared (enforced by
-        // `PartitioningCollection`), so `groupedPartitions` is computed only once.
+        // can only differ in `expressions`; their `partitionKeys` reference and `isCollapsed` flag
+        // are shared (enforced by `PartitioningCollection`), so both `groupedPartitions` and the
+        // new flag are computed once, outside the transform.
         val partitionKeys = groupedPartitions.map(_._1)
+        // `isCollapsed` is sticky: grouping removes the duplicate keys, but it does not make this
+        // partitioning any finer than the layout it came from. Projecting onto the join key
+        // positions, or reducing the keys onto a coarser transform, can collapse keys in its own
+        // right, which `collapsesKeys` answers exactly -- it is computed with the grouping itself,
+        // so it costs nothing per call and needs no comparison of key counts.
+        val isCollapsed = childIsCollapsed || collapsesKeys
         p.transform {
           case k: KeyedPartitioning =>
             val projectedExpressions = joinKeyPositions.fold(k.expressions)(_.map(k.expressions))
-            KeyedPartitioning(projectedExpressions, partitionKeys, isGrouped = isGrouped)
+            KeyedPartitioning(projectedExpressions, partitionKeys, isGrouped, isCollapsed)
         }.asInstanceOf[Partitioning]
       case o => o
     }
@@ -112,26 +119,32 @@ case class GroupPartitionsExec(
     keyMap.toSeq.sorted(keyOrdering.on((t: (InternalRowComparableWrapper, _)) => t._1.row))
   }
 
+  // There must be a `KeyedPartitioning` in the child's output partitioning, as a
+  // `GroupPartitionsExec` node is added to a plan only in that case. A collection's members share
+  // their `partitionKeys` reference and their `isCollapsed` flag, so any one of them will do.
+  @transient private lazy val childKeyedPartitioning: KeyedPartitioning =
+    child.outputPartitioning
+      .asInstanceOf[Partitioning with Expression]
+      .collectFirst { case k: KeyedPartitioning => k }
+      .getOrElse(
+        throw new SparkException("GroupPartitionsExec requires a child with KeyedPartitioning"))
+
   /**
    * Computes the grouped partitions by:
    * 1. Projecting partition keys if joinKeyPositions is specified
    * 2. Reducing keys if reducers are specified
    * 3. Grouping input partition indices by their (possibly projected/reduced) keys
    * 4. Sorting or distributing based on whether partial clustering is enabled
    *
-   * Returns a tuple of (partitions, isGrouped) where:
+   * Returns a tuple of (partitions, isGrouped, collapsesKeys) where:
    * - partitions: sequence of (partitionKey, inputPartitionIndices) pairs representing
    *   how input partitions should be grouped together
    * - isGrouped: whether the output partitioning is grouped (no duplicates in partition keys)
+   * - collapsesKeys: whether any group this node outputs covers more than one of the child's own
+   *   partition keys, i.e. whether the projection or the reduction actually merged keys
    */
   @transient private lazy val groupedPartitionsTuple = {
-    // There must be a `KeyedPartitioning` in child's output partitioning as a
-    // `GroupPartitionsExec` node is added to a plan only in that case.
-    val keyedPartitioning = child.outputPartitioning
-      .asInstanceOf[Partitioning with Expression]
-      .collectFirst { case k: KeyedPartitioning => k }
-      .getOrElse(
-        throw new SparkException("GroupPartitionsExec requires a child with KeyedPartitioning"))
+    val keyedPartitioning = childKeyedPartitioning
 
     // Project partition keys if join key positions are specified
     val (projectedDataTypes, projectedKeys) =
@@ -145,10 +158,26 @@ case class GroupPartitionsExec(
 
     val keyToPartitionIndices = reducedKeys.zipWithIndex.groupMap(_._1)(_._2)
 
+    // Whether this node collapses keys: does any key it keeps stand for more than one of the
+    // child's own partition keys? Counting the child's *keys* rather than its partitions is what
+    // tells a collapse from a source reporting several splits per key, and asking it of the keys
+    // this node keeps is what tells it from `alignToExpectedKeys` dropping keys, which merges
+    // nothing. Ask the key groups, not the partitions finally emitted: `distributePartitions`
+    // spreads a group's splits over one partition each, which would hide the merge, and
+    // replication would ask about the same group repeatedly.
+    val childKeys = keyedPartitioning.partitionKeys.toIndexedSeq
+    def coversSeveralChildKeys(indices: Seq[Int]): Boolean =
+      indices.map(childKeys).distinct.size > 1
+
     if (expectedPartitionKeys.isDefined) {
-      alignToExpectedKeys(keyToPartitionIndices)
+      val (alignedPartitions, grouped) = alignToExpectedKeys(keyToPartitionIndices)
+      val keptGroups = expectedPartitionKeys.get.map { case (key, _) =>
+        keyToPartitionIndices.getOrElse(key, Seq.empty)
+      }
+      (alignedPartitions, grouped, keptGroups.exists(coversSeveralChildKeys))
     } else {
-      (groupAndSortByKeys(keyToPartitionIndices, reducedDataTypes), true)
+      val sorted = groupAndSortByKeys(keyToPartitionIndices, reducedDataTypes)
+      (sorted, true, sorted.map(_._2).exists(coversSeveralChildKeys))
     }
   }
 
@@ -157,6 +186,12 @@ case class GroupPartitionsExec(
 
   @transient lazy val isGrouped: Boolean = groupedPartitionsTuple._2
 
+  /** Whether the grouping this node performs merges keys the child held apart. */
+  @transient private lazy val collapsesKeys: Boolean = groupedPartitionsTuple._3
+
+  /** The child's own flag, from the same `KeyedPartitioning` the grouping was computed from. */
+  @transient private lazy val childIsCollapsed: Boolean = childKeyedPartitioning.isCollapsed
+
   @transient private lazy val hasCoalescing: Boolean = groupedPartitions.exists(_._2.size > 1)
 
   // Whether the child subtree is safe to use with SortedMergeCoalescedRDD (k-way merge).

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

@@ -499,7 +499,7 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession {
     val partitioning = KeyedPartitioning(
       Seq(partAttr),
       Seq(InternalRowComparableWrapper(InternalRow(1), Seq(partAttr))),
-      isGrouped = false)
+      isGrouped = false, isCollapsed = false)
 
     def replanAfterFiltering(afterFilter: Seq[InputPartition]): Unit = {
       val scan = new PartitioningBreakingScan(Seq(KeyedInputPartition(1)), afterFilter)

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

@@ -53,8 +53,9 @@ abstract class DistributionAndOrderingSuiteBase
       plan: QueryPlan[T]): Partitioning = partitioning match {
     case HashPartitioning(exprs, numPartitions) =>
       HashPartitioning(exprs.map(resolveAttrs(_, plan)), numPartitions)
-    case KeyedPartitioning(expressions, partitionKeys, isGrouped, _) =>
-      KeyedPartitioning(expressions.map(resolveAttrs(_, plan)), partitionKeys, isGrouped)
+    case KeyedPartitioning(expressions, partitionKeys, isGrouped, isCollapsed) =>
+      KeyedPartitioning(
+        expressions.map(resolveAttrs(_, plan)), partitionKeys, isGrouped, isCollapsed)
     case PartitioningCollection(partitionings) =>
       PartitioningCollection(partitionings.map(resolvePartitioning(_, plan)))
     case RangePartitioning(ordering, numPartitions) =>