import uuid
from typing import Any, Callable, Type
from google.protobuf.json_format import ( # type: ignore
_WKTJSONMETHODS,
ParseError,
_Parser,
_Printer,
)
from importlib_metadata import version as importlib_version
from packaging import version
from feast.protos.feast.serving.ServingService_pb2 import FeatureList
from feast.protos.feast.types.Value_pb2 import RepeatedValue, Value
ProtoMessage = Any
JsonObject = Any
def _patch_proto_json_encoding(
proto_type: Type[ProtoMessage],
to_json_object: Callable[[_Printer, ProtoMessage], JsonObject],
from_json_object: Callable[[_Parser, JsonObject, ProtoMessage], None],
) -> None:
"""Patch Protobuf JSON Encoder / Decoder for a desired Protobuf type with to_json & from_json methods."""
to_json_fn_name = "_" + uuid.uuid4().hex
from_json_fn_name = "_" + uuid.uuid4().hex
setattr(_Printer, to_json_fn_name, to_json_object)
setattr(_Parser, from_json_fn_name, from_json_object)
_WKTJSONMETHODS[proto_type.DESCRIPTOR.full_name] = [
to_json_fn_name,
from_json_fn_name,
]
def _patch_feast_value_json_encoding():
"""Patch Protobuf JSON Encoder / Decoder with a Feast Value type.
This allows encoding the proto object as a native type, without the dummy structural wrapper.
Here's a before example:
{
"value_1": {
"int64_val": 1
},
"value_2": {
"double_list_val": [1.0, 2.0, 3.0]
},
}
And here's an after example:
{
"value_1": 1,
"value_2": [1.0, 2.0, 3.0]
}
"""
def to_json_object(printer: _Printer, message: ProtoMessage) -> JsonObject:
which = message.WhichOneof("val")
# If the Value message is not set treat as null_value when serialize
# to JSON. The parse back result will be different from original message.
if which is None or which == "null_val":
return None
elif "_list_" in which:
value = list(getattr(message, which).val)
else:
value = getattr(message, which)
return value
def from_json_object(
parser: _Parser, value: JsonObject, message: ProtoMessage
) -> None:
if value is None:
message.null_val = 0
elif isinstance(value, bool):
message.bool_val = value
elif isinstance(value, str):
message.string_val = value
elif isinstance(value, int):
message.int64_val = value
elif isinstance(value, float):
message.double_val = value
elif isinstance(value, list):
if len(value) == 0:
# Clear will mark the struct as modified so it will be created even if there are no values
message.int64_list_val.Clear()
elif isinstance(value[0], bool):
message.bool_list_val.val.extend(value)
elif isinstance(value[0], str):
message.string_list_val.val.extend(value)
elif isinstance(value[0], (float, int, type(None))):
# Identify array as ints if all of the elements are ints
if all(isinstance(item, int) for item in value):
message.int64_list_val.val.extend(value)
# If any of the elements are floats or nulls, then parse it as a float array
else:
# Convert each null as NaN.
message.double_list_val.val.extend(
[item if item is not None else float("nan") for item in value]
)
else:
raise ParseError(
"Value {0} has unexpected type {1}.".format(
value[0], type(value[0])
)
)
else:
raise ParseError(
"Value {0} has unexpected type {1}.".format(value, type(value))
)
def from_json_object_updated(
parser: _Parser, value: JsonObject, message: ProtoMessage, path: str
):
from_json_object(parser, value, message)
# https://github.com/feast-dev/feast/issues/2484 Certain feast users need a higher version of protobuf but the
# parameters of `from_json_object` changes in feast 3.20.1. This change gives users flexibility to use earlier versions.
current_version = importlib_version("protobuf")
if version.parse(current_version) < version.parse("3.20"):
_patch_proto_json_encoding(Value, to_json_object, from_json_object)
else:
_patch_proto_json_encoding(Value, to_json_object, from_json_object_updated)
def _patch_feast_repeated_value_json_encoding():
"""Patch Protobuf JSON Encoder / Decoder with a Feast RepeatedValue type.
This allows list of lists without dummy field name "val".
Here's a before example:
{
"repeated_value": [
{"val": [1,2,3]},
{"val": [4,5,6]}
]
}
And here's an after example:
{
"repeated_value": [
[1,2,3],
[4,5,6]
]
}
"""
def to_json_object(printer: _Printer, message: ProtoMessage) -> JsonObject:
return [printer._MessageToJsonObject(item) for item in message.val]
def from_json_object_updated(
parser: _Parser, value: JsonObject, message: ProtoMessage, path: str
) -> None:
array = value if isinstance(value, list) else value["val"]
for item in array:
parser.ConvertMessage(item, message.val.add(), path)
def from_json_object(
parser: _Parser, value: JsonObject, message: ProtoMessage
) -> None:
array = value if isinstance(value, list) else value["val"]
for item in array:
parser.ConvertMessage(item, message.val.add())
# https://github.com/feast-dev/feast/issues/2484 Certain feast users need a higher version of protobuf but the
# parameters of `from_json_object` changes in feast 3.20.1. This change gives users flexibility to use earlier versions.
current_version = importlib_version("protobuf")
if version.parse(current_version) < version.parse("3.20"):
_patch_proto_json_encoding(RepeatedValue, to_json_object, from_json_object)
else:
_patch_proto_json_encoding(
RepeatedValue, to_json_object, from_json_object_updated
)
def _patch_feast_feature_list_json_encoding():
"""Patch Protobuf JSON Encoder / Decoder with a Feast FeatureList type.
This allows list of lists without dummy field name "features".
Here's a before example:
{
"feature_list": {
"features": [
"feature-1",
"feature-2",
"feature-3"
]
}
}
And here's an after example:
{
"feature_list": [
"feature-1",
"feature-2",
"feature-3"
]
}
"""
def to_json_object(printer: _Printer, message: ProtoMessage) -> JsonObject:
return list(message.val)
def from_json_object(
parser: _Parser, value: JsonObject, message: ProtoMessage
) -> None:
array = value if isinstance(value, list) else value["val"]
message.val.extend(array)
def from_json_object_updated(
parser: _Parser, value: JsonObject, message: ProtoMessage, path: str
) -> None:
from_json_object(parser, value, message)
# https://github.com/feast-dev/feast/issues/2484 Certain feast users need a higher version of protobuf but the
# parameters of `from_json_object` changes in feast 3.20.1. This change gives users flexibility to use earlier versions.
current_version = importlib_version("protobuf")
if version.parse(current_version) < version.parse("3.20"):
_patch_proto_json_encoding(FeatureList, to_json_object, from_json_object)
else:
_patch_proto_json_encoding(
FeatureList, to_json_object, from_json_object_updated
)
[docs]def patch():
"""Patch Protobuf JSON Encoder / Decoder with all desired Feast types."""
_patch_feast_value_json_encoding()
_patch_feast_repeated_value_json_encoding()
_patch_feast_feature_list_json_encoding()