Skip to content

Commit

Permalink
fix conversion of null timestamp from proto to python
Browse files Browse the repository at this point in the history
Signed-off-by: Andrew Pope <apope@nursefly.com>
  • Loading branch information
Andrew Pope committed Jun 17, 2022
1 parent 6b4e8d0 commit 02379aa
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 2 deletions.
16 changes: 14 additions & 2 deletions sdk/python/feast/type_map.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand All @@ -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

Expand Down
28 changes: 28 additions & 0 deletions sdk/python/tests/unit/test_type_map.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 02379aa

Please sign in to comment.