diff --git a/sdk/python/feast/infra/online_stores/redis.py b/sdk/python/feast/infra/online_stores/redis.py index 9a4e908810d..aeeb540b910 100644 --- a/sdk/python/feast/infra/online_stores/redis.py +++ b/sdk/python/feast/infra/online_stores/redis.py @@ -372,14 +372,11 @@ def _convert_redis_values_to_protobuf( redis_values: List[List[ByteString]], feature_view: str, requested_features: List[str], - ): - result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = [] - for values in redis_values: - features = self._get_features_for_entity( - values, feature_view, requested_features - ) - result.append(features) - return result + ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]: + return [ + self._get_features_for_entity(values, feature_view, requested_features) + for values in redis_values + ] def online_read( self, @@ -445,21 +442,21 @@ def _get_features_for_entity( res_val = dict(zip(requested_features, values)) res_ts = Timestamp() - ts_val = res_val.pop(f"_ts:{feature_view}") + ts_key = f"_ts:{feature_view}" + ts_val = res_val.pop(ts_key) if ts_val: - res_ts.ParseFromString(bytes(ts_val)) + res_ts.ParseFromString(ts_val) - res = {} + res: Dict[str, ValueProto] = {} for feature_name, val_bin in res_val.items(): val = ValueProto() if val_bin: - val.ParseFromString(bytes(val_bin)) + val.ParseFromString(val_bin) res[feature_name] = val if not res: return None, None - else: - # reconstruct full timestamp including nanos - total_seconds = res_ts.seconds + res_ts.nanos / 1_000_000_000.0 - timestamp = datetime.fromtimestamp(total_seconds, tz=timezone.utc) - return timestamp, res + + total_seconds = res_ts.seconds + res_ts.nanos / 1_000_000_000.0 + timestamp = datetime.fromtimestamp(total_seconds, tz=timezone.utc) + return timestamp, res diff --git a/sdk/python/tests/unit/infra/online_store/test_redis.py b/sdk/python/tests/unit/infra/online_store/test_redis.py index 83c8d3d61e4..0d9f2cd8739 100644 --- a/sdk/python/tests/unit/infra/online_store/test_redis.py +++ b/sdk/python/tests/unit/infra/online_store/test_redis.py @@ -128,3 +128,127 @@ def test_get_features_for_entity(redis_online_store: RedisOnlineStore, feature_v assert "feature_view_1:feature_11" in features assert features["feature_view_1:feature_10"].int32_val == 1 assert features["feature_view_1:feature_11"].int32_val == 2 + + +def test_get_features_for_entity_with_memoryview( + redis_online_store: RedisOnlineStore, feature_view +): + """Test that _get_features_for_entity handles memoryview inputs correctly. + + Redis may return memoryview objects instead of bytes in some cases. + The optimized code should handle both without unnecessary conversions. + """ + requested_features = [ + "feature_view_1:feature_10", + "feature_view_1:feature_11", + "_ts:feature_view_1", + ] + # Create memoryview objects to simulate redis returning memoryview + val1_bytes = ValueProto(int32_val=100).SerializeToString() + val2_bytes = ValueProto(int32_val=200).SerializeToString() + ts_bytes = Timestamp(seconds=1234567890, nanos=123456789).SerializeToString() + + values = [ + memoryview(val1_bytes), + memoryview(val2_bytes), + memoryview(ts_bytes), + ] + + timestamp, features = redis_online_store._get_features_for_entity( + values=values, + feature_view=feature_view.name, + requested_features=requested_features, + ) + assert features["feature_view_1:feature_10"].int32_val == 100 + assert features["feature_view_1:feature_11"].int32_val == 200 + assert timestamp is not None + + +def test_get_features_for_entity_with_none_values( + redis_online_store: RedisOnlineStore, feature_view +): + """Test that _get_features_for_entity handles None values correctly.""" + requested_features = [ + "feature_view_1:feature_10", + "feature_view_1:feature_11", + "_ts:feature_view_1", + ] + values = [ + ValueProto(int32_val=1).SerializeToString(), + None, # Missing feature value + Timestamp().SerializeToString(), + ] + + timestamp, features = redis_online_store._get_features_for_entity( + values=values, + feature_view=feature_view.name, + requested_features=requested_features, + ) + assert features["feature_view_1:feature_10"].int32_val == 1 + # None value should result in empty ValueProto + assert features["feature_view_1:feature_11"].WhichOneof("val") is None + + +def test_convert_redis_values_to_protobuf_multiple_entities( + redis_online_store: RedisOnlineStore, feature_view +): + """Test batch conversion with multiple entities.""" + requested_features = [ + "feature_view_1:feature_10", + "feature_view_1:feature_11", + "_ts:feature_view_1", + ] + # Multiple entity values + values = [ + [ + ValueProto(int32_val=1).SerializeToString(), + ValueProto(int32_val=2).SerializeToString(), + Timestamp(seconds=1000).SerializeToString(), + ], + [ + ValueProto(int32_val=10).SerializeToString(), + ValueProto(int32_val=20).SerializeToString(), + Timestamp(seconds=2000).SerializeToString(), + ], + [ + ValueProto(int32_val=100).SerializeToString(), + ValueProto(int32_val=200).SerializeToString(), + Timestamp(seconds=3000).SerializeToString(), + ], + ] + + results = redis_online_store._convert_redis_values_to_protobuf( + redis_values=values, + feature_view=feature_view.name, + requested_features=requested_features, + ) + + assert len(results) == 3 + assert results[0][1]["feature_view_1:feature_10"].int32_val == 1 + assert results[1][1]["feature_view_1:feature_10"].int32_val == 10 + assert results[2][1]["feature_view_1:feature_10"].int32_val == 100 + + +def test_get_features_for_entity_with_all_none_values( + redis_online_store: RedisOnlineStore, feature_view +): + """Test that None feature values result in empty ValueProto objects.""" + requested_features = [ + "feature_view_1:feature_10", + "_ts:feature_view_1", + ] + # All None values except timestamp + values = [ + None, + Timestamp().SerializeToString(), + ] + + timestamp, features = redis_online_store._get_features_for_entity( + values=values, + feature_view=feature_view.name, + requested_features=requested_features, + ) + # Even with None value, an empty ValueProto is created + assert features is not None + assert "feature_view_1:feature_10" in features + assert features["feature_view_1:feature_10"].WhichOneof("val") is None