(sedona) branch master updated: [SEDONA-568] Refactor TestBaseScala to allow customizable Spark configurations (#1455)

[jiayu]

This is an automated email from the ASF dual-hosted git repository.

jiayu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/sedona.git


The following commit(s) were added to refs/heads/master by this push:
     new 94cd82b51 [SEDONA-568] Refactor TestBaseScala to allow customizable 
Spark configurations (#1455)
94cd82b51 is described below

commit 94cd82b514170a4c2461765954e909c4f3768093
Author: Feng Zhang <f...@wherobots.com>
AuthorDate: Thu Jun 6 21:44:30 2024 -0700

    [SEDONA-568] Refactor TestBaseScala to allow customizable Spark 
configurations (#1455)
    
    - Added default Spark configurations in `TestBaseScala` trait.
    - Provided a method `defaultSparkConfig` for default configurations.
    - Allowed subclasses to override `sparkConfig` to add or modify Spark 
configurations.
    - Updated initialization of Spark session and Spark context to use the 
provided configurations.
    - Included comments to explain how to override configurations in subclasses.
    - Ensured that default configurations are preserved if not overridden by 
subclasses.
---
 .../org/apache/sedona/sql/TestBaseScala.scala      | 33 +++++++++++++++-------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git 
a/spark/common/src/test/scala/org/apache/sedona/sql/TestBaseScala.scala 
b/spark/common/src/test/scala/org/apache/sedona/sql/TestBaseScala.scala
index eaace3a32..891b0b3ae 100644
--- a/spark/common/src/test/scala/org/apache/sedona/sql/TestBaseScala.scala
+++ b/spark/common/src/test/scala/org/apache/sedona/sql/TestBaseScala.scala
@@ -26,7 +26,8 @@ import org.apache.sedona.common.Functions.{frechetDistance, 
hausdorffDistance}
 import org.apache.sedona.common.Predicates.dWithin
 import org.apache.sedona.common.sphere.{Haversine, Spheroid}
 import org.apache.sedona.spark.SedonaContext
-import org.apache.spark.sql.DataFrame
+import org.apache.spark.SparkContext
+import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.locationtech.jts.geom._
 import org.scalatest.{BeforeAndAfterAll, FunSpec}
 
@@ -39,16 +40,27 @@ trait TestBaseScala extends FunSpec with BeforeAndAfterAll {
   Logger.getLogger("akka").setLevel(Level.WARN)
   Logger.getLogger("org.apache.sedona.core").setLevel(Level.WARN)
 
-  val warehouseLocation = System.getProperty("user.dir") + "/target/"
-  val sparkSession = SedonaContext.builder().
-    master("local[*]").appName("sedonasqlScalaTest")
-    .config("spark.sql.warehouse.dir", warehouseLocation)
-    // We need to be explicit about broadcasting in tests.
-    .config("sedona.join.autoBroadcastJoinThreshold", "-1")
-    .config("spark.kryoserializer.buffer.max", "64m")
-    .getOrCreate()
+  // Default Spark configurations
+  def defaultSparkConfig: Map[String, String] = Map(
+    "spark.sql.warehouse.dir" -> (System.getProperty("user.dir") + "/target/"),
+    "sedona.join.autoBroadcastJoinThreshold" -> "-1",
+    "spark.kryoserializer.buffer.max" -> "64m"
+  )
 
-  val sc = sparkSession.sparkContext
+  // Method to be overridden by subclasses to provide additional configurations
+  def sparkConfig: Map[String, String] = defaultSparkConfig
+
+  // Lazy initialization of Spark session using configurations
+  lazy val sparkSession: SparkSession = {
+    val builder = SedonaContext.builder()
+      .master("local[*]")
+      .appName("sedonasqlScalaTest")
+    sparkConfig.foreach { case (key, value) => builder.config(key, value) }
+    builder.getOrCreate()
+  }
+
+  // Lazy initialization of Spark context from the Spark session
+  lazy val sc: SparkContext = sparkSession.sparkContext
 
   val resourceFolder = System.getProperty("user.dir") + "/src/test/resources/"
   val mixedWkbGeometryInputLocation = resourceFolder + "county_small_wkb.tsv"
@@ -80,6 +92,7 @@ trait TestBaseScala extends FunSpec with BeforeAndAfterAll {
   private val factory = new GeometryFactory()
 
   override def beforeAll(): Unit = {
+    super.beforeAll()
     SedonaContext.create(sparkSession)
   }