-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Expand file tree
/
Copy pathproto_utils.py
More file actions
101 lines (80 loc) · 3.07 KB
/
proto_utils.py
File metadata and controls
101 lines (80 loc) · 3.07 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
# Copyright 2019 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.
"""
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