[SPARK-59054][CORE][SQL] Fix KeyGroupedPartitioner to compare partition keys by value in storage-partitioned join shuffles

This PR fixes a correctness bug in the `KeyedPartitioning` shuffle path (storage-partitioned join with `spark.sql.sources.v2.bucketing.shuffle.enabled=true`) by making the partition-key lookup in `KeyGroupedPartitioner` use the exact same equivalence that grouped the partition keys.

- In the `val part` match inside `ShuffleExchangeExec.prepareShuffleDependency`, the driver-side `valueMap` now uses the partitioning's own `InternalRowComparableWrapper`s as keys (`k.partitionKeys.zipWithIndex.toMap`), instead of `Seq[Any]` from `InternalRow.toSeq`.
- In the local `getPartitionKeyExtractor` of that same method, the executor-side lookup key is built by evaluating the bound partition expressions into a reused `GenericInternalRow` and wrapping it with `InternalRowComparableWrapper.getInternalRowComparableWrapperFactory`, instead of `Seq[Any]` from per-row `eval`.
- `InternalRowComparableWrapper` is now `Serializable`: `structType` and `ordering` cannot cross the wire (the ordering may be generated code), so they are `transient` and re-derived from the shared caches on first use after deserialization. This is what allows the wrappers to be shipped to executors inside the partitioner.
- `KeyGroupedPartitioner` (core) takes `Map[Any, Int]` and looks keys up with `getOrElse`, removing the `ArraySeq` normalization and the `getOrElseUpdate` map mutation. Since lookup keys are no longer retained by the map, the executor side can reuse a single row per task without per-record copies.

Because both the map keys and the lookup keys are wrappers over the same data types, lookups share the `RowOrdering` equivalence (binary keys by content, `-0.0 == 0.0`, NaNs equal, collation-aware strings) that grouped and de-duplicated `partitionKeys` on the driver -- by construction, not by approximation.

The `valueMap` keys and the executor-side lookup keys were `Seq[Any]` compared with Scala `==` element equality, while partition-key grouping/de-duplication uses `InternalRowComparableWrapper` (`RowOrdering`) semantics. The two disagree for `BinaryType`: `Array[Byte]` elements are compared by reference, so every lookup misses and falls back to `nonNegativeMod(hashCode, numPartitions)` -- effectively a random partition per row, since `Array` hash codes are identity-based.

As a result, when the non-keyed side of a storage-partitioned join is shuffled into a table partitioned by e.g. `identity(binary_col)`, rows land in partitions that do not match the keyed side, and the join silently drops matches (wrong results, no error).

An earlier revision of this PR used `UnsafeRow`s produced by identical `UnsafeProjection`s as the shared key representation. As pointed out in review, byte equality is strictly finer than the `RowOrdering` equality that grouped `partitionKeys`: a transform whose `resultType()` is floating point can produce `-0.0` and `0.0` partition keys, which the driver collapses into one partition but byte comparison splits, reintroducing the same lost-match failure (and regressing a case the `Seq[Any]` code handled). Using the wrappers themselves eliminates the mismatch by construction.

Note that this bug also exists in the released 4.0, 4.1 and 4.2 lines, which carry the same `KeyGroupedPartitioner` lookup.

Yes, it is a bug fix. Previously, a storage-partitioned join with a `BinaryType` partition key and `spark.sql.sources.v2.bucketing.shuffle.enabled=true` could silently return fewer rows than expected. Now it returns correct results.

Two new regression tests in `KeyGroupedPartitioningSuite`, each exercising both `V2_BUCKETING_SHUFFLE_ENABLED` states (1 shuffle when on, 2 shuffles when off, same results either way):

- `SPARK-59054: shuffle one side: partition keys with binary type`: joins a v2 table partitioned by `identity` on a binary column with an unpartitioned table. Confirmed to fail before the fix (missing join rows) and pass after it.
- `SPARK-59054: shuffle one side: partition transform collapsing -0.0 and 0.0`: uses a new test connector function `signed_zeros` whose `resultType()` (`DoubleType`) differs from its `inputTypes()` (`LongType`), mapping ids 1 and 2 to `-0.0` and `0.0`. Confirmed to fail against the earlier `UnsafeRow`-based revision (one match lost) and pass with the wrapper-based fix. `InMemoryBaseTable` gained the matching transform whitelist entry and `getKey` case.

The full `KeyGroupedPartitioningSuite` (106 tests) passes.

A NaN-key reproduction through ordinary float/double join keys is not reachable, because `NormalizeFloatingNumbers` wraps such join keys in `KnownFloatingPointNormalized(NormalizeNaNAndZero(...))`, which prevents SPJ from triggering; the wrapper-based comparison handles NaN keys correctly regardless.

Generated-by: Claude Fable 5

Closes apache#58345 from dongjoon-hyun/SPARK-59054.

Authored-by: Dongjoon Hyun <dongjoon@apache.org>
Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
(cherry picked from commit 74b5db8)
Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
(cherry picked from commit b46334b)

(cherry picked from commit 6686830)

Author: dongjoon-hyun

core/src/main/scala/org/apache/spark/Partitioner.scala

@@ -19,7 +19,6 @@ package org.apache.spark
 
 import java.io.{IOException, ObjectInputStream, ObjectOutputStream}
 
-import scala.collection.immutable.ArraySeq
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 import scala.math.log10
@@ -144,16 +143,16 @@ private[spark] class PartitionIdPassthrough(override val numPartitions: Int) ext
  * The `valueMap` is a map that contains tuples of (partition value, partition id). It is generated
  * by [[org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning]], used to partition
  * the other side of a join to make sure records with same partition value are in the same
- * partition.
+ * partition. Keys are looked up with `equals`/`hashCode`, so the caller must supply the map keys
+ * and the per-record lookup keys in a single representation that compares partition values
+ * consistently (the caller uses the partitioning's own `InternalRowComparableWrapper`s). Keys
+ * absent from the map fall back to a partition derived from the key's hash code.
  */
 private[spark] class KeyGroupedPartitioner(
-    valueMap: mutable.Map[Seq[Any], Int],
+    valueMap: Map[Any, Int],
     override val numPartitions: Int) extends Partitioner {
   override def getPartition(key: Any): Int = {
-    val keys = key.asInstanceOf[Seq[Any]]
-    val normalizedKeys = ArraySeq.from(keys)
-    valueMap.getOrElseUpdate(normalizedKeys,
-      Utils.nonNegativeMod(normalizedKeys.hashCode, numPartitions))
+    valueMap.getOrElse(key, Utils.nonNegativeMod(key.hashCode, numPartitions))
   }
 }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapper.scala

@@ -20,7 +20,7 @@ package org.apache.spark.sql.catalyst.util
 import scala.collection.mutable
 
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Expression, Murmur3HashFunction, RowOrdering}
+import org.apache.spark.sql.catalyst.expressions.{BaseOrdering, Expression, Murmur3HashFunction, RowOrdering}
 import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition}
 import org.apache.spark.sql.types.{DataType, StructField, StructType}
 import org.apache.spark.util.NonFateSharingCache
@@ -33,11 +33,27 @@ import org.apache.spark.util.NonFateSharingCache
  *
  * @param dataTypes the data types for the row
  */
-class InternalRowComparableWrapper(val row: InternalRow, val dataTypes: Seq[DataType]) {
-  import InternalRowComparableWrapper._
+class InternalRowComparableWrapper(val row: InternalRow, val dataTypes: Seq[DataType])
+  extends Serializable {
 
-  private val structType = structTypeCache.get(dataTypes)
-  private val ordering = orderingCache.get(dataTypes)
+  // `structType` and `ordering` cannot cross the wire (the ordering may be generated code), so
+  // they are transient and re-derived from the shared caches on first use after deserialization.
+  @transient private var _structType: StructType = _
+  @transient private var _ordering: BaseOrdering = _
+
+  def structType: StructType = {
+    if (_structType == null) {
+      _structType = InternalRowComparableWrapper.structTypeCache.get(dataTypes)
+    }
+    _structType
+  }
+
+  def ordering: BaseOrdering = {
+    if (_ordering == null) {
+      _ordering = InternalRowComparableWrapper.orderingCache.get(dataTypes)
+    }
+    _ordering
+  }
 
   override def hashCode(): Int = Murmur3HashFunction.hash(
     row,

sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala

@@ -109,6 +109,7 @@ abstract class InMemoryBaseTable(
     case _: SortedBucketTransform =>
     case _: ClusterByTransform =>
     case NamedTransform("truncate", Seq(_: NamedReference, _: Literal[_])) =>
+    case NamedTransform("signed_zeros", Seq(_: NamedReference)) =>
     case t if !allowUnsupportedTransforms =>
       throw new IllegalArgumentException(s"Transform $t is not a supported transform")
   }
@@ -221,6 +222,15 @@ abstract class InMemoryBaseTable(
           case (v, t) =>
             throw new IllegalArgumentException(s"Match: unsupported argument(s) type - ($v, $t)")
         }
+      // the result should be consistent with SignedZerosFunction defined at
+      // transformFunctions.scala
+      case NamedTransform("signed_zeros", Seq(ref: NamedReference)) =>
+        extractor(ref.fieldNames, cleanedSchema, row) match {
+          case (value: Long, LongType) =>
+            if (value == 1L) -0.0d else if (value == 2L) 0.0d else value.toDouble
+          case (v, t) =>
+            throw new IllegalArgumentException(s"Match: unsupported argument(s) type - ($v, $t)")
+        }
       case ClusterByTransform(columnNames) =>
         columnNames.map { colName =>
           extractor(colName.fieldNames, cleanedSchema, row)._1

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

@@ -20,7 +20,6 @@ package org.apache.spark.sql.execution.exchange
 import java.util.concurrent.atomic.AtomicReference
 import java.util.function.Supplier
 
-import scala.collection.mutable
 import scala.concurrent.{ExecutionContext, Future, Promise}
 
 import org.apache.spark._
@@ -30,12 +29,13 @@ import org.apache.spark.serializer.Serializer
 import org.apache.spark.shuffle.{ShuffleWriteMetricsReporter, ShuffleWriteProcessor}
 import org.apache.spark.shuffle.sort.SortShuffleManager
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, GenericInternalRow, UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.catalyst.expressions.BindReferences.bindReferences
 import org.apache.spark.sql.catalyst.expressions.codegen.LazilyGeneratedOrdering
 import org.apache.spark.sql.catalyst.plans.logical.Statistics
 import org.apache.spark.sql.catalyst.plans.physical._
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter}
 import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf}
@@ -367,10 +367,14 @@ object ShuffleExchangeExec {
           samplePointsPerPartitionHint = SQLConf.get.rangeExchangeSampleSizePerPartition)
       case SinglePartition => new ConstantPartitioner
       case k @ KeyGroupedPartitioning(expressions, n, _, _, _, _) =>
-        val valueMap = k.uniquePartitionValues.zipWithIndex.map {
-          case (partition, index) => (partition.toSeq(expressions.map(_.dataType)), index)
-        }.toMap
-        new KeyGroupedPartitioner(mutable.Map(valueMap.toSeq: _*), n)
+        // The map keys and the lookup keys produced by `getPartitionKeyExtractor` below are
+        // wrapped over this partitioning's expression data types, so map lookups share the exact
+        // equivalence (`RowOrdering`, e.g. binary keys by content, -0.0 == 0.0, NaNs equal) that
+        // grouped and de-duplicated the partition values.
+        val dataTypes = expressions.map(_.dataType)
+        val valueMap = k.uniquePartitionValues
+          .map(new InternalRowComparableWrapper(_, dataTypes)).zipWithIndex.toMap[Any, Int]
+        new KeyGroupedPartitioner(valueMap, n)
       case _ => throw SparkException.internalError(s"Exchange not implemented for $newPartitioning")
       // TODO: Handle BroadcastPartitioning.
     }
@@ -398,7 +402,20 @@ object ShuffleExchangeExec {
         row => projection(row)
       case SinglePartition => identity
       case KeyGroupedPartitioning(expressions, _, _, _, _, _) =>
-        row => bindReferences(expressions, outputAttributes).map(_.eval(row))
+        // Wrap the evaluated partition key so it compares equal to the KeyGroupedPartitioner's
+        // map keys (wrappers over the same data types) under `RowOrdering` semantics. The wrapped
+        // row is reused across records; KeyGroupedPartitioner does not retain lookup keys.
+        val boundExpressions = bindReferences(expressions, outputAttributes).toArray
+        val dataTypes = expressions.map(_.dataType)
+        val partitionKeyRow = new GenericInternalRow(boundExpressions.length)
+        row => {
+          var i = 0
+          while (i < boundExpressions.length) {
+            partitionKeyRow.update(i, boundExpressions(i).eval(row))
+            i += 1
+          }
+          new InternalRowComparableWrapper(partitionKeyRow, dataTypes)
+        }
       case _ => throw SparkException.internalError(s"Exchange not implemented for $newPartitioning")
     }
 

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

@@ -44,7 +44,8 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
     UnboundYearsFunction,
     UnboundDaysFunction,
     UnboundBucketFunction,
-    UnboundTruncateFunction)
+    UnboundTruncateFunction,
+    UnboundSignedZerosFunction)
 
   override def sparkConf: SparkConf = super.sparkConf
     .set(V2_BUCKETING_ENABLED, true)
@@ -1405,6 +1406,132 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
     }
   }
 
+  test("SPARK-59054: shuffle one side: partition keys with binary type") {
+    val items_partitions = Array(identity("id"))
+    createTable(items, Array(
+      Column.create("id", BinaryType),
+      Column.create("name", StringType),
+      Column.create("price", DoubleType)), items_partitions)
+
+    sql(s"INSERT INTO testcat.ns.$items VALUES " +
+      "(X'0101', 'aa', 40.0), " +
+      "(X'0202', 'bb', 10.0), " +
+      "(X'0303', 'cc', 15.5), " +
+      "(X'0404', 'dd', 20.0)")
+
+    createTable(purchases, Array(
+      Column.create("item_id", BinaryType),
+      Column.create("price", DoubleType)), Array.empty)
+    sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+      "(X'0101', 42.0), (X'0101', 44.0), (X'0202', 11.0), (X'0202', 19.5), " +
+      "(X'0303', 26.0), (X'0303', 30.0), (X'0404', 50.0), (X'0404', 60.0)")
+
+    Seq(true, false).foreach { shuffle =>
+      withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> shuffle.toString) {
+        val df = createJoinTestDF(Seq("id" -> "item_id"))
+        val shuffles = collectShuffles(df.queryExecution.executedPlan)
+        if (shuffle) {
+          assert(shuffles.size == 1, "only shuffle one side not report partitioning")
+        } else {
+          assert(shuffles.size == 2, "should add two side shuffle when bucketing shuffle one " +
+            "side is not enabled")
+        }
+
+        checkAnswer(df, Seq(
+          Row(Array[Byte](1, 1), "aa", 40.0, 42.0),
+          Row(Array[Byte](1, 1), "aa", 40.0, 44.0),
+          Row(Array[Byte](2, 2), "bb", 10.0, 11.0),
+          Row(Array[Byte](2, 2), "bb", 10.0, 19.5),
+          Row(Array[Byte](3, 3), "cc", 15.5, 26.0),
+          Row(Array[Byte](3, 3), "cc", 15.5, 30.0),
+          Row(Array[Byte](4, 4), "dd", 20.0, 50.0),
+          Row(Array[Byte](4, 4), "dd", 20.0, 60.0)))
+      }
+    }
+  }
+
+  test("SPARK-59054: shuffle one side: struct partition keys with different field names") {
+    // Struct equality ignores field names, so joining STRUCT<a:INT> with STRUCT<b:INT> is legal
+    // and SPJ stays eligible. The shuffled side's lookup keys carry its own schema while the
+    // partitioner's map keys come from the keyed side, so key comparison must not depend on
+    // the field names.
+    val items_partitions = Array(identity("id"))
+    createTable(items, Array(
+      Column.create("id", new StructType().add("a", IntegerType)),
+      Column.create("name", StringType),
+      Column.create("price", DoubleType)), items_partitions)
+
+    sql(s"INSERT INTO testcat.ns.$items VALUES " +
+      "(named_struct('a', 1), 'aa', 40.0), " +
+      "(named_struct('a', 2), 'bb', 10.0), " +
+      "(named_struct('a', 3), 'cc', 15.5), " +
+      "(named_struct('a', 4), 'dd', 20.0)")
+
+    createTable(purchases, Array(
+      Column.create("item_id", new StructType().add("b", IntegerType)),
+      Column.create("price", DoubleType)), Array.empty)
+    sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+      "(named_struct('b', 1), 42.0), (named_struct('b', 2), 19.5), " +
+      "(named_struct('b', 3), 26.0), (named_struct('b', 4), 50.0)")
+
+    Seq(true, false).foreach { shuffle =>
+      withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> shuffle.toString) {
+        val df = createJoinTestDF(Seq("id" -> "item_id"))
+        val shuffles = collectShuffles(df.queryExecution.executedPlan)
+        if (shuffle) {
+          assert(shuffles.size == 1, "only shuffle one side not report partitioning")
+        } else {
+          assert(shuffles.size == 2, "should add two side shuffle when bucketing shuffle one " +
+            "side is not enabled")
+        }
+
+        checkAnswer(df, Seq(
+          Row(Row(1), "aa", 40.0, 42.0),
+          Row(Row(2), "bb", 10.0, 19.5),
+          Row(Row(3), "cc", 15.5, 26.0),
+          Row(Row(4), "dd", 20.0, 50.0)))
+      }
+    }
+  }
+
+  test("SPARK-59054: shuffle one side: partition transform collapsing -0.0 and 0.0") {
+    // `signed_zeros` maps id 1 to -0.0 and id 2 to 0.0: two partition keys that are equal
+    // under SQL semantics but distinct bitwise, which the grouped side collapses into one
+    // partition. Rows of both forms on the shuffled side must land in that partition.
+    val items_partitions = Array(
+      Expressions.apply("signed_zeros", Expressions.column("id")))
+    createTable(items, itemsColumns, items_partitions)
+
+    sql(s"INSERT INTO testcat.ns.$items VALUES " +
+      "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+      "(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
+      "(3, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+
+    createTable(purchases, purchasesColumns, Array.empty)
+    sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+      "(1, 42.0, cast('2020-01-01' as timestamp)), " +
+      "(2, 19.5, cast('2020-02-01' as timestamp)), " +
+      "(3, 26.0, cast('2023-01-01' as timestamp))")
+
+    Seq(true, false).foreach { shuffle =>
+      withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> shuffle.toString) {
+        val df = createJoinTestDF(Seq("id" -> "item_id"))
+        val shuffles = collectShuffles(df.queryExecution.executedPlan)
+        if (shuffle) {
+          assert(shuffles.size == 1, "only shuffle one side not report partitioning")
+        } else {
+          assert(shuffles.size == 2, "should add two side shuffle when bucketing shuffle one " +
+            "side is not enabled")
+        }
+
+        checkAnswer(df, Seq(
+          Row(1, "aa", 40.0, 42.0),
+          Row(2, "bb", 10.0, 19.5),
+          Row(3, "cc", 15.5, 26.0)))
+      }
+    }
+  }
+
   test("SPARK-44641: duplicated records when SPJ is not triggered") {
     val items_partitions = Array(bucket(8, "id"))
     createTable(items, itemsColumns, items_partitions)

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

@@ -133,6 +133,32 @@ object StringSelfFunction extends ScalarFunction[UTF8String] {
   }
 }
 
+object UnboundSignedZerosFunction extends UnboundFunction {
+  override def bind(inputType: StructType): BoundFunction = {
+    if (inputType.size == 1 && inputType.head.dataType == LongType) SignedZerosFunction
+    else throw new UnsupportedOperationException("'signed_zeros' only takes a long as input type")
+  }
+  override def description(): String = name()
+  override def name(): String = "signed_zeros"
+}
+
+// A transform whose result type (DoubleType) differs from its input type (LongType), mapping
+// 1 to -0.0 and 2 to 0.0 so that two distinct inputs produce partition keys that are equal
+// under SQL semantics but distinct bitwise. The result should be consistent with the
+// "signed_zeros" NamedTransform defined at InMemoryBaseTable.scala.
+object SignedZerosFunction extends ScalarFunction[Double] {
+  override def inputTypes(): Array[DataType] = Array(LongType)
+  override def resultType(): DataType = DoubleType
+  override def name(): String = "signed_zeros"
+  override def canonicalName(): String = name()
+  override def toString: String = name()
+  override def produceResult(input: InternalRow): Double = input.getLong(0) match {
+    case 1L => -0.0d
+    case 2L => 0.0d
+    case v => v.toDouble
+  }
+}
+
 object UnboundTruncateFunction extends UnboundFunction {
   override def bind(inputType: StructType): BoundFunction = TruncateFunction
   override def description(): String = name()