Return KeyGroupedPartitioning

G-Research · Oct 9, 2023 · 84fb949 · 84fb949
1 parent 10b19c8
commit 84fb949
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/...c/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala b/...c/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala
@@ -21,7 +21,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Expression, RowOrdering, SortOrder}
 import org.apache.spark.sql.catalyst.plans.physical
-import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, KeyGroupedPartitioning, SinglePartition}
+import org.apache.spark.sql.catalyst.plans.physical.{KeyGroupedPartitioning, SinglePartition}
 import org.apache.spark.sql.catalyst.util.{truncatedString, InternalRowComparableWrapper}
 import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, PartitionReaderFactory, Scan}
 import org.apache.spark.sql.execution.{ExplainUtils, LeafExecNode, SQLExecution}
@@ -99,8 +99,9 @@ trait DataSourceV2ScanExecBase extends LeafExecNode {
             val keyGroupedParts = keyGroupedPartsInfo.groupedParts
             KeyGroupedPartitioning(exprs, keyGroupedParts.size, keyGroupedParts.map(_.value),
               keyGroupedPartsInfo.originalParts.map(_.partitionKey()))
+          }.getOrElse {
+            KeyGroupedPartitioning(exprs, inputPartitions.size)
           }
-          .getOrElse(HashPartitioning(exprs, inputPartitions.size))
       case _ =>
         super.outputPartitioning
     }

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
@@ -105,7 +105,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
       val distribution = physical.ClusteredDistribution(
         Seq(TransformExpression(BucketFunction, Seq(attr("ts")), Some(32))))
 
-      checkQueryPlan(df, distribution, physical.HashPartitioning(distribution.clustering, 32))
+      checkQueryPlan(df, distribution, physical.KeyGroupedPartitioning(distribution.clustering, 32))
     }
   }
 
@@ -175,7 +175,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
       val distribution = physical.ClusteredDistribution(
         Seq(TransformExpression(BucketFunction, Seq(attr("ts")), Some(32))))
 
-      checkQueryPlan(df, distribution, physical.HashPartitioning(distribution.clustering, 32))
+      checkQueryPlan(df, distribution, physical.KeyGroupedPartitioning(distribution.clustering, 32))
     }
   }