Skip to content

Commit 8fbda9f

Browse files
abhijeet-dhumalabhijeet-dhumal
committed
[perf] Replace MessageToDict with optimized custom dict builder
Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com>
1 parent 6f5203a commit 8fbda9f

3 files changed

Lines changed: 486 additions & 15 deletions

File tree

‎sdk/python/feast/feature_server.py‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,6 @@
4141
from fastapi.logger import logger
4242
from fastapi.responses import JSONResponse, ORJSONResponse
4343
from fastapi.staticfiles import StaticFiles
44-
from google.protobuf.json_format import MessageToDict
4544
from prometheus_client import Gauge, start_http_server
4645
from pydantic import BaseModel
4746

@@ -53,6 +52,7 @@
5352
FeastError,
5453
)
5554
from feast.feast_object import FeastObject
55+
from feast.feature_server_utils import response_to_dict_fast
5656
from feast.feature_view_utils import get_feature_view_from_feature_store
5757
from feast.permissions.action import WRITE, AuthzedAction
5858
from feast.permissions.security_manager import assert_permissions
@@ -341,13 +341,7 @@ async def get_online_features(request: GetOnlineFeaturesRequest) -> ORJSONRespon
341341
lambda: store.get_online_features(**read_params) # type: ignore
342342
)
343343

344-
# Convert Protobuf to dict, then use ORJSONResponse for faster JSON serialization
345-
response_dict = await run_in_threadpool(
346-
MessageToDict,
347-
response.proto,
348-
preserving_proto_field_name=True,
349-
float_precision=18,
350-
)
344+
response_dict = response_to_dict_fast(response.proto)
351345
return ORJSONResponse(content=response_dict)
352346

353347
@app.post(
@@ -376,13 +370,7 @@ async def retrieve_online_documents(
376370
lambda: store.retrieve_online_documents(**read_params) # type: ignore
377371
)
378372

379-
# Convert Protobuf to dict, then use ORJSONResponse for faster JSON serialization
380-
response_dict = await run_in_threadpool(
381-
MessageToDict,
382-
response.proto,
383-
preserving_proto_field_name=True,
384-
float_precision=18,
385-
)
373+
response_dict = response_to_dict_fast(response.proto)
386374
return ORJSONResponse(content=response_dict)
387375

388376
@app.post("/push", dependencies=[Depends(inject_user_details)])

‎sdk/python/feast/feature_server_utils.py‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
"""Fast serialization utilities for Feature Server responses."""
2+
3+
import base64
4+
from datetime import datetime, timezone
5+
from typing import Any, Dict
6+
7+
from feast.protos.feast.serving.ServingService_pb2 import GetOnlineFeaturesResponse
8+
from feast.protos.feast.types.Value_pb2 import Value
9+
10+
# FieldStatus enum mapping (protos/feast/serving/ServingService.proto)
11+
_STATUS_NAMES: Dict[int, str] = {
12+
0: "INVALID",
13+
1: "PRESENT",
14+
2: "NULL_VALUE",
15+
3: "NOT_FOUND",
16+
4: "OUTSIDE_MAX_AGE",
17+
}
18+
19+
20+
def response_to_dict_fast(response: GetOnlineFeaturesResponse) -> Dict[str, Any]:
21+
"""Convert GetOnlineFeaturesResponse to dict (optimized for this message type)."""
22+
result: Dict[str, Any] = {
23+
"results": [
24+
{
25+
"values": [_value_to_dict(v) for v in feature_vector.values],
26+
"statuses": [
27+
_STATUS_NAMES.get(s, "UNKNOWN") for s in feature_vector.statuses
28+
],
29+
"event_timestamps": [
30+
_timestamp_to_str(ts) for ts in feature_vector.event_timestamps
31+
]
32+
if feature_vector.event_timestamps
33+
else [],
34+
}
35+
for feature_vector in response.results
36+
]
37+
}
38+
39+
if response.HasField("metadata"):
40+
result["metadata"] = _metadata_to_dict(response.metadata)
41+
42+
return result
43+
44+
45+
def _value_to_dict(v: Value) -> Dict[str, Any]:
46+
"""Convert a Value proto to dict based on which field is set."""
47+
which = v.WhichOneof("val")
48+
if which is None:
49+
return {"null_val": 0}
50+
51+
if which == "double_val":
52+
return {"double_val": v.double_val}
53+
elif which == "float_val":
54+
return {"float_val": v.float_val}
55+
elif which == "int64_val":
56+
return {"int64_val": str(v.int64_val)}
57+
elif which == "int32_val":
58+
return {"int32_val": v.int32_val}
59+
elif which == "string_val":
60+
return {"string_val": v.string_val}
61+
elif which == "bool_val":
62+
return {"bool_val": v.bool_val}
63+
elif which == "bytes_val":
64+
return {"bytes_val": base64.b64encode(v.bytes_val).decode("ascii")}
65+
elif which == "double_list_val":
66+
return {"double_list_val": {"val": list(v.double_list_val.val)}}
67+
elif which == "float_list_val":
68+
return {"float_list_val": {"val": list(v.float_list_val.val)}}
69+
elif which == "int64_list_val":
70+
return {"int64_list_val": {"val": [str(x) for x in v.int64_list_val.val]}}
71+
elif which == "int32_list_val":
72+
return {"int32_list_val": {"val": list(v.int32_list_val.val)}}
73+
elif which == "string_list_val":
74+
return {"string_list_val": {"val": list(v.string_list_val.val)}}
75+
elif which == "bool_list_val":
76+
return {"bool_list_val": {"val": list(v.bool_list_val.val)}}
77+
elif which == "unix_timestamp_val":
78+
return {"unix_timestamp_val": v.unix_timestamp_val}
79+
elif which == "unix_timestamp_list_val":
80+
return {"unix_timestamp_list_val": {"val": list(v.unix_timestamp_list_val.val)}}
81+
82+
return {"null_val": 0}
83+
84+
85+
def _timestamp_to_str(ts) -> str:
86+
"""Convert protobuf Timestamp to ISO format string."""
87+
if ts.seconds == 0 and ts.nanos == 0:
88+
return ""
89+
dt = datetime.fromtimestamp(ts.seconds + ts.nanos / 1e9, tz=timezone.utc)
90+
return dt.isoformat()
91+
92+
93+
def _metadata_to_dict(metadata) -> Dict[str, Any]:
94+
"""Convert FeatureResponseMeta to dict."""
95+
result: Dict[str, Any] = {}
96+
if metadata.HasField("feature_names"):
97+
result["feature_names"] = {"val": list(metadata.feature_names.val)}
98+
return result

0 commit comments

Comments
 (0)