diff --git a/sdk/python/feast/type_map.py b/sdk/python/feast/type_map.py index ad3e273d37..bfc3726ef8 100644 --- a/sdk/python/feast/type_map.py +++ b/sdk/python/feast/type_map.py @@ -13,6 +13,7 @@ # limitations under the License. import json +import re from collections import defaultdict from datetime import datetime, timezone from typing import ( @@ -77,9 +78,11 @@ 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) - if v != NULL_TIMESTAMP_INT_VALUE - else None + ( + datetime.fromtimestamp(v, tz=timezone.utc) + if v != NULL_TIMESTAMP_INT_VALUE + else None + ) for v in val ] elif val_attr == "unix_timestamp_val": @@ -295,9 +298,11 @@ def _type_err(item, dtype): ValueType.INT32: ("int32_val", lambda x: int(x), None), ValueType.INT64: ( "int64_val", - lambda x: int(x.timestamp()) - if isinstance(x, pd._libs.tslibs.timestamps.Timestamp) - else int(x), + lambda x: ( + int(x.timestamp()) + if isinstance(x, pd._libs.tslibs.timestamps.Timestamp) + else int(x) + ), None, ), ValueType.FLOAT: ("float_val", lambda x: float(x), None), @@ -390,15 +395,21 @@ def _python_value_to_proto_value( if feast_value_type == ValueType.BOOL_LIST: # ProtoValue does not support conversion of np.bool_ so we need to convert it to support np.bool_. return [ - ProtoValue(**{field_name: proto_type(val=[bool(e) for e in value])}) # type: ignore - if value is not None - else ProtoValue() + ( + ProtoValue( + **{field_name: proto_type(val=[bool(e) for e in value])} + ) # type: ignore + if value is not None + else ProtoValue() + ) for value in values ] return [ - ProtoValue(**{field_name: proto_type(val=value)}) # type: ignore - if value is not None - else ProtoValue() + ( + ProtoValue(**{field_name: proto_type(val=value)}) # type: ignore + if value is not None + else ProtoValue() + ) for value in values ] @@ -433,15 +444,17 @@ def _python_value_to_proto_value( if feast_value_type == ValueType.BOOL: # ProtoValue does not support conversion of np.bool_ so we need to convert it to support np.bool_. return [ - ProtoValue( - **{ - field_name: func( - bool(value) if type(value) is np.bool_ else value # type: ignore - ) - } + ( + ProtoValue( + **{ + field_name: func( + bool(value) if type(value) is np.bool_ else value # type: ignore + ) + } + ) + if not pd.isnull(value) + else ProtoValue() ) - if not pd.isnull(value) - else ProtoValue() for value in values ] if feast_value_type in PYTHON_SCALAR_VALUE_TYPE_TO_PROTO_VALUE: @@ -752,7 +765,7 @@ def _non_empty_value(value: Any) -> bool: def spark_to_feast_value_type(spark_type_as_str: str) -> ValueType: # TODO not all spark types are convertible - # Current non-convertible types: interval, map, struct, structfield, decimal, binary + # Current non-convertible types: interval, map, struct, structfield, binary type_map: Dict[str, ValueType] = { "null": ValueType.UNKNOWN, "byte": ValueType.BYTES, @@ -762,6 +775,7 @@ def spark_to_feast_value_type(spark_type_as_str: str) -> ValueType: "bigint": ValueType.INT64, "long": ValueType.INT64, "double": ValueType.DOUBLE, + "decimal": ValueType.DOUBLE, "float": ValueType.FLOAT, "boolean": ValueType.BOOL, "timestamp": ValueType.UNIX_TIMESTAMP, @@ -774,6 +788,10 @@ def spark_to_feast_value_type(spark_type_as_str: str) -> ValueType: "array": ValueType.BOOL_LIST, "array": ValueType.UNIX_TIMESTAMP_LIST, } + decimal_regex_pattern = r"^decimal\([0-9]{1,2},[0-9]{1,2}\)$" + if re.match(decimal_regex_pattern, spark_type_as_str): + spark_type_as_str = "decimal" + # TODO: Find better way of doing this. if not isinstance(spark_type_as_str, str) or spark_type_as_str not in type_map: return ValueType.NULL