feast-dev · mike0sv · Nov 26, 2020 · Nov 26, 2020 · Nov 26, 2020 · Nov 26, 2020
@@ -154,6 +154,9 @@ class ConfigOptions(metaclass=ConfigMeta):
     #: Directory where Spark is installed
     SPARK_HOME: Optional[str] = None
 
+    #: Addtional config options for Spark
+    SPARK_ADDITIONAL_OPTS: Optional[str] = None
+
     #: Dataproc cluster to run Feast Spark Jobs in
     DATAPROC_CLUSTER_NAME: Optional[str] = None
 

@@ -1,5 +1,6 @@
+import re
 from datetime import datetime
-from typing import TYPE_CHECKING, List, Union
+from typing import TYPE_CHECKING, Dict, List, Union
 
 from feast.config import Config
 from feast.constants import ConfigOptions as opt
@@ -26,7 +27,9 @@ def _standalone_launcher(config: Config) -> JobLauncher:
     from feast.pyspark.launchers import standalone
 
     return standalone.StandaloneClusterLauncher(
-        config.get(opt.SPARK_STANDALONE_MASTER), config.get(opt.SPARK_HOME),
+        config.get(opt.SPARK_STANDALONE_MASTER),
+        config.get(opt.SPARK_HOME),
+        _parse_additional_spark_options(config),
     )
 
 
@@ -41,6 +44,7 @@ def _dataproc_launcher(config: Config) -> JobLauncher:
         executor_instances=config.get(opt.DATAPROC_EXECUTOR_INSTANCES),
         executor_cores=config.get(opt.DATAPROC_EXECUTOR_CORES),
         executor_memory=config.get(opt.DATAPROC_EXECUTOR_MEMORY),
+        additional_options=_parse_additional_spark_options(config),
     )
 
 
@@ -57,6 +61,7 @@ def _get_optional(option):
         new_cluster_template_path=_get_optional(opt.EMR_CLUSTER_TEMPLATE_PATH),
         staging_location=config.get(opt.SPARK_STAGING_LOCATION),
         emr_log_location=config.get(opt.EMR_LOG_LOCATION),
+        additional_options=_parse_additional_spark_options(config),
     )
 
 
@@ -126,6 +131,26 @@ def _feature_table_to_argument(
     }
 
 
+def _quoted_split(string, delimiter):
+    for token in re.findall(f'(?:".*?"|[^{delimiter}])+', string):
+        if token.startswith('"') and token.endswith('"'):
+            token = token[1:-1]
+        yield token
+
+
+def _parse_additional_spark_options(config: Config) -> Dict[str, str]:
+    options_string = config.get(opt.SPARK_ADDITIONAL_OPTS, None)
+    if options_string is None:
+        return {}
+    try:
+        return dict(
+            _quoted_split(opt_val, "=")
+            for opt_val in _quoted_split(options_string, ";")
+        )
+    except ValueError:
+        raise ValueError(f"Cannot parse {opt.SPARK_ADDITIONAL_OPTS}: {options_string}")
+
+
 def start_historical_feature_retrieval_spark_session(
     client: "Client",
     project: str,
@@ -218,7 +243,6 @@ def start_offline_to_online_ingestion(
     start: datetime,
     end: datetime,
 ) -> BatchIngestionJob:
-
     launcher = resolve_launcher(client._config)
 
     return launcher.offline_to_online_ingestion(
@@ -268,7 +292,6 @@ def get_stream_to_online_ingestion_params(
 def start_stream_to_online_ingestion(
     client: "Client", project: str, feature_table: FeatureTable, extra_jars: List[str]
 ) -> StreamIngestionJob:
-
     launcher = resolve_launcher(client._config)
 
     return launcher.start_stream_to_online_ingestion(

@@ -145,6 +145,7 @@ def __init__(
         new_cluster_template_path: Optional[str],
         staging_location: str,
         emr_log_location: str,
+        additional_options: Dict[str, str],
     ):
         """
         Initialize a dataproc job controller client, used internally for job submission and result
@@ -162,6 +163,8 @@ def __init__(
                 An S3 staging location for artifacts.
             emr_log_location:
                 S3 location for EMR logs.
+            additional_options:
+                Additional configuration options for Spark job
         """
 
         assert existing_cluster_id or new_cluster_template_path
@@ -177,6 +180,7 @@ def __init__(
         self._staging_location = staging_location
         self._emr_log_location = emr_log_location
         self._region = region
+        self._additional_options = additional_options
 
     def _emr_client(self):
 
@@ -230,6 +234,7 @@ def historical_feature_retrieval(
 
         step = _historical_retrieval_step(
             pyspark_script_path,
+            conf=self._additional_options,
             args=job_params.get_arguments(),
             output_file_uri=job_params.get_destination_path(),
         )
@@ -260,6 +265,7 @@ def offline_to_online_ingestion(
         step = _sync_offline_to_online_step(
             jar_s3_path,
             ingestion_job_params.get_feature_table_name(),
+            self._additional_options,
             args=ingestion_job_params.get_arguments(),
         )
 
@@ -293,6 +299,7 @@ def start_stream_to_online_ingestion(
             jar_s3_path,
             extra_jar_paths,
             ingestion_job_params.get_feature_table_name(),
+            self._additional_options,
             args=ingestion_job_params.get_arguments(),
             job_hash=job_hash,
         )

@@ -135,7 +135,7 @@ def _upload_jar(jar_s3_prefix: str, local_path: str) -> str:
 
 
 def _sync_offline_to_online_step(
-    jar_path: str, feature_table_name: str, args: List[str],
+    jar_path: str, feature_table_name: str, conf: Dict[str, str], args: List[str],
 ) -> Dict[str, Any]:
 
     return {
@@ -155,6 +155,7 @@ def _sync_offline_to_online_step(
                 "spark-submit",
                 "--class",
                 "feast.ingestion.IngestionJob",
+                *_prepare_conf_args(conf),
                 "--packages",
                 "com.google.cloud.spark:spark-bigquery-with-dependencies_2.12:0.17.2",
                 jar_path,
@@ -344,8 +345,17 @@ def _upload_dataframe(s3prefix: str, df: pandas.DataFrame) -> str:
         )
 
 
+def _prepare_conf_args(conf: Dict[str, str]):
+    return [
+        _ for name, value in conf.items() for _ in ["--conf", f'"{name}"="{value}"']
+    ]
+
+
 def _historical_retrieval_step(
-    pyspark_script_path: str, args: List[str], output_file_uri: str,
+    pyspark_script_path: str,
+    conf: Dict[str, str],
+    args: List[str],
+    output_file_uri: str,
 ) -> Dict[str, Any]:
 
     return {
@@ -361,7 +371,8 @@ def _historical_retrieval_step(
                     "Value": output_file_uri,
                 },
             ],
-            "Args": ["spark-submit", pyspark_script_path] + args,
+            "Args": ["spark-submit", *_prepare_conf_args(conf), pyspark_script_path]
+            + args,
             "Jar": "command-runner.jar",
         },
     }
@@ -371,6 +382,7 @@ def _stream_ingestion_step(
     jar_path: str,
     extra_jar_paths: List[str],
     feature_table_name: str,
+    conf: Dict[str, str],
     args: List[str],
     job_hash: str,
 ) -> Dict[str, Any]:
@@ -395,6 +407,7 @@ def _stream_ingestion_step(
                 {"Key": "feast.step_metadata.job_hash", "Value": job_hash},
             ],
             "Args": ["spark-submit", "--class", "feast.ingestion.IngestionJob"]
+            + _prepare_conf_args(conf)
             + jars_args
             + [
                 "--packages",

@@ -208,6 +208,7 @@ def __init__(
         executor_instances: str,
         executor_cores: str,
         executor_memory: str,
+        additional_options: Dict[str, str] = None,
     ):
         """
         Initialize a dataproc job controller client, used internally for job submission and result
@@ -228,6 +229,8 @@ def __init__(
                 Number of cores for dataproc job.
             executor_memory (str):
                 Amount of memory for dataproc job.
+            additional_options (Dict[str, str]):
+                Additional configuration options for Spark job
         """
 
         self.cluster_name = cluster_name
@@ -247,6 +250,7 @@ def __init__(
         self.executor_instances = executor_instances
         self.executor_cores = executor_cores
         self.executor_memory = executor_memory
+        self.additional_options = additional_options or {}
 
     def _stage_file(self, file_path: str, job_id: str) -> str:
         if not os.path.isfile(file_path):
@@ -285,6 +289,7 @@ def dataproc_submit(
                             "spark.executor.instances": self.executor_instances,
                             "spark.executor.cores": self.executor_cores,
                             "spark.executor.memory": self.executor_memory,
+                            **self.additional_options,
                         },
                     }
                 }

@@ -224,7 +224,12 @@ class StandaloneClusterLauncher(JobLauncher):
 
     BQ_CONNECTOR_VERSION = "2.12:0.17.3"
 
-    def __init__(self, master_url: str, spark_home: str = None):
+    def __init__(
+        self,
+        master_url: str,
+        spark_home: str = None,
+        additional_options: Dict[str, str] = None,
+    ):
         """
         This launcher executes the spark-submit script in a subprocess. The subprocess
         will run until the Pyspark driver exits.
@@ -235,9 +240,12 @@ def __init__(self, master_url: str, spark_home: str = None):
             spark_home (str):
                 Local file path to Spark installation directory. If not provided,
                 the environmental variable `SPARK_HOME` will be used instead.
+            additional_options (Dict[str, str]):
+                Additional configuration options for Spark job
         """
         self.master_url = master_url
         self.spark_home = spark_home if spark_home else os.getenv("SPARK_HOME")
+        self.additional_options = additional_options
 
     @property
     def spark_submit_script_path(self):
@@ -285,6 +293,10 @@ def spark_submit(
             ]
         )
 
+        if self.additional_options is not None:
+            for option, value in self.additional_options.items():
+                submission_cmd.extend(["--conf", f'"{option}"="{value}"'])
+
         submission_cmd.append(job_params.get_main_file_path())
         submission_cmd.extend(job_params.get_arguments())
 

@@ -0,0 +1,29 @@
+import pytest
+
+from feast.config import Config
+from feast.constants import ConfigOptions as opt
+from feast.pyspark.launcher import _parse_additional_spark_options
+
+
+class TestSparkAdditionalOpts:
+    def parse(self, options_string):
+        return _parse_additional_spark_options(
+            Config(options={opt.SPARK_ADDITIONAL_OPTS: options_string})
+        )
+
+    def test_normal_options(self):
+        options_string = "option1=aaaa;option2=bbb"
+        assert self.parse(options_string) == {"option1": "aaaa", "option2": "bbb"}
+
+    def test_value_with_delimiter(self):
+        options_string = 'option1=aaaa;option2="b;b"'
+        assert self.parse(options_string) == {"option1": "aaaa", "option2": "b;b"}
+
+    def test_value_with_another_delimiter(self):
+        options_string = 'option1=aaaa;option2="b=b"'
+        assert self.parse(options_string) == {"option1": "aaaa", "option2": "b=b"}
+
+    def test_error_on_wrong_format(self):
+        options_string = "option1=aaaa;option2"
+        with pytest.raises(ValueError):
+            self.parse(options_string)
diff --git a/tests/integration/fixtures/launchers.py b/tests/integration/fixtures/launchers.py
@@ -20,4 +20,5 @@ def dataproc_launcher(pytestconfig) -> DataprocClusterLauncher:
         executor_instances=executor_instances,
         executor_cores=executor_cores,
         executor_memory=executor_memory,
+        additional_options={}
     )