feast/sdk/python/feast/proto_utils.py at featureview-versioning · feast-dev/feast · GitHub

History

101 lines (80 loc) · 3.07 KB

Raw

100

101

# 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.

"""

Utility functions for protobuf serialization of Feast objects.

"""

from typing import TYPE_CHECKING, Any, Optional, Union

from google.protobuf.message import Message

from feast.protos.feast.core.Transformation_pb2 import (

FeatureTransformationV2 as FeatureTransformationProto,

)

from feast.protos.feast.core.Transformation_pb2 import (

SubstraitTransformationV2 as SubstraitTransformationProto,

)

from feast.protos.feast.core.Transformation_pb2 import (

UserDefinedFunctionV2 as UserDefinedFunctionProto,

)

if TYPE_CHECKING:

from feast.data_source import DataSource

from feast.transformation.mode import TransformationMode

def serialize_data_source(source: Optional["DataSource"]) -> Optional[Message]:

"""Serialize a data source to proto with class type annotation.

Args:

source: The data source to serialize, or None.

Returns:

The serialized proto with data_source_class_type set, or None if source is None.

"""

if source is None:

return None

proto = source.to_proto()

proto.data_source_class_type = (

f"{source.__class__.__module__}.{source.__class__.__name__}"

)

return proto

def transformation_to_proto(

transformation: Optional[Any],

) -> Optional[FeatureTransformationProto]:

"""Convert a transformation to FeatureTransformationProto.

Args:

transformation: The transformation object with a to_proto() method.

Returns:

A FeatureTransformationProto wrapping the transformation, or None.

"""

if transformation is None:

return None

if not hasattr(transformation, "to_proto"):

return None

transformation_proto = transformation.to_proto()

if isinstance(transformation_proto, UserDefinedFunctionProto):

return FeatureTransformationProto(

user_defined_function=transformation_proto,

)

elif isinstance(transformation_proto, SubstraitTransformationProto):

return FeatureTransformationProto(

substrait_transformation=transformation_proto,

)

return None

def mode_to_string(mode: Optional[Union["TransformationMode", str]]) -> str:

"""Convert mode to string value.

Args:

mode: A TransformationMode enum or string, or None.

Returns:

The string representation of the mode, or empty string if None.

"""

from feast.transformation.mode import TransformationMode

if mode is None:

return ""

if isinstance(mode, TransformationMode):

return mode.value

return mode

feast/sdk/python/feast/proto_utils.py at featureview-versioning · feast-dev/feast · GitHub

Uh oh!

Expand file tree

Breadcrumbs

proto_utils.py

Latest commit

History

Breadcrumbs

proto_utils.py

File metadata and controls