[SPARK-55535][SPARK-55092][SQL] Refactor KeyGroupedPartitioning and Storage Partition Join

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

@@ -142,7 +142,7 @@ private[spark] class PartitionIdPassthrough(override val numPartitions: Int) ext
 /**
  * A [[org.apache.spark.Partitioner]] that partitions all records using partition value map.
  * 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
+ * by [[org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning]], used to partition
  * the other side of a join to make sure records with same partition value are in the same
  * partition.
  */

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

@@ -17,8 +17,6 @@
 
 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.{BaseOrdering, Expression, Murmur3HashFunction, RowOrdering}
 import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition}
@@ -101,31 +99,6 @@ object InternalRowComparableWrapper {
     new InternalRowComparableWrapper(partitionRow, partitionExpression.map(_.dataType))
   }
 
-  def mergePartitions(
-      leftPartitioning: Seq[InternalRow],
-      rightPartitioning: Seq[InternalRow],
-      partitionExpression: Seq[Expression],
-      intersect: Boolean = false): Seq[InternalRowComparableWrapper] = {
-    val partitionDataTypes = partitionExpression.map(_.dataType)
-    val leftPartitionSet = new mutable.HashSet[InternalRowComparableWrapper]
-    val internalRowComparableWrapperFactory =
-      getInternalRowComparableWrapperFactory(partitionDataTypes)
-    leftPartitioning
-      .map(internalRowComparableWrapperFactory)
-      .foreach(partition => leftPartitionSet.add(partition))
-    val rightPartitionSet = new mutable.HashSet[InternalRowComparableWrapper]
-    rightPartitioning
-      .map(internalRowComparableWrapperFactory)
-      .foreach(partition => rightPartitionSet.add(partition))
-
-    val result = if (intersect) {
-      leftPartitionSet.intersect(rightPartitionSet)
-    } else {
-      leftPartitionSet.union(rightPartitionSet)
-    }
-    result.toSeq
-  }
-
   /** Creates a shared factory method for a given row schema to avoid excessive cache lookups. */
   def getInternalRowComparableWrapperFactory(
       dataTypes: Seq[DataType]): InternalRow => InternalRowComparableWrapper = {

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapperBenchmark.scala

@@ -17,9 +17,9 @@
 
 package org.apache.spark.sql.catalyst.util
 
+import scala.collection.mutable
+
 import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
-import org.apache.spark.sql.catalyst.expressions.Literal
-import org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning
 import org.apache.spark.sql.connector.catalog.PartitionInternalRow
 import org.apache.spark.sql.types.IntegerType
 
@@ -41,30 +41,28 @@ object InternalRowComparableWrapperBenchmark extends BenchmarkBase {
     val partitionNum = 200_000
     val bucketNum = 4096
     val day = 20240401
-    val partitions = (0 until partitionNum).map { i =>
+    val partitionKeys = (0 until partitionNum).map { i =>
       val bucketId = i % bucketNum
       PartitionInternalRow.apply(Array(day, bucketId));
     }
     val benchmark = new Benchmark("internal row comparable wrapper", partitionNum, output = output)
 
+    val comparableKeyWrapperFactory =
+      InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(
+        Seq(IntegerType, IntegerType))
+    val comparablePartitionKeys = partitionKeys.map(comparableKeyWrapperFactory)
+
     benchmark.addCase("toSet") { _ =>
-      val internalRowComparableWrapperFactory =
-        InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(
-          Seq(IntegerType, IntegerType))
-      val distinct = partitions
-        .map(internalRowComparableWrapperFactory)
-        .toSet
+      val distinct = comparablePartitionKeys.toSet
+
       assert(distinct.size == bucketNum)
     }
 
     benchmark.addCase("mergePartitions") { _ =>
-      // just to mock the data types
-      val expressions = (Seq(Literal(day, IntegerType), Literal(0, IntegerType)))
+      val leftKeySet = mutable.HashSet.from(comparablePartitionKeys)
+      val rightKeySet = mutable.HashSet.from(comparablePartitionKeys)
+      val merged = leftKeySet.union(rightKeySet)
 
-      val leftPartitioning = KeyGroupedPartitioning(expressions, bucketNum, partitions)
-      val rightPartitioning = KeyGroupedPartitioning(expressions, bucketNum, partitions)
-      val merged = InternalRowComparableWrapper.mergePartitions(
-        leftPartitioning.partitionValues, rightPartitioning.partitionValues, expressions)
       assert(merged.size == bucketNum)
     }