Skip to content

Commit 08d6881

Browse files
judahrandjudahrand
authored
Return UNIX_TIMESTAMP as Python datetime (#2244)
* Refactor `UNIX_TIMESTAMP` conversion Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com> * Return `UNIX_TIMESTAMP` types as `datetime` to user Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com> * Fix linting errors Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com> * Rename variable to something more sensible Signed-off-by: Judah Rand <17158624+judahrand@users.noreply.github.com>
1 parent f6cc618 commit 08d6881

File tree

2 files changed

+57
-33
lines changed

2 files changed

+57
-33
lines changed

‎sdk/python/feast/type_map.py‎

Lines changed: 56 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,20 @@
1313
# limitations under the License.
1414

1515
import re
16-
from datetime import datetime
17-
from typing import Any, Dict, List, Optional, Set, Sized, Tuple, Type
16+
from datetime import datetime, timezone
17+
from typing import (
18+
Any,
19+
Dict,
20+
List,
21+
Optional,
22+
Sequence,
23+
Set,
24+
Sized,
25+
Tuple,
26+
Type,
27+
Union,
28+
cast,
29+
)
1830

1931
import numpy as np
2032
import pandas as pd
@@ -49,8 +61,17 @@ def feast_value_type_to_python_type(field_value_proto: ProtoValue) -> Any:
4961
if val_attr is None:
5062
return None
5163
val = getattr(field_value_proto, val_attr)
64+
65+
# If it's a _LIST type extract the list.
5266
if hasattr(val, "val"):
5367
val = list(val.val)
68+
69+
# Convert UNIX_TIMESTAMP values to `datetime`
70+
if val_attr == "unix_timestamp_list_val":
71+
val = [datetime.fromtimestamp(v, tz=timezone.utc) for v in val]
72+
elif val_attr == "unix_timestamp_val":
73+
val = datetime.fromtimestamp(val, tz=timezone.utc)
74+
5475
return val
5576

5677

@@ -240,6 +261,28 @@ def _type_err(item, dtype):
240261
}
241262

242263

264+
def _python_datetime_to_int_timestamp(
265+
values: Sequence[Any],
266+
) -> Sequence[Union[int, np.int_]]:
267+
# Fast path for Numpy array.
268+
if isinstance(values, np.ndarray) and isinstance(values.dtype, np.datetime64):
269+
if values.ndim != 1:
270+
raise ValueError("Only 1 dimensional arrays are supported.")
271+
return cast(Sequence[np.int_], values.astype("datetime64[s]").astype(np.int_))
272+
273+
int_timestamps = []
274+
for value in values:
275+
if isinstance(value, datetime):
276+
int_timestamps.append(int(value.timestamp()))
277+
elif isinstance(value, Timestamp):
278+
int_timestamps.append(int(value.ToSeconds()))
279+
elif isinstance(value, np.datetime64):
280+
int_timestamps.append(value.astype("datetime64[s]").astype(np.int_))
281+
else:
282+
int_timestamps.append(int(value))
283+
return int_timestamps
284+
285+
243286
def _python_value_to_proto_value(
244287
feast_value_type: ValueType, values: List[Any]
245288
) -> List[ProtoValue]:
@@ -275,22 +318,14 @@ def _python_value_to_proto_value(
275318
raise _type_err(first_invalid, valid_types[0])
276319

277320
if feast_value_type == ValueType.UNIX_TIMESTAMP_LIST:
278-
converted_values = []
279-
for value in values:
280-
converted_sub_values = []
281-
for sub_value in value:
282-
if isinstance(sub_value, datetime):
283-
converted_sub_values.append(int(sub_value.timestamp()))
284-
elif isinstance(sub_value, Timestamp):
285-
converted_sub_values.append(int(sub_value.ToSeconds()))
286-
elif isinstance(sub_value, np.datetime64):
287-
converted_sub_values.append(
288-
sub_value.astype("datetime64[s]").astype("int")
289-
)
290-
else:
291-
converted_sub_values.append(sub_value)
292-
converted_values.append(converted_sub_values)
293-
values = converted_values
321+
int_timestamps_lists = (
322+
_python_datetime_to_int_timestamp(value) for value in values
323+
)
324+
return [
325+
# ProtoValue does actually accept `np.int_` but the typing complains.
326+
ProtoValue(unix_timestamp_list_val=Int64List(val=ts)) # type: ignore
327+
for ts in int_timestamps_lists
328+
]
294329

295330
return [
296331
ProtoValue(**{field_name: proto_type(val=value)}) # type: ignore
@@ -302,20 +337,9 @@ def _python_value_to_proto_value(
302337
# Handle scalar types below
303338
else:
304339
if feast_value_type == ValueType.UNIX_TIMESTAMP:
305-
if isinstance(sample, datetime):
306-
return [
307-
ProtoValue(int64_val=int(value.timestamp())) for value in values
308-
]
309-
elif isinstance(sample, Timestamp):
310-
return [
311-
ProtoValue(int64_val=int(value.ToSeconds())) for value in values
312-
]
313-
elif isinstance(sample, np.datetime64):
314-
return [
315-
ProtoValue(int64_val=value.astype("datetime64[s]").astype("int"))
316-
for value in values
317-
]
318-
return [ProtoValue(int64_val=int(value)) for value in values]
340+
int_timestamps = _python_datetime_to_int_timestamp(values)
341+
# ProtoValue does actually accept `np.int_` but the typing complains.
342+
return [ProtoValue(unix_timestamp_val=ts) for ts in int_timestamps] # type: ignore
319343

320344
if feast_value_type in PYTHON_SCALAR_VALUE_TYPE_TO_PROTO_VALUE:
321345
(

‎sdk/python/tests/integration/registration/test_universal_types.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -234,7 +234,7 @@ def test_feature_get_online_features_types_match(online_types_test_fixtures):
234234
"float": float,
235235
"string": str,
236236
"bool": bool,
237-
"datetime": int,
237+
"datetime": datetime,
238238
}
239239
expected_dtype = feature_list_dtype_to_expected_online_response_value_type[
240240
config.feature_dtype

0 commit comments

Comments
 (0)