From 1e0fec72bd0c8a19cbd3be78626b9fa6e4bb2b88 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Wed, 23 Oct 2024 11:50:08 -0400 Subject: [PATCH] fix: Path faiss online return signature Signed-off-by: Francisco Javier Arceo --- .../contrib/faiss_online_store.py | 41 +++++++++++-------- 1 file changed, 23 insertions(+), 18 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/contrib/faiss_online_store.py b/sdk/python/feast/infra/online_stores/contrib/faiss_online_store.py index f69ca899d6..cc2e75800e 100644 --- a/sdk/python/feast/infra/online_stores/contrib/faiss_online_store.py +++ b/sdk/python/feast/infra/online_stores/contrib/faiss_online_store.py @@ -9,8 +9,8 @@ from feast import Entity, FeatureView, RepoConfig from feast.infra.key_encoding_utils import serialize_entity_key from feast.infra.online_stores.online_store import OnlineStore -from feast.protos.feast.types.EntityKey_pb2 import EntityKey -from feast.protos.feast.types.Value_pb2 import Value +from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto +from feast.protos.feast.types.Value_pb2 import Value as ValueProto from feast.repo_config import FeastConfigBaseModel @@ -94,9 +94,9 @@ def online_read( self, config: RepoConfig, table: FeatureView, - entity_keys: List[EntityKey], + entity_keys: List[EntityKeyProto], requested_features: Optional[List[str]] = None, - ) -> List[Tuple[Optional[datetime], Optional[Dict[str, Value]]]]: + ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]: if self._index is None: return [(None, None)] * len(entity_keys) @@ -111,7 +111,7 @@ def online_read( else: feature_vector = self._index.reconstruct(int(idx)) feature_dict = { - name: Value(double_val=value) + name: ValueProto(double_val=value) for name, value in zip( self._in_memory_store.feature_names, feature_vector ) @@ -123,7 +123,9 @@ def online_write_batch( self, config: RepoConfig, table: FeatureView, - data: List[Tuple[EntityKey, Dict[str, Value], datetime, Optional[datetime]]], + data: List[ + Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]] + ], progress: Optional[Callable[[int], Any]], ) -> None: if self._index is None: @@ -181,9 +183,10 @@ def retrieve_online_documents( ) -> List[ Tuple[ Optional[datetime], - Optional[Value], - Optional[Value], - Optional[Value], + Optional[EntityKeyProto], + Optional[ValueProto], + Optional[ValueProto], + Optional[ValueProto], ] ]: if self._index is None: @@ -196,9 +199,10 @@ def retrieve_online_documents( results: List[ Tuple[ Optional[datetime], - Optional[Value], - Optional[Value], - Optional[Value], + Optional[EntityKeyProto], + Optional[ValueProto], + Optional[ValueProto], + Optional[ValueProto], ] ] = [] for i, idx in enumerate(indices[0]): @@ -209,14 +213,15 @@ def retrieve_online_documents( timestamp = Timestamp() timestamp.GetCurrentTime() - - feature_value = Value(string_val=",".join(map(str, feature_vector))) - vector_value = Value(string_val=",".join(map(str, feature_vector))) - distance_value = Value(float_val=distances[0][i]) + entity_value = EntityKeyProto() + feature_value = ValueProto(string_val=",".join(map(str, feature_vector))) + vector_value = ValueProto(string_val=",".join(map(str, feature_vector))) + distance_value = ValueProto(float_val=distances[0][i]) results.append( ( timestamp.ToDatetime(), + entity_value, feature_value, vector_value, distance_value, @@ -229,8 +234,8 @@ async def online_read_async( self, config: RepoConfig, table: FeatureView, - entity_keys: List[EntityKey], + entity_keys: List[EntityKeyProto], requested_features: Optional[List[str]] = None, - ) -> List[Tuple[Optional[datetime], Optional[Dict[str, Value]]]]: + ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]: # Implement async read if needed raise NotImplementedError("Async read is not implemented for FaissOnlineStore")