fix: Added unix_timestamp_val in _serialize_val (#5659)

ntkathole · web-flow · commit 35a84237bb15 · 2025-10-12T22:02:32.000-04:00
diff --git a/sdk/python/feast/infra/key_encoding_utils.py b/sdk/python/feast/infra/key_encoding_utils.py
@@ -22,6 +22,8 @@ def _serialize_val(
         if 0 <= entity_key_serialization_version <= 1:
             return struct.pack("<l", v.int64_val), ValueType.INT64
         return struct.pack("<q", v.int64_val), ValueType.INT64
+    elif value_type == "unix_timestamp_val":
+        return struct.pack("<q", v.unix_timestamp_val), ValueType.UNIX_TIMESTAMP
     else:
         raise ValueError(f"Value type not supported for feast feature store: {v}")
 
@@ -38,6 +40,9 @@ def _deserialize_value(value_type, value_bytes) -> ValueProto:
         return ValueProto(string_val=value)
     elif value_type == ValueType.BYTES:
         return ValueProto(bytes_val=value_bytes)
+    elif value_type == ValueType.UNIX_TIMESTAMP:
+        value = struct.unpack("<q", value_bytes)[0]
+        return ValueProto(unix_timestamp_val=value)
     else:
         raise ValueError(f"Unsupported value type: {value_type}")
 
diff --git a/sdk/python/tests/unit/infra/test_key_encoding_utils.py b/sdk/python/tests/unit/infra/test_key_encoding_utils.py
@@ -79,6 +79,15 @@ def test_serialize_value():
     assert t == ValueType.INT64
     assert v == b"\x01\x00\x00\x00\x00\x00\x00\x00"
 
+    # Test unix_timestamp_val serialization
+    v, t = _serialize_val(
+        "unix_timestamp_val", ValueProto(unix_timestamp_val=1758823656)
+    )
+    assert t == ValueType.UNIX_TIMESTAMP
+    # Verify roundtrip: deserialize the serialized value
+    deserialized = _deserialize_value(ValueType.UNIX_TIMESTAMP, v)
+    assert deserialized.unix_timestamp_val == 1758823656
+
 
 def test_deserialize_value():
     v = _deserialize_value(ValueType.STRING, b"test")
@@ -93,6 +102,33 @@ def test_deserialize_value():
     v = _deserialize_value(ValueType.INT64, b"\x01\x00\x00\x00\x00\x00\x00\x00")
     assert v.int64_val == 1
 
+    timestamp_val = 1758823656
+    serialized_bytes, _ = _serialize_val(
+        "unix_timestamp_val", ValueProto(unix_timestamp_val=timestamp_val)
+    )
+    v = _deserialize_value(ValueType.UNIX_TIMESTAMP, serialized_bytes)
+    assert v.unix_timestamp_val == timestamp_val
+
+
+def test_serialize_deserialize_unix_timestamp_entity():
+    entity_key_proto = EntityKeyProto(
+        join_keys=["e2"],
+        entity_values=[ValueProto(unix_timestamp_val=1758823656)],
+    )
+
+    serialized_key = serialize_entity_key(
+        entity_key_proto,
+        entity_key_serialization_version=3,
+    )
+
+    deserialized_key = deserialize_entity_key(
+        serialized_key,
+        entity_key_serialization_version=3,
+    )
+
+    assert deserialized_key == entity_key_proto
+    assert deserialized_key.entity_values[0].unix_timestamp_val == 1758823656
+
 
 def test_reserialize_entity_v2_key_to_v3():
     entity_key_proto_v2 = EntityKeyProto(
