chore: Clean up entity key serialization (feast-dev#3199)

* Clean up entity key serialization Signed-off-by: Felix Wang <[email protected]> * Remove incorrect docs Signed-off-by: Felix Wang <[email protected]> * Format Signed-off-by: Felix Wang <[email protected]> Signed-off-by: Felix Wang <[email protected]>
franciscojavierarceo · Sep 8, 2022 · 782c759 · 782c759
1 parent 2535024
commit 782c759
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 20 deletions.
diff --git a/docs/how-to-guides/customizing-feast/adding-support-for-a-new-online-store.md b/docs/how-to-guides/customizing-feast/adding-support-for-a-new-online-store.md
@@ -154,7 +154,10 @@ def online_write_batch(
     project = config.project
 
     for entity_key, values, timestamp, created_ts in data:
-        entity_key_bin = serialize_entity_key(entity_key).hex()
+        entity_key_bin = serialize_entity_key(
+            entity_key,
+            entity_key_serialization_version=config.entity_key_serialization_version,
+        ).hex()
         timestamp = _to_naive_utc(timestamp)
         if created_ts is not None:
             created_ts = _to_naive_utc(created_ts)
@@ -184,7 +187,10 @@ def online_read(
 
     project = config.project
     for entity_key in entity_keys:
-        entity_key_bin = serialize_entity_key(entity_key).hex()
+        entity_key_bin = serialize_entity_key(
+            entity_key,
+            entity_key_serialization_version=config.entity_key_serialization_version,
+        ).hex()
         print(f"entity_key_bin: {entity_key_bin}")
 
         cur.execute(
@@ -208,18 +214,6 @@ def online_read(
 ```
 {% endcode %}
 
-### 1.3 Type Mapping
-
-Most online stores will have to perform some custom mapping of online store datatypes to feast value types.
-
-* The function to implement here are `source_datatype_to_feast_value_type` and `get_column_names_and_types` in your `DataSource` class.
-* `source_datatype_to_feast_value_type` is used to convert your DataSource's datatypes to feast value types.
-* `get_column_names_and_types` retrieves the column names and corresponding datasource types.
-
-Add any helper functions for type conversion to `sdk/python/feast/type_map.py`.
-
-* Be sure to implement correct type mapping so that Feast can process your feature columns without casting incorrectly that can potentially cause loss of information or incorrect data.
-
 ## 2. Defining an OnlineStoreConfig class
 
 Additional configuration may be needed to allow the OnlineStore to talk to the backing store. For example, MySQL may need configuration information like the host at which the MySQL instance is running, credentials for connecting to the database, etc.

diff --git a/...python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py b/...python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py
@@ -314,7 +314,8 @@ def online_write_batch(
         project = config.project
         for entity_key, values, timestamp, created_ts in data:
             entity_key_bin = serialize_entity_key(
-                entity_key, entity_key_serialization_version=2
+                entity_key,
+                entity_key_serialization_version=config.entity_key_serialization_version,
             ).hex()
             with tracing_span(name="remote_call"):
                 self._write_rows(
@@ -353,7 +354,8 @@ def online_read(
 
         for entity_key in entity_keys:
             entity_key_bin = serialize_entity_key(
-                entity_key, entity_key_serialization_version=2
+                entity_key,
+                entity_key_serialization_version=config.entity_key_serialization_version,
             ).hex()
 
             with tracing_span(name="remote_call"):

diff --git a/sdk/python/feast/infra/online_stores/snowflake.py b/sdk/python/feast/infra/online_stores/snowflake.py
@@ -94,9 +94,13 @@ def online_write_batch(
 
             for j, (feature_name, val) in enumerate(values.items()):
                 df.loc[j, "entity_feature_key"] = serialize_entity_key(
-                    entity_key, 2
+                    entity_key,
+                    entity_key_serialization_version=config.entity_key_serialization_version,
                 ) + bytes(feature_name, encoding="utf-8")
-                df.loc[j, "entity_key"] = serialize_entity_key(entity_key, 2)
+                df.loc[j, "entity_key"] = serialize_entity_key(
+                    entity_key,
+                    entity_key_serialization_version=config.entity_key_serialization_version,
+                )
                 df.loc[j, "feature_name"] = feature_name
                 df.loc[j, "value"] = val.SerializeToString()
                 df.loc[j, "event_ts"] = timestamp
@@ -162,7 +166,10 @@ def online_read(
                 (
                     "TO_BINARY("
                     + hexlify(
-                        serialize_entity_key(combo[0], 2)
+                        serialize_entity_key(
+                            combo[0],
+                            entity_key_serialization_version=config.entity_key_serialization_version,
+                        )
                         + bytes(combo[1], encoding="utf-8")
                     ).__str__()[1:]
                     + ")"
@@ -184,7 +191,10 @@ def online_read(
             df = execute_snowflake_statement(conn, query).fetch_pandas_all()
 
         for entity_key in entity_keys:
-            entity_key_bin = serialize_entity_key(entity_key, 2)
+            entity_key_bin = serialize_entity_key(
+                entity_key,
+                entity_key_serialization_version=config.entity_key_serialization_version,
+            )
             res = {}
             res_ts = None
             for index, row in df[df["entity_key"] == entity_key_bin].iterrows():