diff --git a/sdk/python/feast/field.py b/sdk/python/feast/field.py index 498130e2099..2198d8a3f1b 100644 --- a/sdk/python/feast/field.py +++ b/sdk/python/feast/field.py @@ -15,6 +15,7 @@ from feast.feature import Feature from feast.protos.feast.core.Feature_pb2 import FeatureSpecV2 as FieldProto from feast.types import FeastType, from_value_type +from feast.value_type import ValueType class Field: @@ -61,7 +62,8 @@ def __str__(self): def to_proto(self) -> FieldProto: """Converts a Field object to its protobuf representation.""" - return FieldProto(name=self.name, value_type=self.dtype.to_value_type()) + value_type = self.dtype.to_value_type() + return FieldProto(name=self.name, value_type=value_type.value) @classmethod def from_proto(cls, field_proto: FieldProto): @@ -71,7 +73,8 @@ def from_proto(cls, field_proto: FieldProto): Args: field_proto: FieldProto protobuf object """ - return cls(name=field_proto.name, dtype=from_value_type(field_proto.value_type)) + value_type = ValueType(field_proto.value_type) + return cls(name=field_proto.name, dtype=from_value_type(value_type)) @classmethod def from_feature(cls, feature: Feature): @@ -81,4 +84,4 @@ def from_feature(cls, feature: Feature): Args: feature: Feature object to convert. """ - return cls(name=feature.name, dtype=from_value_type(feature.dtype.value)) + return cls(name=feature.name, dtype=from_value_type(feature.dtype)) diff --git a/sdk/python/feast/types.py b/sdk/python/feast/types.py index ebe02e31ad1..29a49a72c08 100644 --- a/sdk/python/feast/types.py +++ b/sdk/python/feast/types.py @@ -15,17 +15,17 @@ from enum import Enum from typing import Dict, Union -from feast.protos.feast.types.Value_pb2 import ValueType as ValueTypeProto +from feast.value_type import ValueType PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES = { - "INVALID": "INVALID", - "STRING": "STRING", + "INVALID": "UNKNOWN", "BYTES": "BYTES", - "BOOL": "BOOL", + "STRING": "STRING", "INT32": "INT32", "INT64": "INT64", - "FLOAT32": "FLOAT", "FLOAT64": "DOUBLE", + "FLOAT32": "FLOAT", + "BOOL": "BOOL", "UNIX_TIMESTAMP": "UNIX_TIMESTAMP", } @@ -40,14 +40,14 @@ def __init__(self): pass @abstractmethod - def to_value_type(self) -> ValueTypeProto.Enum: + def to_value_type(self) -> ValueType: """ - Converts a ComplexFeastType object to the corresponding ValueTypeProto.Enum value. + Converts a ComplexFeastType object to the corresponding ValueType enum. """ raise NotImplementedError def __hash__(self): - return hash(self.to_value_type()) + return hash(self.to_value_type().value) def __eq__(self, other): return self.to_value_type() == other.to_value_type() @@ -57,8 +57,7 @@ class PrimitiveFeastType(Enum): """ A PrimitiveFeastType represents a primitive type in Feast. - Note that these values must match the values in ValueTypeProto.Enum. See - /feast/protos/types/Value.proto for the exact values. + Note that these values must match the values in /feast/protos/types/Value.proto. """ INVALID = 0 @@ -71,12 +70,12 @@ class PrimitiveFeastType(Enum): BOOL = 7 UNIX_TIMESTAMP = 8 - def to_value_type(self) -> ValueTypeProto.Enum: + def to_value_type(self) -> ValueType: """ - Converts a PrimitiveFeastType object to the corresponding ValueTypeProto.Enum value. + Converts a PrimitiveFeastType object to the corresponding ValueType enum. """ value_type_name = PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES[self.name] - return ValueTypeProto.Enum.Value(value_type_name) + return ValueType[value_type_name] def __str__(self): return PRIMITIVE_FEAST_TYPES_TO_STRING[self.name] @@ -136,11 +135,11 @@ def __init__(self, base_type: Union[PrimitiveFeastType, ComplexFeastType]): self.base_type = base_type - def to_value_type(self) -> int: + def to_value_type(self) -> ValueType: assert isinstance(self.base_type, PrimitiveFeastType) value_type_name = PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES[self.base_type.name] value_type_list_name = value_type_name + "_LIST" - return ValueTypeProto.Enum.Value(value_type_list_name) + return ValueType[value_type_list_name] def __str__(self): return f"Array({self.base_type})" @@ -149,33 +148,33 @@ def __str__(self): FeastType = Union[ComplexFeastType, PrimitiveFeastType] -VALUE_TYPES_TO_FEAST_TYPES: Dict["ValueTypeProto.Enum", FeastType] = { - ValueTypeProto.Enum.INVALID: Invalid, - ValueTypeProto.Enum.BYTES: Bytes, - ValueTypeProto.Enum.STRING: String, - ValueTypeProto.Enum.INT32: Int32, - ValueTypeProto.Enum.INT64: Int64, - ValueTypeProto.Enum.DOUBLE: Float64, - ValueTypeProto.Enum.FLOAT: Float32, - ValueTypeProto.Enum.BOOL: Bool, - ValueTypeProto.Enum.UNIX_TIMESTAMP: UnixTimestamp, - ValueTypeProto.Enum.BYTES_LIST: Array(Bytes), - ValueTypeProto.Enum.STRING_LIST: Array(String), - ValueTypeProto.Enum.INT32_LIST: Array(Int32), - ValueTypeProto.Enum.INT64_LIST: Array(Int64), - ValueTypeProto.Enum.DOUBLE_LIST: Array(Float64), - ValueTypeProto.Enum.FLOAT_LIST: Array(Float32), - ValueTypeProto.Enum.BOOL_LIST: Array(Bool), - ValueTypeProto.Enum.UNIX_TIMESTAMP_LIST: Array(UnixTimestamp), +VALUE_TYPES_TO_FEAST_TYPES: Dict["ValueType", FeastType] = { + ValueType.UNKNOWN: Invalid, + ValueType.BYTES: Bytes, + ValueType.STRING: String, + ValueType.INT32: Int32, + ValueType.INT64: Int64, + ValueType.DOUBLE: Float64, + ValueType.FLOAT: Float32, + ValueType.BOOL: Bool, + ValueType.UNIX_TIMESTAMP: UnixTimestamp, + ValueType.BYTES_LIST: Array(Bytes), + ValueType.STRING_LIST: Array(String), + ValueType.INT32_LIST: Array(Int32), + ValueType.INT64_LIST: Array(Int64), + ValueType.DOUBLE_LIST: Array(Float64), + ValueType.FLOAT_LIST: Array(Float32), + ValueType.BOOL_LIST: Array(Bool), + ValueType.UNIX_TIMESTAMP_LIST: Array(UnixTimestamp), } -def from_value_type(value_type: ValueTypeProto.Enum,) -> FeastType: +def from_value_type(value_type: ValueType,) -> FeastType: """ - Converts a ValueTypeProto.Enum to a Feast type. + Converts a ValueType enum to a Feast type. Args: - value_type: The ValueTypeProto.Enum to be converted. + value_type: The ValueType to be converted. Raises: ValueError: The conversion could not be performed. diff --git a/sdk/python/tests/unit/test_types.py b/sdk/python/tests/unit/test_types.py index 5a721737f68..af490b4f3a9 100644 --- a/sdk/python/tests/unit/test_types.py +++ b/sdk/python/tests/unit/test_types.py @@ -1,25 +1,25 @@ import pytest -from feast.protos.feast.types.Value_pb2 import ValueType as ValueTypeProto from feast.types import Array, Float32, String, from_value_type +from feast.value_type import ValueType def test_primitive_feast_type(): - assert String.to_value_type() == ValueTypeProto.Enum.Value("STRING") + assert String.to_value_type() == ValueType.STRING assert from_value_type(String.to_value_type()) == String - assert Float32.to_value_type() == ValueTypeProto.Enum.Value("FLOAT") + assert Float32.to_value_type() == ValueType.FLOAT assert from_value_type(Float32.to_value_type()) == Float32 def test_array_feast_type(): - array_float_32 = Array(Float32) - assert array_float_32.to_value_type() == ValueTypeProto.Enum.Value("FLOAT_LIST") - assert from_value_type(array_float_32.to_value_type()) == array_float_32 - array_string = Array(String) - assert array_string.to_value_type() == ValueTypeProto.Enum.Value("STRING_LIST") + assert array_string.to_value_type() == ValueType.STRING_LIST assert from_value_type(array_string.to_value_type()) == array_string + array_float_32 = Array(Float32) + assert array_float_32.to_value_type() == ValueType.FLOAT_LIST + assert from_value_type(array_float_32.to_value_type()) == array_float_32 + with pytest.raises(ValueError): _ = Array(Array) @@ -28,8 +28,7 @@ def test_array_feast_type(): def test_all_value_types(): - values = ValueTypeProto.Enum.values() - for value in values: + for value in ValueType: # We do not support the NULL type. - if value != ValueTypeProto.Enum.Value("NULL"): + if value != ValueType.NULL: assert from_value_type(value).to_value_type() == value