fix: Conversion of null timestamp from proto to python (#2814)

Andrew Pope · web-flow · commit cb23648da71c · 2022-06-17T12:07:06.000-07:00
Signed-off-by: Andrew Pope &lt;apope@nursefly.com&gt;
diff --git 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
@@ -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
