Refine test

apache · May 6, 2024 · 4e4d528 · 4e4d528
1 parent d62262a
commit 4e4d528
Showing 1 changed file with 15 additions and 9 deletions.
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
@@ -67,15 +67,21 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
   setupTestData()
 
   test("Disable Comet shuffle with AQE coalesce partitions enabled") {
-    withSQLConf(
-      CometConf.COMET_EXEC_ENABLED.key -> "true",
-      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
-      SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
-      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
-      val df = sql(
-        "SELECT * FROM (SELECT * FROM testData WHERE key = 0) t1 FULL JOIN " +
-          "testData2 t2 ON t1.key = t2.a")
-      checkShuffleAnswer(df, 0)
+    Seq(true, false).foreach { coalescePartitionsEnabled =>
+      withSQLConf(
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+        SQLConf.COALESCE_PARTITIONS_ENABLED.key -> coalescePartitionsEnabled.toString,
+        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+        val df = sql(
+          "SELECT * FROM (SELECT * FROM testData WHERE key = 0) t1 FULL JOIN " +
+            "testData2 t2 ON t1.key = t2.a")
+        if (coalescePartitionsEnabled) {
+          checkShuffleAnswer(df, 0)
+        } else {
+          checkShuffleAnswer(df, 2)
+        }
+      }
     }
   }