add spark helper

.gitignore

@@ -1,4 +1,4 @@
 .DS_Store
 .idea
 .coverage
-
+*_pycache__

Dockerfile

@@ -12,6 +12,9 @@ RUN apt-get update && apt-get install -y \
 # Install Jupyterlab and other python dependencies
 RUN pip3 install jupyterlab==4.2.0 pyspark==3.5.1
 
+COPY ./src/ /src
+ENV PYTHONPATH "${PYTHONPATH}:/src"
+
 COPY scripts/entrypoint.sh /opt/
 RUN chmod a+x /opt/entrypoint.sh
 

README.md

@@ -63,23 +63,37 @@ When running Spark in the Jupyter notebook container, the default `spark.driver.
 the hostname (`SPARK_DRIVER_HOST`) of the container. 
 In addition, the environment variable `SPARK_MASTER_URL` should also be configured.
 
+#### Using Predefined SparkSession from `spark.utils.get_spark_session` method
+```python
+from spark.utils import get_spark_session
+
+spark = get_spark_session('TestApp')
+```
+
+#### Manually Configuring SparkSession/SparkContext
+
+If you want to configure the SparkSession manually, you can do so as follows:
+
 #### Example SparkSession Configuration
 ```python
 spark = SparkSession.builder \
+    .master(os.environ['SPARK_MASTER_URL']) \
     .appName("TestSparkJob") \
     .getOrCreate()
 ```
 
 #### Example SparkContext Configuration
 ```python
-conf = SparkConf(). \
-    setAppName("TestSparkJob")
+conf = SparkConf() \
+    .setMaster(os.environ['SPARK_MASTER_URL']) \
+    .setAppName("TestSparkJob")
 sc = SparkContext(conf=conf)
 ```
 
 #### Submitting a Job Using Terminal
 ```bash
 /opt/bitnami/spark/bin/spark-submit \
+    --master $SPARK_MASTER_URL \
     /opt/bitnami/spark/examples/src/main/python/pi.py 10 \
     2>/dev/null
 ```

src/spark/utils.py

@@ -0,0 +1,32 @@
+import os
+
+from pyspark.conf import SparkConf
+from pyspark.sql import SparkSession
+
+
+def get_spark_session(app_name: str, local: bool = False) -> SparkSession:
+    """
+    Helper to get and manage the `SparkSession` and keep all of our spark configuration params in one place.
+
+    :param app_name: The name of the application
+    :param local: Whether to run the spark session locally or not
+
+    :return: A `SparkSession` object
+    """
+
+    if local:
+        return SparkSession.builder.appName(app_name).getOrCreate()
+
+    spark_conf = SparkConf()
+
+    spark_conf.setAll(
+        [
+            (
+                "spark.master",
+                os.environ.get("SPARK_MASTER_URL", "spark://spark-master:7077"),
+            ),
+            ("spark.app.name", app_name),
+        ]
+    )
+
+    return SparkSession.builder.config(conf=spark_conf).getOrCreate()