Source code for feast.types

# 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 Dict, Union

from feast.value_type import ValueType

PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES = {
    "INVALID": "UNKNOWN",
    "BYTES": "BYTES",
    "STRING": "STRING",
    "INT32": "INT32",
    "INT64": "INT64",
    "FLOAT64": "DOUBLE",
    "FLOAT32": "FLOAT",
    "BOOL": "BOOL",
    "UNIX_TIMESTAMP": "UNIX_TIMESTAMP",
}


[docs]class ComplexFeastType(ABC): """ A ComplexFeastType represents a structured type that is recognized by Feast. """ def __init__(self): """Creates a ComplexFeastType object.""" pass
[docs] @abstractmethod def to_value_type(self) -> ValueType: """ Converts a ComplexFeastType object to the corresponding ValueType enum. """ raise NotImplementedError
def __hash__(self): return hash(self.to_value_type().value) def __eq__(self, other): return self.to_value_type() == other.to_value_type()
[docs]class PrimitiveFeastType(Enum): """ A PrimitiveFeastType represents a primitive type in Feast. Note that these values must match the values in /feast/protos/types/Value.proto. """ INVALID = 0 BYTES = 1 STRING = 2 INT32 = 3 INT64 = 4 FLOAT64 = 5 FLOAT32 = 6 BOOL = 7 UNIX_TIMESTAMP = 8
[docs] def to_value_type(self) -> ValueType: """ Converts a PrimitiveFeastType object to the corresponding ValueType enum. """ value_type_name = PRIMITIVE_FEAST_TYPES_TO_VALUE_TYPES[self.name] return ValueType[value_type_name]
def __str__(self): return PRIMITIVE_FEAST_TYPES_TO_STRING[self.name] def __eq__(self, other): if isinstance(other, PrimitiveFeastType): return self.value == other.value else: return False def __hash__(self): return hash((PRIMITIVE_FEAST_TYPES_TO_STRING[self.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, ] PRIMITIVE_FEAST_TYPES_TO_STRING = { "INVALID": "Invalid", "STRING": "String", "BYTES": "Bytes", "BOOL": "Bool", "INT32": "Int32", "INT64": "Int64", "FLOAT32": "Float32", "FLOAT64": "Float64", "UNIX_TIMESTAMP": "UnixTimestamp", }
[docs]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
[docs] 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 ValueType[value_type_list_name]
def __str__(self): return f"Array({self.base_type})"
FeastType = Union[ComplexFeastType, PrimitiveFeastType] 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), }
[docs]def from_value_type( value_type: ValueType, ) -> FeastType: """ Converts a ValueType enum to a Feast type. Args: value_type: The ValueType to be converted. Raises: ValueError: The conversion could not be performed. """ if value_type in VALUE_TYPES_TO_FEAST_TYPES: return VALUE_TYPES_TO_FEAST_TYPES[value_type] raise ValueError(f"Could not convert value type {value_type} to FeastType.")