Fix bug of using un-projected expression in sorting

szehon-ho · Feb 27, 2024 · 0016169 · 0016169
1 parent 0754a23
commit 0016169
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala
@@ -172,7 +172,7 @@ case class BatchScanExec(
                 KeyGroupedShuffleSpec.reducePartitionValue(row, partExpressions, reducers)
               }.map { case (wrapper, splits) => (wrapper.row, splits.flatMap(_._2)) }.toSeq
               val rowOrdering = RowOrdering.createNaturalAscendingOrdering(
-                expressions.map(_.dataType))
+                partExpressions.map(_.dataType))
               result.sorted(rowOrdering.on((t: (InternalRow, _)) => t._1))
             case _ => groupedPartitions
           }

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -69,7 +69,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
       .add("data", StringType)
       .add("ts", TimestampType)
     private val schema2 = new StructType()
-      .add("store_id", IntegerType)
+      .add("store_id", LongType)
       .add("dept_id", IntegerType)
       .add("data", StringType)
 
@@ -1482,7 +1482,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
       case (table1buckets, table2buckets) =>
         catalog.clearTables()
 
-        val partition1 = Array(bucket(3, "store_id"),
+        val partition1 = Array(identity("data"),
           bucket(table1buckets, "dept_id"))
         val partition2 = Array(bucket(3, "store_id"),
           bucket(table2buckets, "dept_id"))