[SPARK-59285][SQL] Hold a KeyedPartitioning's shared partition layout in one value

### What changes were proposed in this pull request?

`KeyedPartitioning` grows a `KeyLayout`, and everything its members share moves into it.

**One value for what a partitioning's members share.** `KeyLayout(partitionKeys, dataTypes, isGrouped, isCollapsed, mayContainUnknownPartitionKeys)` holds everything about the partitions a `KeyedPartitioning` describes except the expressions naming them, so `KeyedPartitioning` is `(expressions, layout)`. The members of a `PartitioningCollection` name one layout with their own expressions and share the object by reference, so:

- the collection's invariant is one `eq` on the layout in place of a clause per shared field, and it now covers `isGrouped`, which the field-by-field check left out;
- `fromPartitionings` merges one canonical layout instead of interning the keys and ORing two flags, and refuses a member that describes another key space, which interning would otherwise retype;
- `KeyedShuffleSpec.createPartitioning` has one thing to decide rather than three, the unknown-keys marker it must set;
- `GroupPartitionsExec`'s `PartitionGrouping` is the layout it will report plus the child partitions each of its own is built from.

**One derivation of a key schema's ordering.** `InternalRowComparableWrapper`'s `Factory` also answers with the `ordering` it holds, and `KeyedPartitioning.apply` takes a `sortKeys` flag that sorts with it. `DataSourceV2ScanExecBase` used to derive an ordering of its own to sort the splits and then hand the keys to `apply`, which derived the type list and the ordering again. It now passes `sortKeys = true`, so one factory answers for the types, the ordering and the wrappers. That is dongjoon-hyun's point on #58523 (#58523 (comment)), which needed this PR's shape to fix cleanly.

**One answer for the key types, including where no key row is left.** `keyDataTypes` reads the layout rather than sampling the first key row and falling back to the partition expressions. A layout is given the types its keys were built at, at the four places one is built. `EnsureRequirements` therefore drops the exception SPARK-59176 added for a side with no key row, since the layout answers for it.

### Why are the changes needed?

Two things, one structural and one a defect the structure hides.

**The shared part of a `KeyedPartitioning` is four fields that every member of a collection has to agree on by hand.** `checkKeyedPartitioningInvariant` compares them clause by clause, and it left `isGrouped` out. Four places put a partitioning's expressions over keys they did not build, and each has to carry the shared fields forward correctly:

1. `GroupPartitionsExec.outputPartitioning` reports the keys `EnsureRequirements` merged, and picks its member with `collectFirst`, which need not be the member the planner merged from.
2. `PartitioningPreservingUnaryExecNode.projectKeyedPartitionings` projects `kps.head` once and stamps every alias alternative onto it with `copy(expressions = ...)`.
3. `KeyedShuffleSpec.createPartitioning` puts the other child's expressions over these keys.
4. `KeyedPartitioning.concat`, for a `UnionExec`.

Sharing one layout by reference is what makes all four correct rather than merely lucky, and it turns the invariant into one `eq`. It also makes the field count stop mattering: SPARK-59050's `mayContainUnknownPartitionKeys` is the fifth shared field, and with the layout it is one more `copy` argument rather than one more clause in the invariant, one more OR in `fromPartitionings` and one more argument at every copy site.

**A side with no key row answers from its partition expressions, and after a both-sides reduce that is a type no key of it holds.** The reduce leaves keys that are `r1(f1(x))` = `r2(f2(x))`, a space neither transform names, so the reported expression is marked and its own type is the un-reduced one. SPARK-59176 worked around it by leaving such a side out of the co-partition type check, which left the check not checking for the shape most likely to need it, and every other reader of `keyDataTypes` still getting the wrong answer. The layout now carries what the reduce produced, so the workaround goes.

The types are on the layout rather than derived, because a partitioning whose partitions were all pruned has no row to read them off and its expressions do not describe a reduced key space. They are on the *layout* rather than on `KeyedPartitioning`, because that is what makes them shared: two members that share a layout share its keys, so the collection's `eq` covers them, and no consumer can pair one member's rows with another member's types.

Two alternatives were tried and dropped. An independent `keyDataTypes` field on `KeyedPartitioning` has to be decided at each of the four sites above, and two of them mix members, which is how it produced two reachable regressions in review. A `TypedKeys(dataTypes, keys)` value object does not settle it either, since two members can still hold two different pairs.

Sharing by reference is a constraint as well as a convenience, and one site had to change for it. `ShuffledJoin`'s marker clearing used to rewrite each member on its own; with the layout that would give each a different one, so it now builds one cleared layout for the whole input.

The scan sorts through `apply` rather than sorting before it, because the ordering it wants is the one its keys are compared at. Deriving it separately worked only because a generated comparator is name-blind, so the raw and the erased type lists happen to give the same order. Taking it off the factory makes that an identity rather than a coincidence.

No plan string changes. `KeyedPartitioning.stringArgs` prints the layout's contents where the value object would print, and deliberately leaves the key types out of that list: they have their field names, nullability and metadata erased (SPARK-59187), so printing them would put a struct field named `0` into a plan that appears nowhere in the query.

### Does this PR introduce _any_ user-facing change?

No. It adds no behaviour of its own beyond making a pruned side report its own key types truthfully.

### How was this patch tested?

Four new tests, plus SPARK-59176's two existing ones, which now pass with its exception removed.

Ablation: with the exception removed and `keyDataTypes` derived from the rows and expressions again, "SPARK-59176: a leg reduced onto no key at all still joins" fails with the error SPARK-59176 was filed for.

- `DistributionSuite`, "fromPartitionings refuses a member that disagrees on isGrouped", for the layout itself.
- `GroupPartitionsExecSuite`, "a reduced key space's type reaches the reported partitioning with no key left": a both-sides reduce onto `LongType` under a `DateType` transform, with keys and without.
- `KeyGroupedPartitioningSuite`, "two sides whose partitions were all pruned are not one layout": two legs pruned to nothing, one `identity(id)` on `LongType` and one `bucket(4, id)` on `IntegerType`, joined and then joined again through a FULL OUTER that brings real keys in. It asserts that no node reports two key spaces as one layout, that the plan passes `ValidateRequirements`, and the answer.
- `KeyGroupedPartitioningSuite`, "two legs whose struct field names differ are still co-partitioned", the shape where two exact type lists differ while the space does not.

SPARK-59050's two collection tests now assert on the layout reference rather than on a marker clause of their own, which is the same guarantee reached by the structure instead of by a check.

`KeyGroupedPartitioningSuite`, `KeyGroupedPartitioningRuntimeFilterSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `ProjectedOrderingAndPartitioningSuite`, `DataSourceV2CatalystRuntimeFilterSuite`, `PlannerSuite`, `DataFrameSetOperationsSuite`, `DistributionSuite`, `ShuffleSpecSuite`, `TransformExpressionSuite` and `InternalRowComparableWrapperSuite`, `DataSourceV2Suite`, 561 tests. Scalastyle and scalafmt clean.

### Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 5)

Author: peter-toth

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

@@ -446,7 +446,7 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
 
     val combined = PartitioningCollection.fromPartitionings(Seq(nested, kpY))
     val interned = combined.partitionings.last.asInstanceOf[KeyedPartitioning]
-    assert(interned.partitionKeys eq kpX.partitionKeys)
+    assert(interned.layout eq kpX.layout, "the whole layout is interned, keys and all")
   }
 
   test("SPARK-59050: a marked one-partition layout keeps the global ordering claim") {
@@ -457,9 +457,9 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
     val a = AttributeReference("a", IntegerType)()
     val ordered = OrderedDistribution(Seq(SortOrder(a, Ascending)))
     val markedOne = KeyedPartitioning(Seq(a), Seq(InternalRow(1)))
-      .copy(mayContainUnknownPartitionKeys = true)
+      .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
     val markedTwo = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(2)))
-      .copy(mayContainUnknownPartitionKeys = true)
+      .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
     withSQLConf(SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
       checkSatisfied(markedOne, ordered, true)
       checkSatisfied(markedTwo, ordered, false)
@@ -470,7 +470,7 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
     val x = AttributeReference("x", IntegerType)()
     val y = AttributeReference("y", IntegerType)()
     val marked = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(2)))
-      .copy(mayContainUnknownPartitionKeys = true)
+      .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
     val plain = KeyedPartitioning(Seq(y), Seq(InternalRow(1), InternalRow(2)))
     val combined = PartitioningCollection.fromPartitionings(Seq(marked, plain))
     // The conservative direction: an unmarked member must never excuse marked data, because
@@ -485,21 +485,23 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
     val y = AttributeReference("y", IntegerType)()
     val keys = Seq(InternalRow(1), InternalRow(2))
     val base = KeyedPartitioning(Seq(x), keys)
-    val marked = base.copy(mayContainUnknownPartitionKeys = true)
-    // Same keys reference, arity and isCollapsed, only the marker disagrees: it has to reach
-    // the marker require rather than the reference check ahead of it.
-    val disagree = marked.copy(expressions = Seq(y), mayContainUnknownPartitionKeys = false)
+    val marked = base.withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+    // The marker lives on the layout, so two members that disagree on it hold two layouts, and the
+    // one `eq` on the layout refuses them. There is no clause of its own to reach.
+    val disagree =
+      KeyedPartitioning(Seq(y), marked.layout.copy(mayContainUnknownPartitionKeys = false))
     val err = intercept[IllegalArgumentException] {
       PartitioningCollection(Seq(marked, disagree))
     }
-    assert(err.getMessage.contains("agree on mayContainUnknownPartitionKeys"))
+    assert(err.getMessage.contains("share the same KeyLayout reference"))
   }
 
   test("SPARK-59050: fromPartitionings normalizes the unknown-keys marker through nesting") {
     val x = AttributeReference("x", IntegerType)()
     val y = AttributeReference("y", IntegerType)()
     val keys = Seq(InternalRow(1), InternalRow(2))
-    val marked = KeyedPartitioning(Seq(x), keys).copy(mayContainUnknownPartitionKeys = true)
+    val marked =
+      KeyedPartitioning(Seq(x), keys).withLayout(_.copy(mayContainUnknownPartitionKeys = true))
     val plainNested = PartitioningCollection.fromPartitionings(
       Seq(KeyedPartitioning(Seq(y), keys)))
     val combined = PartitioningCollection.fromPartitionings(Seq(marked, plainNested))
@@ -509,8 +511,8 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
       .collect { case k: KeyedPartitioning => k }
     assert(leaves.length === 2, combined.toString)
     assert(leaves.forall(_.mayContainUnknownPartitionKeys), leaves.toString)
-    assert(leaves.forall(_.partitionKeys eq leaves.head.partitionKeys),
-      "the interned keys must survive the rebuild")
+    assert(leaves.forall(_.layout eq leaves.head.layout),
+      "the interned layout must survive the rebuild")
   }
 
   test("SPARK-59050: PartitioningCollection requires a nested collection to agree on the " +
@@ -520,15 +522,15 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
     val keys = Seq(InternalRow(1), InternalRow(2))
     val base = KeyedPartitioning(Seq(x), keys)
     val markedNested = PartitioningCollection.fromPartitionings(Seq(
-      base.copy(mayContainUnknownPartitionKeys = true)))
+      base.withLayout(_.copy(mayContainUnknownPartitionKeys = true))))
     val plain = base.copy(expressions = Seq(y))
-    // The nested collection is internally uniform, so its own construction passes; the outer
-    // constructor must still catch its representative against the unmarked sibling. Same keys
-    // reference, arity and isCollapsed, only the marker disagrees.
+    // The nested collection is internally uniform, so its own construction passes. The outer
+    // constructor must still catch its representative against the unmarked sibling, and the marker
+    // living on the layout is what makes that one `eq` enough.
     val err = intercept[IllegalArgumentException] {
       PartitioningCollection(Seq(markedNested, plain))
     }
-    assert(err.getMessage.contains("agree on mayContainUnknownPartitionKeys"))
+    assert(err.getMessage.contains("share the same KeyLayout reference"))
   }
 
   test("SPARK-56877: PartitioningCollection enforces the invariant through nesting") {
@@ -542,7 +544,7 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
     val refMismatch = intercept[IllegalArgumentException] {
       PartitioningCollection(Seq(nested, kpY))
     }
-    assert(refMismatch.getMessage.contains("share the same partitionKeys reference"))
+    assert(refMismatch.getMessage.contains("share the same KeyLayout reference"))
 
     val kpXY = KeyedPartitioning(Seq(x, y), Seq(InternalRow(1, 1), InternalRow(2, 2)))
     val arityMismatch = intercept[IllegalArgumentException] {
@@ -551,12 +553,28 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
     assert(arityMismatch.getMessage.contains("matching expression arity"))
   }
 
+  test("SPARK-59285: fromPartitionings refuses a member that disagrees on isGrouped") {
+    // Interning now replaces a member's layout whole, where it used to keep each member's own
+    // `isGrouped`. Whether the keys are unique is a property of the keys, so a member that
+    // disagrees about it over equal keys is wrong, and interning would hide that.
+    val x = AttributeReference("x", IntegerType)()
+    val y = AttributeReference("y", IntegerType)()
+    val keys = Seq(InternalRow(1), InternalRow(2))
+
+    val kpX = KeyedPartitioning(Seq(x), keys)
+    val ungrouped = KeyedPartitioning(Seq(y), keys).withLayout(_.copy(isGrouped = false))
+    val mismatch = intercept[IllegalArgumentException] {
+      PartitioningCollection.fromPartitionings(Seq(kpX, ungrouped))
+    }
+    assert(mismatch.getMessage.contains("must agree on isGrouped"))
+  }
+
   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)
+      .withLayout(_.copy(isCollapsed = true))
     assert(collapsedKP.toGrouped.isCollapsed, "toGrouped must keep isCollapsed sticky")
 
     val spec = KeyedShuffleSpec(collapsedKP, ClusteredDistribution(Seq(x)))

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

@@ -547,7 +547,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
         keys: Seq[Int],
         hasUnknown: Boolean = false): KeyedShuffleSpec = KeyedShuffleSpec(
       KeyedPartitioning(Seq(a), keys.map(k => InternalRow(k)))
-        .copy(mayContainUnknownPartitionKeys = hasUnknown), distribution)
+        .withLayout(_.copy(mayContainUnknownPartitionKeys = hasUnknown)), distribution)
 
     // A partitioning with unknown partition keys (e.g. a side re-shuffled onto a keyed layout by
     // `KeyedShuffleSpec.createPartitioning`) only guarantees co-location for its declared keys, so
@@ -602,12 +602,12 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
         keys: Seq[Long],
         hasUnknown: Boolean = false): KeyedShuffleSpec = KeyedShuffleSpec(
       KeyedPartitioning(Seq(bucket(4, a)), keys.map(k => InternalRow(k)))
-        .copy(mayContainUnknownPartitionKeys = hasUnknown), distribution)
+        .withLayout(_.copy(mayContainUnknownPartitionKeys = hasUnknown)), distribution)
     def keyedSpec(
         keys: Seq[Long],
         hasUnknown: Boolean = false): KeyedShuffleSpec = KeyedShuffleSpec(
       KeyedPartitioning(Seq(a), keys.map(k => InternalRow(k)))
-        .copy(mayContainUnknownPartitionKeys = hasUnknown), distribution)
+        .withLayout(_.copy(mayContainUnknownPartitionKeys = hasUnknown)), distribution)
 
     // `isExpressionCompatible` admits an identity-vs-transform pair when compatible transforms
     // are allowed, but then the two sides' partition keys live in different domains: raw values
@@ -636,7 +636,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
     val a = $"a".int
     val b = $"b".int
     val marked = KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 2), InternalRow(3, 4)))
-      .copy(mayContainUnknownPartitionKeys = true)
+      .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
     withSQLConf(
         SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
         SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
@@ -682,7 +682,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
         KeyedPartitioning(
           Seq(TransformExpression(fn, Seq(a), Some(numBuckets))),
           Seq(InternalRow(0L), InternalRow(1L)))
-          .copy(mayContainUnknownPartitionKeys = hasUnknown),
+          .withLayout(_.copy(mayContainUnknownPartitionKeys = hasUnknown)),
         ClusteredDistribution(Seq(a)))
 
     // `allowCompatibleTransforms` lets a differing-bucket-count pair through

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

@@ -98,15 +98,15 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
    * achievable granularity. Positions that cannot be expressed in the output are dropped.
    *
    * The resulting [[KeyedPartitioning]]s are the cross-product of the per-position alternatives
-   * restricted to the projectable positions. All share the same `partitionKeys` object (projected
-   * to the same subset of positions), preserving the invariant required by
+   * restricted to the projectable positions. All share the same `KeyLayout` object (projected to
+   * the same subset of positions), preserving the invariant required by
    * [[PartitioningCollection]].
    */
   private def projectKeyedPartitionings(
       kps: Seq[KeyedPartitioning]): LazyList[KeyedPartitioning] = {
     if (kps.isEmpty) return LazyList.empty
-    // All input KPs share the same `partitionKeys` reference and matching arity by the
-    // [[PartitioningCollection]] invariant (the only producer of multi-KP inputs here).
+    // All input KPs have matching arity by the [[PartitioningCollection]] invariant (the only
+    // producer of multi-KP inputs here).
     val numPositions = kps.head.expressions.length
 
     val alternativesPerPosition: IndexedSeq[LazyList[Expression]] =
@@ -136,20 +136,16 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode
 
     if (projectablePositions.isEmpty) return LazyList.empty
 
-    // `PartitioningCollection` requires its members to agree on the marker, so the head
-    // represents them all.
-    val mayContainUnknownPartitionKeys = kps.head.mayContainUnknownPartitionKeys
-
     // Dropping a key position coarsens the declared set, which an unknown-keyed claim cannot
-    // survive.
-    if (projectablePositions.length < numPositions && mayContainUnknownPartitionKeys) {
+    // survive. The members share one layout, so the head answers for the marker.
+    if (projectablePositions.length < numPositions && kps.head.mayContainUnknownPartitionKeys) {
       return LazyList.empty
     }
 
-    // All input KPs share the same partitionKeys and flags by invariant, so the first one
-    // projects the keys for every combination below; only the expressions differ. The marker
-    // rides the copies unchanged: the guard above turned away the one shape that could not, a
-    // narrowing projection of a marked collection.
+    // All input KPs share one layout by invariant, so the first one projects it for every
+    // combination below. Only the expressions differ, and the marker rides the copies unchanged:
+    // the guard above turned away the one shape that could not, a narrowing projection of a marked
+    // collection.
     val projected = kps.head.project(projectablePositions)
 
     // Cross-product the per-position alternatives to produce all concrete KPs.

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

@@ -952,12 +952,17 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup
     (left, right) match {
       case (SinglePartition, SinglePartition) => true
       case (l: HashPartitioningLike, r: HashPartitioningLike) => l == r
-      // For `KeyedPartitioning`, only the partition expressions must match (both sides'
-      // expressions have already been remapped to this union's output attributes by
-      // `prepareOutputPartitioning`). The partition keys are intentionally not compared here:
-      // children typically carry different key sets, and `outputPartitioning` merges them.
+      // For `KeyedPartitioning`, the partition expressions must match (both sides' expressions
+      // have already been remapped to this union's output attributes by
+      // `prepareOutputPartitioning`), and so must the types their keys were built with. The
+      // partition keys themselves are intentionally not compared: children typically carry
+      // different key sets, and `outputPartitioning` merges them. Their types cannot be merged the
+      // same way, since one type list has to stand for every row of the concatenation, and a
+      // wrapper compares its types before its values, so keys of two types would never be found
+      // equal to one another.
       case (l: KeyedPartitioning, r: KeyedPartitioning) =>
         l.expressions.length == r.expressions.length &&
+          l.keyDataTypes == r.keyDataTypes &&
           l.expressions.zip(r.expressions).forall { case (le, re) => le.semanticEquals(re) }
       // Note: two `RangePartitioning`s with even same ordering and number of partitions
       // are not equal, because they might have different partition bounds.

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

@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2
 
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Ascending, Expression, RowOrdering, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, Expression, SortOrder}
 import org.apache.spark.sql.catalyst.plans.physical
 import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning
 import org.apache.spark.sql.catalyst.util.truncatedString
@@ -104,11 +104,13 @@ trait DataSourceV2ScanExecBase
     keyGroupedPartitioning match {
       case Some(exprs) if conf.v2BucketingEnabled && KeyedPartitioning.supportsExpressions(exprs) &&
           inputPartitions.nonEmpty && inputPartitions.forall(_.isInstanceOf[HasPartitionKey]) =>
-        val dataTypes = exprs.map(_.dataType)
-        val rowOrdering = RowOrdering.createNaturalAscendingOrdering(dataTypes)
-        val partitionKeys =
-          inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).sorted(rowOrdering)
-        Some(KeyedPartitioning(exprs, partitionKeys))
+        // `sortKeys` rather than sorting here, so the ordering and the type list come off the one
+        // factory that also wraps the keys, instead of this deriving a second ordering over the
+        // same schema.
+        Some(KeyedPartitioning(
+          exprs,
+          inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()),
+          sortKeys = true))
       case _ => None
     }
   }