fix ut of Support transform count(1) with table cache

apache · Oct 23, 2023 · f7a7fed · f7a7fed
1 parent e25a7f0
commit f7a7fed
Showing 1 changed file with 13 additions and 14 deletions.
diff --git a/backends-velox/src/test/scala/io/glutenproject/execution/VeloxColumnarCacheSuite.scala b/backends-velox/src/test/scala/io/glutenproject/execution/VeloxColumnarCacheSuite.scala
@@ -40,6 +40,7 @@ class VeloxColumnarCacheSuite extends VeloxWholeStageTransformerSuite with Adapt
     super.sparkConf
       .set("spark.shuffle.manager", "org.apache.spark.shuffle.sort.ColumnarShuffleManager")
       .set("spark.sql.shuffle.partitions", "3")
+      .set(GlutenConfig.COLUMNAR_TABLE_CACHE_ENABLED.key, "true")
   }
 
   private def checkColumnarTableCache(plan: SparkPlan): Unit = {
@@ -110,20 +111,18 @@ class VeloxColumnarCacheSuite extends VeloxWholeStageTransformerSuite with Adapt
   }
 
   test("Support transform count(1) with table cache") {
-    withSQLConf(GlutenConfig.COLUMNAR_TABLE_CACHE_ENABLED.key -> "true") {
-      val cached = spark.table("lineitem").cache()
-      try {
-        val df = spark.sql("SELECT COUNT(*) FROM lineitem")
-        checkAnswer(df, Row(60175))
-        assert(
-          find(df.queryExecution.executedPlan) {
-            case _: RowToVeloxColumnarExec => true
-            case _ => false
-          }.isEmpty
-        )
-      } finally {
-        cached.unpersist()
-      }
+    val cached = spark.table("lineitem").cache()
+    try {
+      val df = spark.sql("SELECT COUNT(*) FROM lineitem")
+      checkAnswer(df, Row(60175))
+      assert(
+        find(df.queryExecution.executedPlan) {
+          case _: RowToVeloxColumnarExec => true
+          case _ => false
+        }.isEmpty
+      )
+    } finally {
+      cached.unpersist()
     }
   }