add spark session for remote fs

apache · Mar 21, 2024 · 17f73fe · 17f73fe
1 parent 14d94d0
commit 17f73fe
Showing 1 changed file with 44 additions and 0 deletions.
diff --git a/.../glutenproject/execution/GlutenClickHouseMergeTreeWriteOnObjectStorageAbstractSuite.scala b/.../glutenproject/execution/GlutenClickHouseMergeTreeWriteOnObjectStorageAbstractSuite.scala
@@ -16,15 +16,25 @@
  */
 package io.glutenproject.execution
 
+import io.glutenproject.GlutenConfig
+
+import org.apache.spark.sql.SparkSession
+
 import _root_.org.apache.spark.{SPARK_VERSION_SHORT, SparkConf}
 import _root_.org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.commons.io.FileUtils
+
+import java.io.File
 
 // Some sqls' line length exceeds 100
 // scalastyle:off line.size.limit
 
 class GlutenClickHouseMergeTreeWriteOnObjectStorageAbstractSuite
   extends GlutenClickHouseTPCHAbstractSuite
   with AdaptiveSparkPlanHelper {
+  private var _spark: SparkSession = _
+
+  override protected def spark: SparkSession = _spark
 
   override protected val needCopyParquetToTablePath = true
 
@@ -52,8 +62,19 @@ class GlutenClickHouseMergeTreeWriteOnObjectStorageAbstractSuite
   val S3_ACCESS_KEY = "BypTYzcXOlfr03FFIvt4"
   val S3_SECRET_KEY = "K9MDaGItPSaphorZM8t4hXf30gHF9dBWi6L2dK5E"
 
+  override protected def initializeSession(): Unit = {
+    if (_spark == null) {
+      _spark = SparkSession
+        .builder()
+        .appName("Gluten-UT-RemoteHS")
+        .config(sparkConf)
+        .getOrCreate()
+    }
+  }
+
   override protected def sparkConf: SparkConf = {
     super.sparkConf
+      .setMaster("local[2]")
       .set("spark.shuffle.manager", "org.apache.spark.shuffle.sort.ColumnarShuffleManager")
       .set("spark.io.compression.codec", "LZ4")
       .set("spark.sql.shuffle.partitions", "5")
@@ -140,5 +161,28 @@ class GlutenClickHouseMergeTreeWriteOnObjectStorageAbstractSuite
     createNotNullTPCHTablesInParquet(tablesPath)
   }
 
+  override protected def afterAll(): Unit = {
+    try {
+      super.afterAll()
+    } finally {
+      try {
+        if (_spark != null) {
+          try {
+            _spark.sessionState.catalog.reset()
+          } finally {
+            _spark.stop()
+            _spark = null
+          }
+        }
+      } finally {
+        SparkSession.clearActiveSession()
+        SparkSession.clearDefaultSession()
+      }
+    }
+
+    FileUtils.forceDelete(new File(basePath))
+    // init GlutenConfig in the next beforeAll
+    GlutenConfig.ins = null
+  }
 }
 // scalastyle:off line.size.limit