SPARKNLP-743: Add parameter to SparkNLP.start (#13510)

- Added params parameter which can supply custom
  configurations to the SparkSession

Author: DevinTDHa

python/sparknlp/__init__.py

@@ -107,11 +107,13 @@ def start(gpu=False,
         for WordEmbeddings. By default, this locations is the location of
         `hadoop.tmp.dir` set via Hadoop configuration for Apache Spark. NOTE: `S3` is
         not supported and it must be local, HDFS, or DBFS.
+    params : dict, optional
+        Custom parameters to set for the Spark configuration, by default None.
     cluster_tmp_dir : str, optional
         The location to save logs from annotators during training. If not set, it will
         be in the users home directory under `annotator_logs`.
     real_time_output : bool, optional
-        Whether to output in real time, by default False
+        Whether to read and print JVM output in real time, by default False
     output_level : int, optional
         Output level for logs, by default 1
 

src/main/scala/com/johnsnowlabs/nlp/SparkNLP.scala

@@ -47,6 +47,8 @@ object SparkNLP {
     * @param cluster_tmp_dir
     *   The location to save logs from annotators during training (By default, it will be in the
     *   users home directory under `annotator_logs`.)
+    * @param params
+    *   Custom parameters to set for the Spark configuration (Default: `Map.empty`)
     * @return
     *   SparkSession
     */
@@ -57,9 +59,13 @@ object SparkNLP {
       memory: String = "16G",
       cache_folder: String = "",
       log_folder: String = "",
-      cluster_tmp_dir: String = ""): SparkSession = {
+      cluster_tmp_dir: String = "",
+      params: Map[String, String] = Map.empty): SparkSession = {
 
-    val build = SparkSession
+    if (SparkSession.getActiveSession.isDefined)
+      println("Warning: Spark Session already created, some configs may not be applied.")
+
+    val builder = SparkSession
       .builder()
       .appName("Spark NLP")
       .master("local[*]")
@@ -68,26 +74,33 @@ object SparkNLP {
       .config("spark.kryoserializer.buffer.max", "2000M")
       .config("spark.driver.maxResultSize", "0")
 
-    if (apple_silicon) {
-      build.config("spark.jars.packages", MavenSparkSilicon)
-    } else if (aarch64) {
-      build.config("spark.jars.packages", MavenSparkAarch64)
-    } else if (gpu) {
-      build.config("spark.jars.packages", MavenGpuSpark3)
-    } else {
-      build.config("spark.jars.packages", MavenSpark3)
+    val sparkNlpJar =
+      if (apple_silicon) MavenSparkSilicon
+      else if (aarch64) MavenSparkAarch64
+      else if (gpu) MavenGpuSpark3
+      else MavenSpark3
+
+    if (!params.contains("spark.jars.packages")) {
+      builder.config("spark.jars.packages", sparkNlpJar)
+    }
+
+    params.foreach {
+      case (key, value) if key == "spark.jars.packages" =>
+        builder.config(key, sparkNlpJar + "," + value)
+      case (key, value) =>
+        builder.config(key, value)
     }
 
     if (cache_folder.nonEmpty)
-      build.config("spark.jsl.settings.pretrained.cache_folder", cache_folder)
+      builder.config("spark.jsl.settings.pretrained.cache_folder", cache_folder)
 
     if (log_folder.nonEmpty)
-      build.config("spark.jsl.settings.annotator.log_folder", log_folder)
+      builder.config("spark.jsl.settings.annotator.log_folder", log_folder)
 
     if (cluster_tmp_dir.nonEmpty)
-      build.config("spark.jsl.settings.storage.cluster_tmp_dir", cluster_tmp_dir)
+      builder.config("spark.jsl.settings.storage.cluster_tmp_dir", cluster_tmp_dir)
 
-    build.getOrCreate()
+    builder.getOrCreate()
   }
 
   def version(): String = {

src/test/java/com/johnsnowlabs/nlp/GeneralAnnotationsTest.java

@@ -27,6 +27,7 @@
 import org.apache.spark.sql.Encoders;
 import org.apache.spark.sql.Row;
 import org.apache.spark.sql.SparkSession;
+import scala.collection.immutable.HashMap;
 
 import java.util.LinkedList;
 
@@ -46,14 +47,13 @@ public static void main(String[] args) {
         Pipeline pipeline = new Pipeline();
         pipeline.setStages(new PipelineStage[]{document, tokenizer});
 
-        SparkSession spark = com.johnsnowlabs.nlp.SparkNLP.start(
-                false,
+        SparkSession spark = com.johnsnowlabs.nlp.SparkNLP.start(false,
                 false,
                 false,
                 "16G",
                 "",
                 "",
-                "");
+                "", new HashMap<>());
 
         LinkedList<String> text = new java.util.LinkedList<>();
 

src/test/scala/com/johnsnowlabs/nlp/SparkNLPTestSpec.scala

@@ -0,0 +1,27 @@
+package com.johnsnowlabs.nlp
+
+import com.johnsnowlabs.tags.SlowTest
+import com.johnsnowlabs.util.ConfigHelper.{awsJavaSdkVersion, hadoopAwsVersion}
+import org.scalatest.flatspec.AnyFlatSpec
+
+class SparkNLPTestSpec extends AnyFlatSpec {
+
+  behavior of "SparkNLPTestSpec"
+
+  it should "start with extra parameters" taggedAs SlowTest ignore {
+    val extraParams: Map[String, String] = Map(
+      "spark.jars.packages" -> ("org.apache.hadoop:hadoop-aws:" + hadoopAwsVersion + ",com.amazonaws:aws-java-sdk:" + awsJavaSdkVersion),
+      "spark.hadoop.fs.s3a.path.style.access" -> "true")
+
+    val spark = SparkNLP.start(params = extraParams)
+
+    assert(spark.conf.get("spark.hadoop.fs.s3a.path.style.access") == "true")
+
+    Seq(
+      "com.johnsnowlabs.nlp:spark-nlp",
+      "org.apache.hadoop:hadoop-aws",
+      "com.amazonaws:aws-java-sdk").foreach { pkg =>
+      assert(spark.conf.get("spark.jars.packages").contains(pkg))
+    }
+  }
+}