feast-dev · franciscojavierarceo · Jun 18, 2024 · Jun 16, 2024 · Jun 16, 2024 · Jun 16, 2024
@@ -23,6 +23,22 @@ def _serialize_val(
         raise ValueError(f"Value type not supported for Firestore: {v}")
 
 
+def _deserialize_value(value_type, value_bytes) -> ValueProto:
+    if value_type == ValueType.INT64:
+        value = struct.unpack("<q", value_bytes)[0]
+        return ValueProto(int64_val=value)
+    if value_type == ValueType.INT32:
+        value = struct.unpack("<i", value_bytes)[0]
+        return ValueProto(int32_val=value)
+    elif value_type == ValueType.STRING:
+        value = value_bytes.decode("utf-8")
+        return ValueProto(string_val=value)
+    elif value_type == ValueType.BYTES:
+        return ValueProto(bytes_val=value_bytes)
+    else:
+        raise ValueError(f"Unsupported value type: {value_type}")
+
+
 def serialize_entity_key_prefix(entity_keys: List[str]) -> bytes:
     """
     Serialize keys to a bytestring, so it can be used to prefix-scan through items stored in the online store
@@ -58,6 +74,8 @@ def serialize_entity_key(
     output: List[bytes] = []
     for k in sorted_keys:
         output.append(struct.pack("<I", ValueType.STRING))
+        if entity_key_serialization_version > 2:
+            output.append(struct.pack("<I", len(k)))
         output.append(k.encode("utf8"))
     for v in sorted_values:
         val_bytes, value_type = _serialize_val(
@@ -74,6 +92,57 @@ def serialize_entity_key(
     return b"".join(output)
 
 
+def deserialize_entity_key(
+    serialized_entity_key: bytes, entity_key_serialization_version=3
+) -> EntityKeyProto:
+    """
+    Deserialize entity key from a bytestring. This function can only be used with entity_key_serialization_version > 2.
+    Args:
+        entity_key_serialization_version: version of the entity key serialization
+        serialized_entity_key: serialized entity key bytes
+
+    Returns: EntityKeyProto
+
+    """
+    if entity_key_serialization_version <= 2:
+        raise ValueError(
+            "Deserialization of entity key with version <= 2 is not supported. Please use version > 2."
+        )
+    offset = 0
+    keys = []
+    values = []
+    while offset < len(serialized_entity_key):
+        key_type = struct.unpack_from("<I", serialized_entity_key, offset)[0]
+        offset += 4
+
+        # Read the length of the key
+        key_length = struct.unpack_from("<I", serialized_entity_key, offset)[0]
+        offset += 4
+
+        if key_type == ValueType.STRING:
+            key = struct.unpack_from(f"<{key_length}s", serialized_entity_key, offset)[
+                0
+            ]
+            keys.append(key.decode("utf-8").rstrip("\x00"))
+            offset += key_length
+        else:
+            raise ValueError(f"Unsupported key type: {key_type}")
+
+        (value_type,) = struct.unpack_from("<I", serialized_entity_key, offset)
+        offset += 4
+
+        (value_length,) = struct.unpack_from("<I", serialized_entity_key, offset)
+        offset += 4
+
+        # Read the value based on its type and length
+        value_bytes = serialized_entity_key[offset : offset + value_length]
+        value = _deserialize_value(value_type, value_bytes)
+        values.append(value)
+        offset += value_length
+
+    return EntityKeyProto(join_keys=keys, entity_values=values)
+
+
 def get_list_val_str(val):
     accept_value_types = [
         "float_list_val",

@@ -175,10 +175,12 @@ class RepoConfig(FeastBaseModel):
     used when writing data to the online store.
     A value <= 1 uses the serialization scheme used by feast up to Feast 0.22.
     A value of 2 uses a newer serialization scheme, supported as of Feast 0.23.
-    The main difference between the two scheme is that the serialization scheme v1 stored `long` values as `int`s,
-    which would result in errors trying to serialize a range of values.
-    v2 fixes this error, but v1 is kept around to ensure backwards compatibility - specifically the ability to read
+    A value of 3 uses the latest serialization scheme, supported as of Feast 0.38.
+    The main difference between the three schema is that
+    v1: the serialization scheme v1 stored `long` values as `int`s, which would result in errors trying to serialize a range of values.
+    v2: fixes this error, but v1 is kept around to ensure backwards compatibility - specifically the ability to read
     feature values for entities that have already been written into the online store.
+    v3: add entity_key value length to serialized bytes to enable deserialization, which can be used in retrieval of entity_key in document retrieval.
     """
 
     coerce_tz_aware: Optional[bool] = True

@@ -1,6 +1,6 @@
 import pytest
 
-from feast.infra.key_encoding_utils import serialize_entity_key
+from feast.infra.key_encoding_utils import deserialize_entity_key, serialize_entity_key
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 
@@ -28,3 +28,19 @@ def test_serialize_entity_key():
                 join_keys=["user"], entity_values=[ValueProto(int64_val=int(2**31))]
             ),
         )
+
+
+def test_deserialize_entity_key():
+    serialized_entity_key = serialize_entity_key(
+        EntityKeyProto(
+            join_keys=["user"], entity_values=[ValueProto(int64_val=int(2**15))]
+        ),
+        entity_key_serialization_version=3,
+    )
+
+    deserialized_entity_key = deserialize_entity_key(
+        serialized_entity_key, entity_key_serialization_version=3
+    )
+    assert deserialized_entity_key == EntityKeyProto(
+        join_keys=["user"], entity_values=[ValueProto(int64_val=int(2**15))]
+    )