fix: Snowflake_online_read fix

sfc-gh-madkins · sfc-gh-madkins · commit 49e2bf00f50f · 2022-08-01T09:23:29.000-05:00
Signed-off-by: Miles Adkins &lt;miles.adkins@snowflake.com&gt;
diff --git a/sdk/python/feast/infra/online_stores/snowflake.py b/sdk/python/feast/infra/online_stores/snowflake.py
@@ -96,11 +96,11 @@ def online_write_batch(
             for j, (feature_name, val) in enumerate(values.items()):
                 df.loc[j, "entity_feature_key"] = serialize_entity_key(
                     entity_key,
-                    entity_key_serialization_version=entity_key_serialization_version,
+                    entity_key_serialization_version,
                 ) + bytes(feature_name, encoding="utf-8")
                 df.loc[j, "entity_key"] = serialize_entity_key(
                     entity_key,
-                    entity_key_serialization_version=entity_key_serialization_version,
+                    entity_key_serialization_version,
                 )
                 df.loc[j, "feature_name"] = feature_name
                 df.loc[j, "value"] = val.SerializeToString()
@@ -156,12 +156,18 @@ def online_read(
 
         result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []
 
+        entity_key_serialization_version = (
+            config.entity_key_serialization_version
+            if config.entity_key_serialization_version
+            else 2
+        )
+
         entity_fetch_str = ",".join(
             [
                 (
                     "TO_BINARY("
                     + hexlify(
-                        serialize_entity_key(combo[0])
+                        serialize_entity_key(combo[0], entity_key_serialization_version)
                         + bytes(combo[1], encoding="utf-8")
                     ).__str__()[1:]
                     + ")"
@@ -187,15 +193,10 @@ def online_read(
                 .fetch_pandas_all()
             )
 
-        entity_key_serialization_version = (
-            config.entity_key_serialization_version
-            if config.entity_key_serialization_version
-            else 2
-        )
         for entity_key in entity_keys:
             entity_key_bin = serialize_entity_key(
                 entity_key,
-                entity_key_serialization_version=entity_key_serialization_version,
+                entity_key_serialization_version,
             )
             res = {}
             res_ts = None
diff --git a/sdk/python/feast/templates/snowflake/feature_store.yaml b/sdk/python/feast/templates/snowflake/feature_store.yaml
@@ -9,6 +9,7 @@ offline_store:
     role: SNOWFLAKE_ROLE
     warehouse: SNOWFLAKE_WAREHOUSE
     database: SNOWFLAKE_DATABASE
+entity_key_serialization_version: 2
 online_store:
     type: snowflake.online
     account: SNOWFLAKE_DEPLOYMENT_URL
