For comments.

apache · Jun 24, 2020 · 7ceaebc · 7ceaebc
1 parent 4b9b0e8
commit 7ceaebc
Showing 1 changed file with 26 additions and 4 deletions.
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
@@ -1032,6 +1032,7 @@ class AdaptiveQueryExecSuite
         val df1 = spark.range(10).repartition($"id")
         val df2 = spark.range(10).repartition(10, $"id")
         val df3 = spark.range(10).repartition(10)
+        val df4 = spark.range(10).repartitionByRange(10, $"id".asc)
 
         val partitionsNum1 = df1.rdd.collectPartitions().length
         if (enableAQE) {
@@ -1050,6 +1051,7 @@ class AdaptiveQueryExecSuite
 
         assert(df2.rdd.collectPartitions().length == 10)
         assert(df3.rdd.collectPartitions().length == 10)
+        assert(df4.rdd.collectPartitions().length == 10)
       }
     }
   }
@@ -1064,19 +1066,39 @@ class AdaptiveQueryExecSuite
 
         val partitionsNum1 = (1 to 10).toDF.repartition($"value")
           .rdd.collectPartitions().length
-        val partitionsNum2 = (1 to 10).toDF.repartitionByRange($"value".asc)
+        val partitionsNum2 = (1 to 10).toDF.repartition($"value" + 1)
           .rdd.collectPartitions().length
-        val partitionsNum3 = (1 to 10).toDF.repartition($"value" + 1)
+
+        if (enableAQE) {
+          assert(partitionsNum1 < 10)
+          assert(partitionsNum2 < 10)
+        } else {
+          assert(partitionsNum1 === 10)
+          assert(partitionsNum2 === 10)
+        }
+      }
+    }
+  }
+
+  test("SPARK-32056 coalesce partitions for repartition by range when AQE is enabled") {
+    Seq(true, false).foreach { enableAQE =>
+      withSQLConf(
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> enableAQE.toString,
+        SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+        SQLConf.COALESCE_PARTITIONS_INITIAL_PARTITION_NUM.key -> "50",
+        SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
+
+        val partitionsNum1 = (1 to 10).toDF.repartitionByRange($"value".asc)
+          .rdd.collectPartitions().length
+        val partitionsNum2 = (1 to 10).toDF.repartitionByRange(($"value" + 1).asc)
           .rdd.collectPartitions().length
 
         if (enableAQE) {
           assert(partitionsNum1 < 10)
           assert(partitionsNum2 < 10)
-          assert(partitionsNum3 < 10)
         } else {
           assert(partitionsNum1 === 10)
           assert(partitionsNum2 === 10)
-          assert(partitionsNum3 === 10)
         }
       }
     }