Add setting to tune spark

pola-rs · Apr 5, 2024 · 3fd91a1 · 3fd91a1
1 parent 1c9d39f
commit 3fd91a1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/queries/pyspark/utils.py b/queries/pyspark/utils.py
@@ -21,10 +21,13 @@
 
 def get_or_create_spark() -> SparkSession:
     spark = (
-        SparkSession.builder.appName("spark_queries").master("local[*]").getOrCreate()
+        SparkSession.builder.appName("spark_queries")
+        .master("local[*]")
+        .config("spark.driver.memory", settings.run.spark_driver_memory)
+        .config("spark.executor.memory", settings.run.spark_executor_memory)
+        .config("spark.log.level", settings.run.spark_log_level)
+        .getOrCreate()
     )
-    spark.sparkContext.setLogLevel(settings.run.spark_log_level)
-
     return spark
 
 

diff --git a/settings.py b/settings.py
@@ -26,10 +26,13 @@ class Run(BaseSettings):
     show_results: bool = False
     check_results: bool = False  # Only available for SCALE_FACTOR=1
 
-    spark_log_level: str = "ERROR"
     polars_show_plan: bool = False
     polars_streaming: bool = False
 
+    spark_driver_memory: str = "1g"  # Tune as needed for optimal performance
+    spark_executor_memory: str = "1g"  # Tune as needed for optimal performance
+    spark_log_level: str = "ERROR"
+
     model_config = SettingsConfigDict(
         env_prefix="run_", env_file=".env", extra="ignore"
     )