diff --git a/sdk/python/feast/infra/key_encoding_utils.py b/sdk/python/feast/infra/key_encoding_utils.py index ca834f1917..1f9ffeef14 100644 --- a/sdk/python/feast/infra/key_encoding_utils.py +++ b/sdk/python/feast/infra/key_encoding_utils.py @@ -20,7 +20,23 @@ def _serialize_val( return struct.pack(" ValueProto: + if value_type == ValueType.INT64: + value = struct.unpack(" bytes: @@ -50,6 +66,15 @@ def serialize_entity_key( serialize to the same byte string[1]. [1] https://developers.google.com/protocol-buffers/docs/encoding + + Args: + entity_key_serialization_version: version of the entity key serialization + version 1: int64 values are serialized as 4 bytes + version 2: int64 values are serialized as 8 bytes + version 3: entity_key size is added to the serialization for deserialization purposes + entity_key: EntityKeyProto + + Returns: bytes of the serialized entity key """ sorted_keys, sorted_values = zip( *sorted(zip(entity_key.join_keys, entity_key.entity_values)) @@ -58,6 +83,8 @@ def serialize_entity_key( output: List[bytes] = [] for k in sorted_keys: output.append(struct.pack(" 2: + output.append(struct.pack(" 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 by setting entity_key_serialization_version=3" + ) + offset = 0 + keys = [] + values = [] + while offset < len(serialized_entity_key): + key_type = struct.unpack_from("