[SPARK-56164][SQL] Fix SPJ merged key ordering

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

Fix a bug in `EnsureRequirements` where the ordering used to merge and dedup partition keys after reduction was based on the original partition key types rather than the reduced types.

When compatible partition transform reducers are applied (e.g. reducing days keys to years keys), the resulting partition key values may have a different `DataType` than the originals.

Please note that #54884 fixed the issue when left and right side reduced key types are not equal, but this PR fixes the issue when the common reduced types differ to the left side original types.

### Why are the changes needed?

Without this fix, when reducers change the data type of partition keys (e.g. `DateType` → `LongType`), the ordering used to merge partition values is built for the wrong type, which can produce incorrect merge results or runtime failures during storage-partitioned joins with compatible transform reducers.

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

This PR fixes a theoretical bug, a real world example is unlikely to exist.

### How was this patch tested?

Added a new test case.

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

Generated-by: Claude Sonnet 4.6

Closes #54961 from peter-toth/SPARK-56164-fix-spj-merged-key-ordering.

Authored-by: Peter Toth <peter.toth@gmail.com>
Signed-off-by: Peter Toth <peter.toth@gmail.com>

Author: peter-toth

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

@@ -517,17 +517,24 @@ case class EnsureRequirements(
         val (rightReducedDataTypes, rightReducedKeys) = rightReducers.fold(
           (rightPartitioning.expressionDataTypes, rightPartitioning.partitionKeys)
         )(rightPartitioning.reduceKeys)
-        if (leftReducedDataTypes != rightReducedDataTypes) {
+        val reducedDataTypes = if (leftReducedDataTypes == rightReducedDataTypes) {
+          leftReducedDataTypes
+        } else {
           throw QueryExecutionErrors.storagePartitionJoinIncompatibleReducedTypesError(
             leftReducers = leftReducers,
             leftReducedDataTypes = leftReducedDataTypes,
             rightReducers = rightReducers,
             rightReducedDataTypes = rightReducedDataTypes)
         }
 
+        val reducedKeyRowOrdering = RowOrdering.createNaturalAscendingOrdering(reducedDataTypes)
+        val reducedKeyOrdering =
+          reducedKeyRowOrdering.on((t: InternalRowComparableWrapper) => t.row)
+
         // merge values on both sides
-        var mergedPartitionKeys = mergeAndDedupPartitions(leftReducedKeys, rightReducedKeys,
-          joinType, leftPartitioning.keyOrdering).map((_, 1))
+        var mergedPartitionKeys =
+          mergeAndDedupPartitions(leftReducedKeys, rightReducedKeys, joinType, reducedKeyOrdering)
+            .map((_, 1))
 
         logInfo(log"After merging, there are " +
           log"${MDC(LogKeys.NUM_PARTITIONS, mergedPartitionKeys.size)} partitions")

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

@@ -75,17 +75,23 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
       Column.create("dept_id", IntegerType),
       Column.create("data", StringType))
 
-  def withFunction[T](fn: UnboundFunction)(f: => T): T = {
-    val id = Identifier.of(Array.empty, fn.name())
-    val oldFn = Option.when(catalog.listFunctions(Array.empty).contains(id)) {
-      val fn = catalog.loadFunction(id)
-      catalog.dropFunction(id)
-      fn
-    }
-    catalog.createFunction(id, fn)
+  def withFunction[T](fns: UnboundFunction*)(f: => T): T = {
+    val fnIds = catalog.listFunctions(Array.empty)
+    val oldFns = fns.map { fn =>
+      val id = Identifier.of(Array.empty, fn.name())
+      val oldFn = Option.when(fnIds.contains(id)) {
+        val fn = catalog.loadFunction(id)
+        catalog.dropFunction(id)
+        fn
+      }
+      catalog.createFunction(id, fn)
+      (id, oldFn)
+    }
     try f finally {
-      catalog.dropFunction(id)
-      oldFn.foreach(catalog.createFunction(id, _))
+      oldFns.foreach { case (id, oldFn) =>
+        catalog.dropFunction(id)
+        oldFn.foreach(catalog.createFunction(id, _))
+      }
     }
   }
 
@@ -3441,7 +3447,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
   }
 
   test("SPARK-56046: Reducers with different result types") {
-    withFunction(UnboundDaysFunctionWithIncompatibleResultTypeReducer) {
+    withFunction(UnboundDaysFunctionWithToYearsReducerWithDateResult) {
       val items_partitions = Array(days("arrive_time"))
       createTable(items, itemsColumns, items_partitions)
       sql(s"INSERT INTO testcat.ns.$items VALUES " +
@@ -3478,4 +3484,48 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
       }
     }
   }
+
+  test("SPARK-56164: Reducers with different result types to original keys") {
+    withFunction(
+      UnboundDaysFunctionWithToYearsReducerWithLongResult,
+      UnboundYearsFunctionWithToYearsReducerWithLongResult) {
+      val items_partitions = Array(days("arrive_time"))
+      createTable(items, itemsColumns, items_partitions)
+      sql(s"INSERT INTO testcat.ns.$items VALUES " +
+        s"(0, 'aa', 39.0, cast('2020-01-01' as timestamp)), " +
+        s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+        s"(2, 'bb', 41.0, cast('2021-01-03' as timestamp)), " +
+        s"(3, 'bb', 42.0, cast('2021-01-04' as timestamp))")
+
+      val purchases_partitions = Array(years("time"))
+      createTable(purchases, purchasesColumns, purchases_partitions)
+      sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+        s"(1, 42.0, cast('2020-01-01' as timestamp)), " +
+        s"(5, 44.0, cast('2020-01-15' as timestamp)), " +
+        s"(7, 46.5, cast('2021-02-08' as timestamp))")
+
+      withSQLConf(
+        SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+        SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+        Seq(
+          s"testcat.ns.$items i JOIN testcat.ns.$purchases p ON p.time = i.arrive_time",
+          s"testcat.ns.$purchases p JOIN testcat.ns.$items i ON i.arrive_time = p.time"
+        ).foreach { joinString =>
+          val df = sql(
+            s"""
+               |${selectWithMergeJoinHint("i", "p")} id, item_id
+               |FROM $joinString
+               |ORDER BY id, item_id
+               |""".stripMargin)
+
+          val shuffles = collectShuffles(df.queryExecution.executedPlan)
+          assert(shuffles.isEmpty, "should not add shuffle for both sides of the join")
+          val groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan)
+          assert(groupPartitions.forall(_.outputPartitioning.numPartitions == 2))
+
+          checkAnswer(df, Seq(Row(0, 1), Row(1, 1)))
+        }
+      }
+    }
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala

@@ -24,23 +24,34 @@ import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
-object UnboundYearsFunction extends UnboundFunction {
+abstract class UnboundYearsFunctionBase extends UnboundFunction {
+  protected def isValidType(dt: DataType): Boolean = dt match {
+    case DateType | TimestampType => true
+    case _ => false
+  }
+
+  override def description(): String = name()
+  override def name(): String = "years"
+}
+
+object UnboundYearsFunction extends UnboundYearsFunctionBase {
   override def bind(inputType: StructType): BoundFunction = {
     if (inputType.size == 1 && isValidType(inputType.head.dataType)) YearsFunction
     else throw new UnsupportedOperationException(
       "'years' only take date or timestamp as input type")
   }
+}
 
-  private def isValidType(dt: DataType): Boolean = dt match {
-    case DateType | TimestampType => true
-    case _ => false
+object UnboundYearsFunctionWithToYearsReducerWithLongResult extends UnboundYearsFunctionBase {
+  override def bind(inputType: StructType): BoundFunction = {
+    if (inputType.size == 1 && isValidType(inputType.head.dataType)) {
+      YearsFunctionWithToYearsReducerWithLongResult
+    } else throw new UnsupportedOperationException(
+      "'years' only take date or timestamp as input type")
   }
-
-  override def description(): String = name()
-  override def name(): String = "years"
 }
 
-object YearsFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] {
+abstract class YearsFunctionBase[O] extends ScalarFunction[Int] with ReducibleFunction[Int, O] {
   override def inputTypes(): Array[DataType] = Array(TimestampType)
   override def resultType(): DataType = IntegerType
   override def name(): String = "years"
@@ -49,14 +60,37 @@ object YearsFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int
   val UTC: ZoneId = ZoneId.of("UTC")
   val EPOCH_LOCAL_DATE: LocalDate = Instant.EPOCH.atZone(UTC).toLocalDate
 
-  def invoke(ts: Long): Int = {
+  protected def doInvoke(ts: Long): Long = {
     val localDate = DateTimeUtils.microsToInstant(ts).atZone(UTC).toLocalDate
-    ChronoUnit.YEARS.between(EPOCH_LOCAL_DATE, localDate).toInt
+    ChronoUnit.YEARS.between(EPOCH_LOCAL_DATE, localDate)
   }
+}
 
+// This `years` function reduces `IntegerType` partition keys to `IntegerType` partition keys when
+// partitions are reduced to partitions of a `days` function, which produces `DateType` keys.
+object YearsFunction extends YearsFunctionBase[Int]  {
+  def invoke(ts: Long): Int = doInvoke(ts).toInt
   override def reducer(otherFunction: ReducibleFunction[_, _]): Reducer[Int, Int] = null
 }
 
+// This `years` function reduces `IntegerType` partition keys to `LongType` partition keys when
+// partitions are reduced to partitions of a `days` function, which produces `DateType` keys.
+object YearsFunctionWithToYearsReducerWithLongResult extends YearsFunctionBase[Long] {
+  def invoke(ts: Long): Int = doInvoke(ts).toInt
+  override def reducer(otherFunction: ReducibleFunction[_, _]): Reducer[Int, Long] = {
+    if (otherFunction == DaysFunctionWithToYearsReducerWithLongResult) {
+      YearsToYearsReducerWithLongResult()
+    } else {
+      null
+    }
+  }
+}
+
+case class YearsToYearsReducerWithLongResult() extends Reducer[Int, Long] {
+  override def resultType(): DataType = LongType
+  override def reduce(days: Int): Long = days.toLong
+}
+
 abstract class UnboundDaysFunctionBase extends UnboundFunction {
   protected def isValidType(dt: DataType): Boolean = dt match {
     case DateType | TimestampType => true
@@ -75,16 +109,25 @@ object UnboundDaysFunction extends UnboundDaysFunctionBase {
   }
 }
 
-object UnboundDaysFunctionWithIncompatibleResultTypeReducer extends UnboundDaysFunctionBase {
+object UnboundDaysFunctionWithToYearsReducerWithDateResult extends UnboundDaysFunctionBase {
+  override def bind(inputType: StructType): BoundFunction = {
+    if (inputType.size == 1 && isValidType(inputType.head.dataType)) {
+      DaysFunctionWithToYearsReducerWithDateResult
+    } else throw new UnsupportedOperationException(
+      "'days' only take date or timestamp as input type")
+  }
+}
+
+object UnboundDaysFunctionWithToYearsReducerWithLongResult extends UnboundDaysFunctionBase {
   override def bind(inputType: StructType): BoundFunction = {
     if (inputType.size == 1 && isValidType(inputType.head.dataType)) {
-      DaysFunctionWithIncompatibleResultTypeReducer
+      DaysFunctionWithToYearsReducerWithLongResult
     } else throw new UnsupportedOperationException(
       "'days' only take date or timestamp as input type")
   }
 }
 
-abstract class DaysFunctionBase extends ScalarFunction[Int] with ReducibleFunction[Int, Int] {
+abstract class DaysFunctionBase[O] extends ScalarFunction[Int] with ReducibleFunction[Int, O] {
   override def inputTypes(): Array[DataType] = Array(TimestampType)
   override def resultType(): DataType = DateType
   override def name(): String = "days"
@@ -93,7 +136,7 @@ abstract class DaysFunctionBase extends ScalarFunction[Int] with ReducibleFuncti
 
 // This `days` function reduces `DateType` partition keys to `IntegerType` partition keys when
 // partitions are reduced to partitions of a `years` function, which produces `IntegerType` keys.
-object DaysFunction extends DaysFunctionBase {
+object DaysFunction extends DaysFunctionBase[Int] {
   override def reducer(otherFunc: ReducibleFunction[_, _]): Reducer[Int, Int] = {
     if (otherFunc == YearsFunction) {
       DaysToYearsReducer()
@@ -105,32 +148,51 @@ object DaysFunction extends DaysFunctionBase {
 
 // This `days` function reduces `DateType` partition keys to `DateType` partition keys when
 // partitions are reduced to partitions of a `years` function, which produces `IntegerType` keys.
-object DaysFunctionWithIncompatibleResultTypeReducer extends DaysFunctionBase {
+object DaysFunctionWithToYearsReducerWithDateResult extends DaysFunctionBase[Int] {
   override def reducer(otherFunc: ReducibleFunction[_, _]): Reducer[Int, Int] = {
     if (otherFunc == YearsFunction) {
-      DaysToYearsReducerWithIncompatibleResultType()
+      DaysToYearsReducerWithDateResult()
     } else {
       null
     }
   }
 }
 
-abstract class DaysToYearsReducerBase extends Reducer[Int, Int] {
+// This `days` function reduces `DateType` partition keys to `LongType` partition keys when
+// partitions are reduced to partitions of a `years` function, which produces `IntegerType` keys.
+object DaysFunctionWithToYearsReducerWithLongResult extends DaysFunctionBase[Long] {
+  override def reducer(otherFunc: ReducibleFunction[_, _]): Reducer[Int, Long] = {
+    if (otherFunc == YearsFunctionWithToYearsReducerWithLongResult) {
+      DaysToYearsReducerWithLongResult()
+    } else {
+      null
+    }
+  }
+}
+
+abstract class DaysToYearsReducerBase[O] extends Reducer[Int, O] {
   val UTC: ZoneId = ZoneId.of("UTC")
   val EPOCH_LOCAL_DATE: LocalDate = Instant.EPOCH.atZone(UTC).toLocalDate
 
-  override def reduce(days: Int): Int = {
+  protected def doReduce(days: Int): Long = {
     val localDate = EPOCH_LOCAL_DATE.plusDays(days)
-    ChronoUnit.YEARS.between(EPOCH_LOCAL_DATE, localDate).toInt
+    ChronoUnit.YEARS.between(EPOCH_LOCAL_DATE, localDate)
   }
 }
 
-case class DaysToYearsReducer() extends DaysToYearsReducerBase {
+case class DaysToYearsReducer() extends DaysToYearsReducerBase[Int] {
   override def resultType(): DataType = IntegerType
+  override def reduce(days: Int): Int = doReduce(days).toInt
 }
 
-case class DaysToYearsReducerWithIncompatibleResultType() extends DaysToYearsReducerBase {
+case class DaysToYearsReducerWithDateResult() extends DaysToYearsReducerBase[Int] {
   override def resultType(): DataType = DateType
+  override def reduce(days: Int): Int = doReduce(days).toInt
+}
+
+case class DaysToYearsReducerWithLongResult() extends DaysToYearsReducerBase[Long] {
+  override def resultType(): DataType = LongType
+  override def reduce(days: Int): Long = doReduce(days)
 }
 
 object UnboundBucketFunction extends UnboundFunction {