chore: Updating entity key serialization for multiple entities (feast-dev#4967)

Author: franciscojavierarceo

sdk/python/feast/infra/key_encoding_utils.py

@@ -83,6 +83,8 @@ def serialize_entity_key(
     )
 
     output: List[bytes] = []
+    if entity_key_serialization_version > 2:
+        output.append(struct.pack("<I", len(sorted_keys)))
     for k in sorted_keys:
         output.append(struct.pack("<I", ValueType.STRING))
         if entity_key_serialization_version > 2:
@@ -122,7 +124,11 @@ def deserialize_entity_key(
     offset = 0
     keys = []
     values = []
-    while offset < len(serialized_entity_key):
+
+    num_keys = struct.unpack_from("<I", serialized_entity_key, offset)[0]
+    offset += 4
+
+    for _ in range(num_keys):
         key_type = struct.unpack_from("<I", serialized_entity_key, offset)[0]
         offset += 4
 
@@ -139,6 +145,7 @@ def deserialize_entity_key(
         else:
             raise ValueError(f"Unsupported key type: {key_type}")
 
+    while offset < len(serialized_entity_key):
         (value_type,) = struct.unpack_from("<I", serialized_entity_key, offset)
         offset += 4
 

sdk/python/tests/unit/infra/test_key_encoding_utils.py

@@ -52,6 +52,24 @@ def test_deserialize_entity_key():
     )
 
 
+def test_deserialize_multiple_entity_keys():
+    entity_key_proto = EntityKeyProto(
+        join_keys=["customer", "user"],
+        entity_values=[ValueProto(string_val="test"), ValueProto(int64_val=int(2**15))],
+    )
+
+    serialized_entity_key = serialize_entity_key(
+        entity_key_proto,
+        entity_key_serialization_version=3,
+    )
+
+    deserialized_entity_key = deserialize_entity_key(
+        serialized_entity_key,
+        entity_key_serialization_version=3,
+    )
+    assert deserialized_entity_key == entity_key_proto
+
+
 def test_serialize_value():
     v, t = _serialize_val("string_val", ValueProto(string_val="test"))
     assert t == ValueType.STRING