[SPARK-50496][Core]Change partitioning to SinglePartition if partitio…

…n number is 1
guixiaowen · Dec 6, 2024 · 04705aa · 04705aa
1 parent af4f37c
commit 04705aa
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/...st/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala b/...st/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala
@@ -1980,7 +1980,11 @@ case class RebalancePartitions(
   override def output: Seq[Attribute] = child.output
   override val nodePatterns: Seq[TreePattern] = Seq(REBALANCE_PARTITIONS)
 
-  override val partitioning: Partitioning = super.partitioning
+  override val partitioning: Partitioning = if (numPartitions == 1) {
+    SinglePartition
+  } else {
+    super.partitioning
+  }
 
   override protected def withNewChildInternal(newChild: LogicalPlan): RebalancePartitions =
     copy(child = newChild)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -32,7 +32,8 @@ import org.apache.spark.sql.{DataFrame, Dataset, QueryTest, Row, SparkSession, S
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
-import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, RepartitionOperation}
+import org.apache.spark.sql.catalyst.plans.physical.SinglePartition
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
 import org.apache.spark.sql.execution.columnar.{InMemoryTableScanExec, InMemoryTableScanLike}
@@ -2236,6 +2237,25 @@ class AdaptiveQueryExecSuite
     }
   }
 
+  test("SPARK-50496: Change rebalance partitioning to SinglePartition if partition number is 1") {
+    def checkSinglePartitioning(df: DataFrame): Unit = {
+      assert(
+        df.queryExecution.analyzed.collect {
+          case r: RepartitionOperation => r
+        }.size == 1)
+
+      assert(
+        collect(df.queryExecution.executedPlan) {
+          case s: ShuffleExchangeExec if s.outputPartitioning == SinglePartition => s
+        }.size == 1)
+    }
+
+    checkSinglePartitioning(sql("SELECT /*+ REBALANCE(1) */ * FROM VALUES(1),(2),(3) AS t(c)"))
+    checkSinglePartitioning(sql("SELECT /*+ REBALANCE(1, c) */ * FROM VALUES(1),(2),(3) AS t(c)"))
+  }
+
+
+
   test("SPARK-35725: Support optimize skewed partitions in RebalancePartitions") {
     withTempView("v") {
       withSQLConf(