Skip to content

Commit 80b4b37

Browse files
franciscojavierarceoemgeee
authored andcommitted
fix: Patch FAISS online return signature (feast-dev#4671)
fix: Path faiss online return signature Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
1 parent b492941 commit 80b4b37

File tree

1 file changed

+23
-18
lines changed

1 file changed

+23
-18
lines changed

sdk/python/feast/infra/online_stores/contrib/faiss_online_store.py

+23-18
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@
99
from feast import Entity, FeatureView, RepoConfig
1010
from feast.infra.key_encoding_utils import serialize_entity_key
1111
from feast.infra.online_stores.online_store import OnlineStore
12-
from feast.protos.feast.types.EntityKey_pb2 import EntityKey
13-
from feast.protos.feast.types.Value_pb2 import Value
12+
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
13+
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
1414
from feast.repo_config import FeastConfigBaseModel
1515

1616

@@ -94,9 +94,9 @@ def online_read(
9494
self,
9595
config: RepoConfig,
9696
table: FeatureView,
97-
entity_keys: List[EntityKey],
97+
entity_keys: List[EntityKeyProto],
9898
requested_features: Optional[List[str]] = None,
99-
) -> List[Tuple[Optional[datetime], Optional[Dict[str, Value]]]]:
99+
) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
100100
if self._index is None:
101101
return [(None, None)] * len(entity_keys)
102102

@@ -111,7 +111,7 @@ def online_read(
111111
else:
112112
feature_vector = self._index.reconstruct(int(idx))
113113
feature_dict = {
114-
name: Value(double_val=value)
114+
name: ValueProto(double_val=value)
115115
for name, value in zip(
116116
self._in_memory_store.feature_names, feature_vector
117117
)
@@ -123,7 +123,9 @@ def online_write_batch(
123123
self,
124124
config: RepoConfig,
125125
table: FeatureView,
126-
data: List[Tuple[EntityKey, Dict[str, Value], datetime, Optional[datetime]]],
126+
data: List[
127+
Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
128+
],
127129
progress: Optional[Callable[[int], Any]],
128130
) -> None:
129131
if self._index is None:
@@ -181,9 +183,10 @@ def retrieve_online_documents(
181183
) -> List[
182184
Tuple[
183185
Optional[datetime],
184-
Optional[Value],
185-
Optional[Value],
186-
Optional[Value],
186+
Optional[EntityKeyProto],
187+
Optional[ValueProto],
188+
Optional[ValueProto],
189+
Optional[ValueProto],
187190
]
188191
]:
189192
if self._index is None:
@@ -196,9 +199,10 @@ def retrieve_online_documents(
196199
results: List[
197200
Tuple[
198201
Optional[datetime],
199-
Optional[Value],
200-
Optional[Value],
201-
Optional[Value],
202+
Optional[EntityKeyProto],
203+
Optional[ValueProto],
204+
Optional[ValueProto],
205+
Optional[ValueProto],
202206
]
203207
] = []
204208
for i, idx in enumerate(indices[0]):
@@ -209,14 +213,15 @@ def retrieve_online_documents(
209213

210214
timestamp = Timestamp()
211215
timestamp.GetCurrentTime()
212-
213-
feature_value = Value(string_val=",".join(map(str, feature_vector)))
214-
vector_value = Value(string_val=",".join(map(str, feature_vector)))
215-
distance_value = Value(float_val=distances[0][i])
216+
entity_value = EntityKeyProto()
217+
feature_value = ValueProto(string_val=",".join(map(str, feature_vector)))
218+
vector_value = ValueProto(string_val=",".join(map(str, feature_vector)))
219+
distance_value = ValueProto(float_val=distances[0][i])
216220

217221
results.append(
218222
(
219223
timestamp.ToDatetime(),
224+
entity_value,
220225
feature_value,
221226
vector_value,
222227
distance_value,
@@ -229,8 +234,8 @@ async def online_read_async(
229234
self,
230235
config: RepoConfig,
231236
table: FeatureView,
232-
entity_keys: List[EntityKey],
237+
entity_keys: List[EntityKeyProto],
233238
requested_features: Optional[List[str]] = None,
234-
) -> List[Tuple[Optional[datetime], Optional[Dict[str, Value]]]]:
239+
) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
235240
# Implement async read if needed
236241
raise NotImplementedError("Async read is not implemented for FaissOnlineStore")

0 commit comments

Comments
 (0)