From 0daffeffe6d1415bb29ae38674c5027548eebd1e Mon Sep 17 00:00:00 2001 From: Bhargav Dodla Date: Wed, 29 Jan 2025 17:02:32 -0800 Subject: [PATCH] refactor: Move unserialization of artifacts outside the loop in spark_materialization_engine --- .../contrib/spark/spark_materialization_engine.py | 14 +++++++------- 1 file changed, 7 insertions(+), 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 index cf30a42d729..a39c0be3247 100644 --- 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