refactor: Move unserialization of artifacts outside the loop in spark…

…_materialization_engine
ExpediaGroup · Jan 30, 2025 · 0daffef · 0daffef
1 parent 780bb90
commit 0daffef
Showing 1 changed file with 7 additions and 7 deletions.
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
@@ -234,6 +234,13 @@ def _map_by_partition(
     spark_serialized_artifacts: _SparkSerializedArtifacts,
 ):
     """Load pandas df to online store"""
+
+    # unserialize artifacts
+    (
+        feature_view,
+        online_store,
+        repo_config,
+    ) = spark_serialized_artifacts.unserialize()
     for pdf in iterator:
         pdf_row_count = pdf.shape[0]
         start_time = time.time()
@@ -244,13 +251,6 @@ def _map_by_partition(
 
         table = pyarrow.Table.from_pandas(pdf)
 
-        # unserialize artifacts
-        (
-            feature_view,
-            online_store,
-            repo_config,
-        ) = spark_serialized_artifacts.unserialize()
-
         if feature_view.batch_source.field_mapping is not None:
             # Spark offline store does the field mapping during pull_latest_from_table_or_query
             # This is for the case where the offline store is not spark