feat: Add warning suppression flag to spark materialization engine (#299)

omirandadev · omiranda · web-flow · commit 66f7f75968de · 2025-09-18T10:02:01.000-05:00
* feat: add warning suppression flag to spark materialization engine

* chore: rearrange imports

* small reformat

---------

Co-authored-by: omiranda &lt;omiranda@expediagroup.com&gt;
diff --git a/sdk/python/feast/infra/materialization/contrib/spark/spark_materialization_engine.py b/sdk/python/feast/infra/materialization/contrib/spark/spark_materialization_engine.py
@@ -47,6 +47,9 @@ class SparkMaterializationEngineConfig(FeastConfigBaseModel):
     partitions: int = 0
     """Number of partitions to use when writing data to online store. If 0, no repartitioning is done"""
 
+    suppress_warnings: bool = True
+    """Suppress Pyspark deprecation warnings during materialization"""
+
 
 @dataclass
 class SparkMaterializationJob(MaterializationJob):
@@ -267,6 +270,25 @@ def _map_by_partition(
 ):
     feature_view, online_store, repo_config = spark_serialized_artifacts.unserialize()
 
+    if (
+        hasattr(repo_config.batch_engine, "suppress_warnings")
+        and repo_config.batch_engine.suppress_warnings
+    ):
+        import os
+        import warnings
+
+        os.environ["PYTHONWARNINGS"] = "ignore::DeprecationWarning"
+        warnings.filterwarnings("ignore", category=DeprecationWarning)
+        warnings.filterwarnings(
+            "ignore", message=".*is_categorical_dtype is deprecated.*"
+        )
+        warnings.filterwarnings(
+            "ignore", message=".*is_datetime64tz_dtype is deprecated.*"
+        )
+        warnings.filterwarnings(
+            "ignore", message=".*distutils Version classes are deprecated.*"
+        )
+
     total_batches = 0
     total_time = 0.0
     min_time = float("inf")
