SPARK_HOME detection supports Spark 4

apache · May 23, 2024 · 20e71fd · 20e71fd
1 parent 4285abc
commit 20e71fd
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 28 deletions.
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilder.scala
@@ -117,11 +117,11 @@ class SparkProcessBuilder(
   }
 
   override protected lazy val engineHomeDirFilter: FileFilter = file => {
-    val r = SCALA_COMPILE_VERSION match {
-      case "2.12" => SPARK_HOME_REGEX_SCALA_212
-      case "2.13" => SPARK_HOME_REGEX_SCALA_213
+    val patterns = SCALA_COMPILE_VERSION match {
+      case "2.12" => Seq(SPARK3_HOME_REGEX_SCALA_212)
+      case "2.13" => Seq(SPARK3_HOME_REGEX_SCALA_213, SPARK4_HOME_REGEX_SCALA_213)
     }
-    file.isDirectory && r.findFirstMatchIn(file.getName).isDefined
+    file.isDirectory && patterns.exists(_.findFirstMatchIn(file.getName).isDefined)
   }
 
   override protected[kyuubi] lazy val commands: Iterable[String] = {
@@ -364,11 +364,14 @@ object SparkProcessBuilder {
   final private val SPARK_CONF_FILE_NAME = "spark-defaults.conf"
 
   final private[kyuubi] val SPARK_CORE_SCALA_VERSION_REGEX =
-    """^spark-core_(\d\.\d+).*.jar$""".r
+    """^spark-core_(\d\.\d+)-.*\.jar$""".r
 
-  final private[kyuubi] val SPARK_HOME_REGEX_SCALA_212 =
-    """^spark-\d+\.\d+\.\d+-bin-hadoop\d+(\.\d+)?$""".r
+  final private[kyuubi] val SPARK3_HOME_REGEX_SCALA_212 =
+    """^spark-3\.\d+\.\d+-bin-hadoop\d+(\.\d+)?$""".r
 
-  final private[kyuubi] val SPARK_HOME_REGEX_SCALA_213 =
-    """^spark-\d+\.\d+\.\d+-bin-hadoop\d(\.\d+)?+-scala\d+(\.\d+)?$""".r
+  final private[kyuubi] val SPARK3_HOME_REGEX_SCALA_213 =
+    """^spark-3\.\d+\.\d+-bin-hadoop\d(\.\d+)?+-scala2\.13$""".r
+
+  final private[kyuubi] val SPARK4_HOME_REGEX_SCALA_213 =
+    """^spark-4\.\d+\.\d+(-\w*)?-bin-hadoop\d(\.\d+)?+$""".r
 }
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/engine/spark/SparkProcessBuilderSuite.scala
@@ -403,36 +403,43 @@ class SparkProcessBuilderSuite extends KerberizedTestHelper with MockitoSugar {
       "spark-core_2.13-3.5.0-abc-20230921.jar",
       "spark-core_2.13-3.5.0-xyz-1.2.3.jar",
       "spark-core_2.13-3.5.0.1.jar",
-      "spark-core_2.13.2-3.5.0.jar").foreach { f =>
+      "spark-core_2.13-4.0.0-preview1.jar",
+      "spark-core_2.13-4.0.0.jar").foreach { f =>
       assertResult("2.13")(builder.extractSparkCoreScalaVersion(Seq(f)))
     }
 
     Seq(
       "spark-dummy_2.13-3.5.0.jar",
       "spark-core_2.13-3.5.0.1.zip",
-      "yummy-spark-core_2.13-3.5.0.jar").foreach { f =>
+      "yummy-spark-core_2.13-3.5.0.jar",
+      "spark-core_2.13.2-3.5.0.jar").foreach { f =>
       assertThrows[KyuubiException](builder.extractSparkCoreScalaVersion(Seq(f)))
     }
   }
 
   test("match scala version of spark home") {
-    SCALA_COMPILE_VERSION match {
-      case "2.12" => Seq(
-          "spark-3.2.4-bin-hadoop3.2",
-          "spark-3.2.4-bin-hadoop2.7",
-          "spark-3.4.1-bin-hadoop3")
-          .foreach { sparkHome =>
-            assertMatches(sparkHome, SPARK_HOME_REGEX_SCALA_212)
-            assertNotMatches(sparkHome, SPARK_HOME_REGEX_SCALA_213)
-          }
-      case "2.13" => Seq(
-          "spark-3.2.4-bin-hadoop3.2-scala2.13",
-          "spark-3.4.1-bin-hadoop3-scala2.13",
-          "spark-3.5.0-bin-hadoop3-scala2.13")
-          .foreach { sparkHome =>
-            assertMatches(sparkHome, SPARK_HOME_REGEX_SCALA_213)
-            assertNotMatches(sparkHome, SPARK_HOME_REGEX_SCALA_212)
-          }
+    Seq(
+      "spark-3.2.4-bin-hadoop3.2",
+      "spark-3.2.4-bin-hadoop2.7",
+      "spark-3.4.1-bin-hadoop3").foreach { SPARK3_HOME_SCALA_212 =>
+      assertMatches(SPARK3_HOME_SCALA_212, SPARK3_HOME_REGEX_SCALA_212)
+      assertNotMatches(SPARK3_HOME_SCALA_212, SPARK3_HOME_REGEX_SCALA_213)
+      assertNotMatches(SPARK3_HOME_SCALA_212, SPARK4_HOME_REGEX_SCALA_213)
+    }
+    Seq(
+      "spark-3.2.4-bin-hadoop3.2-scala2.13",
+      "spark-3.4.1-bin-hadoop3-scala2.13",
+      "spark-3.5.0-bin-hadoop3-scala2.13").foreach { SPARK3_HOME_SCALA_213 =>
+      assertMatches(SPARK3_HOME_SCALA_213, SPARK3_HOME_REGEX_SCALA_213)
+      assertNotMatches(SPARK3_HOME_SCALA_213, SPARK3_HOME_REGEX_SCALA_212)
+      assertNotMatches(SPARK3_HOME_SCALA_213, SPARK4_HOME_REGEX_SCALA_213)
+    }
+    Seq(
+      "spark-4.0.0-preview1-bin-hadoop3",
+      "spark-4.0.0-bin-hadoop3").foreach { SPARK4_HOME_SCALA_213 =>
+      assertMatches(SPARK4_HOME_SCALA_213, SPARK4_HOME_REGEX_SCALA_213)
+      assertNotMatches(SPARK4_HOME_SCALA_213, SPARK3_HOME_REGEX_SCALA_212)
+      assertNotMatches(SPARK4_HOME_SCALA_213, SPARK3_HOME_REGEX_SCALA_213)
     }
   }