fix: take spark dirver mem max from env var

bzkf · Dec 19, 2024 · eb6e683 · eb6e683
1 parent e7838b2
commit eb6e683
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/.github/workflows/benchmark.yaml b/.github/workflows/benchmark.yaml
@@ -22,6 +22,8 @@ jobs:
           version: "3.38.0"
 
       - name: Run benchmark
+        env:
+          SPARK_DRIVER_MEMORY: "8g"
         run: |
           task run
 

diff --git a/src/pathling_benchmark.py b/src/pathling_benchmark.py
@@ -1,4 +1,5 @@
 import datetime
+import os
 import time
 from pathling import PathlingContext, Expression as exp
 from pyspark.sql import SparkSession, DataFrame
@@ -29,7 +30,7 @@ def _init_pc(self):
             )
             .config(
                 "spark.driver.memory",
-                "64g",
+                os.getenv("SPARK_DRIVER_MEMORY", "64g"),
             )
             .config(
                 "spark.hadoop.fs.s3a.endpoint",