Skip to content

Commit 59dfdb8

Browse files
abhijeet-dhumalabhijeet-dhumal
authored
perf: Optimize protobuf parsing in Redis online store (feast-dev#6023)
* perf: Optimize protobuf parsing in Redis online store Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com> * test: add unit tests for Redis protobuf parsing optimization Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com> * fix: Remove redundant isinstance checks for Redis protobuf parsing Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com> --------- Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com>
1 parent 1dffd18 commit 59dfdb8

2 files changed

Lines changed: 138 additions & 17 deletions

File tree

‎sdk/python/feast/infra/online_stores/redis.py‎

Lines changed: 14 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -372,14 +372,11 @@ def _convert_redis_values_to_protobuf(
372372
redis_values: List[List[ByteString]],
373373
feature_view: str,
374374
requested_features: List[str],
375-
):
376-
result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []
377-
for values in redis_values:
378-
features = self._get_features_for_entity(
379-
values, feature_view, requested_features
380-
)
381-
result.append(features)
382-
return result
375+
) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
376+
return [
377+
self._get_features_for_entity(values, feature_view, requested_features)
378+
for values in redis_values
379+
]
383380

384381
def online_read(
385382
self,
@@ -445,21 +442,21 @@ def _get_features_for_entity(
445442
res_val = dict(zip(requested_features, values))
446443

447444
res_ts = Timestamp()
448-
ts_val = res_val.pop(f"_ts:{feature_view}")
445+
ts_key = f"_ts:{feature_view}"
446+
ts_val = res_val.pop(ts_key)
449447
if ts_val:
450-
res_ts.ParseFromString(bytes(ts_val))
448+
res_ts.ParseFromString(ts_val)
451449

452-
res = {}
450+
res: Dict[str, ValueProto] = {}
453451
for feature_name, val_bin in res_val.items():
454452
val = ValueProto()
455453
if val_bin:
456-
val.ParseFromString(bytes(val_bin))
454+
val.ParseFromString(val_bin)
457455
res[feature_name] = val
458456

459457
if not res:
460458
return None, None
461-
else:
462-
# reconstruct full timestamp including nanos
463-
total_seconds = res_ts.seconds + res_ts.nanos / 1_000_000_000.0
464-
timestamp = datetime.fromtimestamp(total_seconds, tz=timezone.utc)
465-
return timestamp, res
459+
460+
total_seconds = res_ts.seconds + res_ts.nanos / 1_000_000_000.0
461+
timestamp = datetime.fromtimestamp(total_seconds, tz=timezone.utc)
462+
return timestamp, res

‎sdk/python/tests/unit/infra/online_store/test_redis.py‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,3 +128,127 @@ def test_get_features_for_entity(redis_online_store: RedisOnlineStore, feature_v
128128
assert "feature_view_1:feature_11" in features
129129
assert features["feature_view_1:feature_10"].int32_val == 1
130130
assert features["feature_view_1:feature_11"].int32_val == 2
131+
132+
133+
def test_get_features_for_entity_with_memoryview(
134+
redis_online_store: RedisOnlineStore, feature_view
135+
):
136+
"""Test that _get_features_for_entity handles memoryview inputs correctly.
137+
138+
Redis may return memoryview objects instead of bytes in some cases.
139+
The optimized code should handle both without unnecessary conversions.
140+
"""
141+
requested_features = [
142+
"feature_view_1:feature_10",
143+
"feature_view_1:feature_11",
144+
"_ts:feature_view_1",
145+
]
146+
# Create memoryview objects to simulate redis returning memoryview
147+
val1_bytes = ValueProto(int32_val=100).SerializeToString()
148+
val2_bytes = ValueProto(int32_val=200).SerializeToString()
149+
ts_bytes = Timestamp(seconds=1234567890, nanos=123456789).SerializeToString()
150+
151+
values = [
152+
memoryview(val1_bytes),
153+
memoryview(val2_bytes),
154+
memoryview(ts_bytes),
155+
]
156+
157+
timestamp, features = redis_online_store._get_features_for_entity(
158+
values=values,
159+
feature_view=feature_view.name,
160+
requested_features=requested_features,
161+
)
162+
assert features["feature_view_1:feature_10"].int32_val == 100
163+
assert features["feature_view_1:feature_11"].int32_val == 200
164+
assert timestamp is not None
165+
166+
167+
def test_get_features_for_entity_with_none_values(
168+
redis_online_store: RedisOnlineStore, feature_view
169+
):
170+
"""Test that _get_features_for_entity handles None values correctly."""
171+
requested_features = [
172+
"feature_view_1:feature_10",
173+
"feature_view_1:feature_11",
174+
"_ts:feature_view_1",
175+
]
176+
values = [
177+
ValueProto(int32_val=1).SerializeToString(),
178+
None, # Missing feature value
179+
Timestamp().SerializeToString(),
180+
]
181+
182+
timestamp, features = redis_online_store._get_features_for_entity(
183+
values=values,
184+
feature_view=feature_view.name,
185+
requested_features=requested_features,
186+
)
187+
assert features["feature_view_1:feature_10"].int32_val == 1
188+
# None value should result in empty ValueProto
189+
assert features["feature_view_1:feature_11"].WhichOneof("val") is None
190+
191+
192+
def test_convert_redis_values_to_protobuf_multiple_entities(
193+
redis_online_store: RedisOnlineStore, feature_view
194+
):
195+
"""Test batch conversion with multiple entities."""
196+
requested_features = [
197+
"feature_view_1:feature_10",
198+
"feature_view_1:feature_11",
199+
"_ts:feature_view_1",
200+
]
201+
# Multiple entity values
202+
values = [
203+
[
204+
ValueProto(int32_val=1).SerializeToString(),
205+
ValueProto(int32_val=2).SerializeToString(),
206+
Timestamp(seconds=1000).SerializeToString(),
207+
],
208+
[
209+
ValueProto(int32_val=10).SerializeToString(),
210+
ValueProto(int32_val=20).SerializeToString(),
211+
Timestamp(seconds=2000).SerializeToString(),
212+
],
213+
[
214+
ValueProto(int32_val=100).SerializeToString(),
215+
ValueProto(int32_val=200).SerializeToString(),
216+
Timestamp(seconds=3000).SerializeToString(),
217+
],
218+
]
219+
220+
results = redis_online_store._convert_redis_values_to_protobuf(
221+
redis_values=values,
222+
feature_view=feature_view.name,
223+
requested_features=requested_features,
224+
)
225+
226+
assert len(results) == 3
227+
assert results[0][1]["feature_view_1:feature_10"].int32_val == 1
228+
assert results[1][1]["feature_view_1:feature_10"].int32_val == 10
229+
assert results[2][1]["feature_view_1:feature_10"].int32_val == 100
230+
231+
232+
def test_get_features_for_entity_with_all_none_values(
233+
redis_online_store: RedisOnlineStore, feature_view
234+
):
235+
"""Test that None feature values result in empty ValueProto objects."""
236+
requested_features = [
237+
"feature_view_1:feature_10",
238+
"_ts:feature_view_1",
239+
]
240+
# All None values except timestamp
241+
values = [
242+
None,
243+
Timestamp().SerializeToString(),
244+
]
245+
246+
timestamp, features = redis_online_store._get_features_for_entity(
247+
values=values,
248+
feature_view=feature_view.name,
249+
requested_features=requested_features,
250+
)
251+
# Even with None value, an empty ValueProto is created
252+
assert features is not None
253+
assert "feature_view_1:feature_10" in features
254+
assert features["feature_view_1:feature_10"].WhichOneof("val") is None

0 commit comments

Comments
 (0)