[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

@@ -401,11 +401,11 @@ class DistributionSuite extends SparkFunSuite {
     assert(!nonGroupedKP.isGrouped)
     // satisfies() must return false: the partitions are not yet grouped.
     checkSatisfied(nonGroupedKP, ClusteredDistribution(Seq(x)), false)
-    // groupedSatisfies() returns true: it CAN satisfy once GroupPartitionsExec groups them.
-    assert(nonGroupedKP.groupedSatisfies(ClusteredDistribution(Seq(x))))
+    // mayGroupToSatisfy() returns true, because grouping them makes it satisfy.
+    assert(nonGroupedKP.mayGroupToSatisfy(ClusteredDistribution(Seq(x))))
 
     // Grouped: all distinct keys, so isGrouped=true and satisfies() delegates to
-    // groupedSatisfies().
+    // keysSatisfy().
     val groupedKP = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(2), InternalRow(3)))
     assert(groupedKP.isGrouped)
     checkSatisfied(groupedKP, ClusteredDistribution(Seq(x)), true)
@@ -466,4 +466,17 @@ 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)
+    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")
+  }
 }

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

@@ -131,18 +131,10 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
 
     if (projectablePositions.isEmpty) return LazyList.empty
 
-    // All input KPs share the same partitionKeys by invariant; use the first as the key source.
-    val keySource = kps.head
-    val sharedKeys =
-      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
+    // All input KPs share the same partitionKeys and isCollapsed flag by invariant, so the first
+    // one projects the keys and both flags for every combination below. Only the expressions
+    // differ.
+    val projected = kps.head.project(projectablePositions)
 
     // Cross-product the per-position alternatives to produce all concrete KPs.
     // Note: generateCartesianProduct expects thunks () => Seq[T], but wrapping LazyLists in thunks
@@ -151,8 +143,7 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
     // so all cross-product combinations are distinct by construction.
     MultiTransform.generateCartesianProduct(
       projectablePositions.map(i => () => alternativesPerPosition(i)))
-      .map(projectedExprs =>
-        new KeyedPartitioning(projectedExprs, sharedKeys, isGrouped, isNarrowed))
+      .map(projectedExprs => projected.copy(expressions = projectedExprs))
   }
 }
 

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

@@ -986,11 +986,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup
       // The `KeyedPartitioning`s must agree on the partition expressions to merge.
       val compatible = kps.forall(comparePartitioning(_, headKp))
       if (compatible) {
-        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)
+        return KeyedPartitioning.concat(kps)
       } else {
         return super.outputPartitioning
       }

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

@@ -68,21 +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.
-        val partitionKeys = groupedPartitions.map(_._1)
+        // can only differ in `expressions`. Their `partitionKeys` reference and `isCollapsed` flag
+        // are shared (enforced by `PartitioningCollection`), so the grouping is computed once.
+        val partitionKeys = grouping.partitions.map(_._1)
         p.transform {
           case k: KeyedPartitioning =>
             val projectedExpressions = joinKeyPositions.fold(k.expressions)(_.map(k.expressions))
-            KeyedPartitioning(projectedExpressions, partitionKeys, isGrouped = isGrouped)
+            KeyedPartitioning(
+              projectedExpressions, partitionKeys, grouping.isGrouped, grouping.isCollapsed)
         }.asInstanceOf[Partitioning]
       case o => o
     }
   }
 
-  /**
-   * Aligns partitions based on `expectedPartitionKeys` and clustering mode.
-   */
+  /** Aligns partitions based on `expectedPartitionKeys` and clustering mode. */
   private def alignToExpectedKeys(keyMap: Map[InternalRowComparableWrapper, Seq[Int]]) = {
     var isGrouped = true
     val alignedPartitions = expectedPartitionKeys.get.flatMap { case (key, numSplits) =>
@@ -119,15 +118,14 @@ case class GroupPartitionsExec(
    * 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:
-   * - 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)
+   * `isCollapsed` says whether the output stands for more than one of the child's partition keys.
+   * Two things set it: the child's own flag, and a merge this node performs.
    */
-  @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
+  @transient private lazy val grouping: PartitionGrouping = {
+    // There must be a `KeyedPartitioning` in the child's output partitioning, as a
+    // `GroupPartitionsExec` node is added to a plan only in that case. Any member will do, see
+    // `outputPartitioning` above.
+    val childKp = child.outputPartitioning
       .asInstanceOf[Partitioning with Expression]
       .collectFirst { case k: KeyedPartitioning => k }
       .getOrElse(
@@ -136,26 +134,47 @@ case class GroupPartitionsExec(
     // Project partition keys if join key positions are specified
     val (projectedDataTypes, projectedKeys) =
       joinKeyPositions.fold(
-        (keyedPartitioning.keyDataTypes, keyedPartitioning.partitionKeys)
-      )(keyedPartitioning.projectKeys)
+        (childKp.keyDataTypes, childKp.partitionKeys)
+      )(childKp.projectKeys)
 
     // Reduce keys if reducers are specified
     val (reducedDataTypes, reducedKeys) = reducers.fold((projectedDataTypes, projectedKeys))(
       KeyedPartitioning.reduceKeys(projectedKeys, projectedDataTypes, _))
 
     val keyToPartitionIndices = reducedKeys.zipWithIndex.groupMap(_._1)(_._2)
 
-    if (expectedPartitionKeys.isDefined) {
+    val (partitions, isGrouped) = if (expectedPartitionKeys.isDefined) {
       alignToExpectedKeys(keyToPartitionIndices)
     } else {
       (groupAndSortByKeys(keyToPartitionIndices, reducedDataTypes), true)
     }
+
+    // Both cheap terms come first, so the scan below runs only where a merge is possible. A
+    // grouping that left the keys as they are groups the child's own key values, and one of those
+    // groups can only ever cover the one key it was built from.
+    val keysChanged =
+      joinKeyPositions.exists(_.length < childKp.expressions.length) || reducers.isDefined
+    val isCollapsed = childKp.isCollapsed || keysChanged && {
+      // The groups this node keeps are the ones that can merge keys of the child, and asking the
+      // child's keys rather than its partitions is what tells such a merge from a source that
+      // reports several splits per key.
+      val keptGroups = expectedPartitionKeys match {
+        case Some(expected) => expected.view.flatMap { case (key, _) =>
+          keyToPartitionIndices.get(key)
+        }
+        case None => keyToPartitionIndices.values.view
+      }
+      val childKeys = childKp.partitionKeys.toArray
+      keptGroups.exists { group =>
+        val first = childKeys(group.head)
+        group.tail.exists(childKeys(_) != first)
+      }
+    }
+    PartitionGrouping(partitions, isGrouped, isCollapsed)
   }
 
   @transient lazy val groupedPartitions: Seq[(InternalRowComparableWrapper, Seq[Int])] =
-    groupedPartitionsTuple._1
-
-  @transient lazy val isGrouped: Boolean = groupedPartitionsTuple._2
+    grouping.partitions
 
   @transient private lazy val hasCoalescing: Boolean = groupedPartitions.exists(_._2.size > 1)
 
@@ -335,6 +354,12 @@ case class GroupPartitionsExec(
   }
 }
 
+/** What a [[GroupPartitionsExec]] computes once and reports from several members. */
+private case class PartitionGrouping(
+    partitions: Seq[(InternalRowComparableWrapper, Seq[Int])],
+    isGrouped: Boolean,
+    isCollapsed: Boolean)
+
 /**
  * A PartitionCoalescer that groups partitions according to a pre-computed grouping plan.
  *

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

@@ -73,20 +73,20 @@ case class EnsureRequirements(
           child
         } else {
           // Check KeyedPartitioning satisfaction conditions
-          val groupedSatisfies = grouped.find(_.satisfies(distribution))
+          val satisfyingGrouped = grouped.find(_.satisfies(distribution))
           val nonGroupedSatisfiesAsIs = nonGrouped.exists(_.nonGroupedSatisfies(distribution))
-          val nonGroupedSatisfiesWhenGrouped = nonGrouped.find(_.groupedSatisfies(distribution))
+          val groupableNonGrouped = nonGrouped.find(_.mayGroupToSatisfy(distribution))
 
           // Check if any KeyedPartitioning satisfies the distribution
-          if (groupedSatisfies.isDefined || nonGroupedSatisfiesAsIs
-              || nonGroupedSatisfiesWhenGrouped.isDefined) {
+          if (satisfyingGrouped.isDefined || nonGroupedSatisfiesAsIs
+              || groupableNonGrouped.isDefined) {
             distribution match {
               case o: OrderedDistribution =>
                 // OrderedDistribution requires grouped KeyedPartitioning with sorted keys
                 // according to the distribution's ordering.
-                // Find any KeyedPartitioning that satisfies via groupedSatisfies.
+                // Find any KeyedPartitioning that satisfies, grouped or groupable.
                 val satisfyingKeyedPartitioning =
-                  groupedSatisfies.orElse(nonGroupedSatisfiesWhenGrouped).get
+                  satisfyingGrouped.orElse(groupableNonGrouped).get
                 // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees
                 // one attribute per partition expression.
                 val attrs = satisfyingKeyedPartitioning.expressions.flatMap(_.references)
@@ -107,7 +107,7 @@ case class EnsureRequirements(
                   )
                 }
 
-              case _ if groupedSatisfies.isDefined =>
+              case _ if satisfyingGrouped.isDefined =>
                 // Grouped KeyedPartitioning already satisfies
                 child
 

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

@@ -594,7 +594,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,8 @@ 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 kp: KeyedPartitioning =>
+      kp.copy(expressions = kp.expressions.map(resolveAttrs(_, plan)))
     case PartitioningCollection(partitionings) =>
       PartitioningCollection(partitionings.map(resolvePartitioning(_, plan)))
     case RangePartitioning(ordering, numPartitions) =>