fix: Configuration to stop coercion of tz for entity_df (#3255)

ammarar · web-flow · commit 97b7ab906081 · 2022-09-28T15:19:48.000-07:00
add option to skip coercion

Signed-off-by: ammarar &lt;ammar.alrashed@gmail.com&gt;

Signed-off-by: ammarar &lt;ammar.alrashed@gmail.com&gt;
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -1111,7 +1111,8 @@ def get_historical_features(
 
         # Check that the right request data is present in the entity_df
         if type(entity_df) == pd.DataFrame:
-            entity_df = utils.make_df_tzaware(cast(pd.DataFrame, entity_df))
+            if self.config.coerce_tz_aware:
+                entity_df = utils.make_df_tzaware(cast(pd.DataFrame, entity_df))
             for fv in request_feature_views:
                 for feature in fv.features:
                     if feature.name not in entity_df.columns:
diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py
@@ -167,6 +167,9 @@ class RepoConfig(FeastBaseModel):
     feature values for entities that have already been written into the online store.
     """
 
+    coerce_tz_aware: Optional[bool] = True
+    """ If True, coerces entity_df timestamp columns to be timezone aware (to UTC by default). """
+
     def __init__(self, **data: Any):
         super().__init__(**data)
 
