Merge pull request #361 from s22s/feature/353-conf-to-create-spark-sess

Pass kwargs in python create_rf_spark_session to spark conf

Author: metasim

pyrasterframes/src/main/python/pyrasterframes/__init__.py

@@ -55,11 +55,11 @@ def _kryo_init(builder):
     return builder
 
 
-def get_spark_session(master="local[*]"):
+def get_spark_session(master="local[*]", **kwargs):
     """ Create a SparkSession with pyrasterframes enabled and configured. """
     from pyrasterframes.utils import create_rf_spark_session
 
-    return create_rf_spark_session(master)
+    return create_rf_spark_session(master, **kwargs)
 
 
 def _convert_df(df, sp_key=None, metadata=None):

pyrasterframes/src/main/python/pyrasterframes/utils.py

@@ -20,6 +20,7 @@
 
 import glob
 from pyspark.sql import SparkSession
+from pyspark import SparkConf
 import os
 import sys
 from . import RFContext
@@ -76,15 +77,22 @@ def find_pyrasterframes_assembly():
     return jarpath[0]
 
 
-def create_rf_spark_session(master="local[*]"):
+def create_rf_spark_session(master="local[*]", **kwargs):
     """ Create a SparkSession with pyrasterframes enabled and configured. """
     jar_path = find_pyrasterframes_assembly()
 
+    if 'spark.jars' in kwargs.keys():
+        if 'pyrasterframes' not in kwargs['spark.jars']:
+            raise UserWarning("spark.jars config is set, but it seems to be missing the pyrasterframes assembly jar.")
+
+    conf = SparkConf().setAll([(k, kwargs[k]) for k in kwargs])
+
     spark = (SparkSession.builder
              .master(master)
              .appName("RasterFrames")
              .config('spark.jars', jar_path)
              .withKryoSerialization()
+             .config(conf=conf)  # user can override the defaults
              .getOrCreate())
 
     try:

pyrasterframes/src/main/python/tests/PyRasterFramesTests.py

@@ -28,6 +28,14 @@
 from . import TestEnvironment
 
 
+class UtilTest(TestEnvironment):
+
+    def test_spark_confs(self):
+        from . import app_name
+        self.assertEqual(self.spark.conf.get('spark.app.name'), app_name)
+        self.assertEqual(self.spark.conf.get('spark.ui.enabled'), 'false')
+
+
 class CellTypeHandling(unittest.TestCase):
 
     def test_is_raw(self):

pyrasterframes/src/main/python/tests/__init__.py

@@ -31,6 +31,7 @@
 else:
     import __builtin__ as builtins
 
+app_name = 'pyrasterframes test suite'
 
 def resource_dir():
     def pdir(curr):
@@ -46,7 +47,10 @@ def pdir(curr):
 
 
 def spark_test_session():
-    spark = create_rf_spark_session()
+    spark = create_rf_spark_session(**{
+        'spark.ui.enabled': 'false',
+        'spark.app.name': app_name
+    })
     spark.sparkContext.setLogLevel('ERROR')
 
     print("Spark Version: " + spark.version)