diff --git a/sdk/python/feast/type_map.py b/sdk/python/feast/type_map.py index d22c51d5963..ed4b7cba594 100644 --- a/sdk/python/feast/type_map.py +++ b/sdk/python/feast/type_map.py @@ -46,6 +46,9 @@ from feast.protos.feast.types.Value_pb2 import Value as ProtoValue from feast.value_type import ListType, ValueType +# null timestamps get converted to -9223372036854775808 +NULL_TIMESTAMP_INT_VALUE = np.datetime64("NaT").astype(int) + def feast_value_type_to_python_type(field_value_proto: ProtoValue) -> Any: """ @@ -69,9 +72,18 @@ def feast_value_type_to_python_type(field_value_proto: ProtoValue) -> Any: # Convert UNIX_TIMESTAMP values to `datetime` if val_attr == "unix_timestamp_list_val": - val = [datetime.fromtimestamp(v, tz=timezone.utc) for v in val] + val = [ + datetime.fromtimestamp(v, tz=timezone.utc) + if v != NULL_TIMESTAMP_INT_VALUE + else None + for v in val + ] elif val_attr == "unix_timestamp_val": - val = datetime.fromtimestamp(val, tz=timezone.utc) + val = ( + datetime.fromtimestamp(val, tz=timezone.utc) + if val != NULL_TIMESTAMP_INT_VALUE + else None + ) return val diff --git a/sdk/python/tests/unit/test_type_map.py b/sdk/python/tests/unit/test_type_map.py new file mode 100644 index 00000000000..0ba259ab73f --- /dev/null +++ b/sdk/python/tests/unit/test_type_map.py @@ -0,0 +1,28 @@ +import numpy as np + +from feast.type_map import ( + feast_value_type_to_python_type, + python_values_to_proto_values, +) +from feast.value_type import ValueType + + +def test_null_unix_timestamp(): + """Test that null UnixTimestamps get converted from proto correctly.""" + + data = np.array(["NaT"], dtype="datetime64") + protos = python_values_to_proto_values(data, ValueType.UNIX_TIMESTAMP) + converted = feast_value_type_to_python_type(protos[0]) + + assert converted is None + + +def test_null_unix_timestamp_list(): + """Test that UnixTimestamp lists with a null get converted from proto + correctly.""" + + data = np.array([["NaT"]], dtype="datetime64") + protos = python_values_to_proto_values(data, ValueType.UNIX_TIMESTAMP_LIST) + converted = feast_value_type_to_python_type(protos[0]) + + assert converted[0] is None