diff --git a/sdk/python/feast/feature.py b/sdk/python/feast/feature.py index b37e0f562b..57f75c90d7 100644 --- a/sdk/python/feast/feature.py +++ b/sdk/python/feast/feature.py @@ -15,7 +15,7 @@ from typing import Dict, Optional from feast.protos.feast.core.Feature_pb2 import FeatureSpecV2 as FeatureSpecProto -from feast.protos.feast.types import Value_pb2 as ValueTypeProto +from feast.protos.feast.types.Value_pb2 import ValueType as ValueTypeProto from feast.value_type import ValueType @@ -88,7 +88,7 @@ def to_proto(self) -> FeatureSpecProto: Returns: A FeatureSpecProto protobuf. """ - value_type = ValueTypeProto.ValueType.Enum.Value(self.dtype.name) + value_type = ValueTypeProto.Enum.Value(self.dtype.name) return FeatureSpecProto( name=self.name, value_type=value_type, labels=self.labels, diff --git a/sdk/python/feast/types.py b/sdk/python/feast/types.py new file mode 100644 index 0000000000..fe74bd38bd --- /dev/null +++ b/sdk/python/feast/types.py @@ -0,0 +1,117 @@ +# Copyright 2022 The Feast Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from abc import ABC, abstractmethod +from enum import Enum +from typing import Union + +from feast.protos.feast.types.Value_pb2 import ValueType as ValueTypeProto + +PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES = { + "INVALID": "INVALID", + "STRING": "STRING", + "BYTES": "BYTES", + "BOOL": "BOOL", + "INT32": "INT32", + "INT64": "INT64", + "FLOAT32": "FLOAT", + "FLOAT64": "DOUBLE", + "UNIX_TIMESTAMP": "UNIX_TIMESTAMP", +} + + +class ComplexFeastType(ABC): + """ + A ComplexFeastType represents a structured type that is recognized by Feast. + """ + + def __init__(self): + """Creates a ComplexFeastType object.""" + pass + + @abstractmethod + def to_int(self) -> int: + """ + Converts a ComplexFeastType object to the appropriate int value corresponding to + the correct ValueTypeProto.Enum value. + """ + raise NotImplementedError + + +class PrimitiveFeastType(Enum): + """ + A PrimitiveFeastType represents a primitive type in Feast. + """ + + INVALID = 0 + BYTES = 1 + STRING = 2 + INT32 = 3 + INT64 = 4 + FLOAT32 = 5 + FLOAT64 = 6 + BOOL = 7 + UNIX_TIMESTAMP = 8 + + def to_int(self) -> int: + value_type_name = PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES[self.name] + return ValueTypeProto.Enum.Value(value_type_name) + + +Invalid = PrimitiveFeastType.INVALID +Bytes = PrimitiveFeastType.BYTES +String = PrimitiveFeastType.STRING +Bool = PrimitiveFeastType.BOOL +Int32 = PrimitiveFeastType.INT32 +Int64 = PrimitiveFeastType.INT64 +Float32 = PrimitiveFeastType.FLOAT32 +Float64 = PrimitiveFeastType.FLOAT64 +UnixTimestamp = PrimitiveFeastType.UNIX_TIMESTAMP + + +SUPPORTED_BASE_TYPES = [ + Invalid, + String, + Bytes, + Bool, + Int32, + Int64, + Float32, + Float64, + UnixTimestamp, +] + + +class Array(ComplexFeastType): + """ + An Array represents a list of types. + + Attributes: + base_type: The base type of the array. + """ + + base_type: Union[PrimitiveFeastType, ComplexFeastType] + + def __init__(self, base_type: Union[PrimitiveFeastType, ComplexFeastType]): + if base_type not in SUPPORTED_BASE_TYPES: + raise ValueError( + f"Type {type(base_type)} is currently not supported as a base type for Array." + ) + + self.base_type = base_type + + def to_int(self) -> int: + 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) diff --git a/sdk/python/tests/unit/test_types.py b/sdk/python/tests/unit/test_types.py new file mode 100644 index 0000000000..8252a0e181 --- /dev/null +++ b/sdk/python/tests/unit/test_types.py @@ -0,0 +1,23 @@ +import pytest + +from feast.protos.feast.types.Value_pb2 import ValueType as ValueTypeProto +from feast.types import Array, Float32, String + + +def test_primitive_feast_type(): + assert String.to_int() == ValueTypeProto.Enum.Value("STRING") + assert Float32.to_int() == ValueTypeProto.Enum.Value("FLOAT") + + +def test_array_feast_type(): + array_float_32 = Array(Float32) + assert array_float_32.to_int() == ValueTypeProto.Enum.Value("FLOAT_LIST") + + array_string = Array(String) + assert array_string.to_int() == ValueTypeProto.Enum.Value("STRING_LIST") + + with pytest.raises(ValueError): + _ = Array(Array) + + with pytest.raises(ValueError): + _ = Array(Array(String))