Skip to content

Commit ba9f4ef

Browse files
feat: Enable other distance metrics for Vector DB and Update docs (#4170)
* updated PGVector docs Signed-off-by: Francisco Javier Arceo <[email protected]> * adding distance metric to arguments and defaulting to L2 Signed-off-by: Francisco Javier Arceo <[email protected]> * linter Signed-off-by: Francisco Javier Arceo <[email protected]> * testing other distance metric Signed-off-by: Francisco Javier Arceo <[email protected]> * updated default Signed-off-by: Francisco Javier Arceo <[email protected]> * linter Signed-off-by: Francisco Javier Arceo <[email protected]> * fixed some copy Signed-off-by: Francisco Javier Arceo <[email protected]> * updated Signed-off-by: Francisco Javier Arceo <[email protected]> --------- Signed-off-by: Francisco Javier Arceo <[email protected]>
1 parent 67bea4c commit ba9f4ef

8 files changed

Lines changed: 64 additions & 8 deletions

File tree

‎docs/reference/online-stores/postgres.md‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,10 +65,16 @@ To compare this set of functionality against other online stores, please see the
6565
6666
## PGVector
6767
The Postgres online store supports the use of [PGVector](https://github.com/pgvector/pgvector) for storing feature values.
68-
To enable PGVector, set `pgvector_enabled: true` in the online store configuration.
68+
To enable PGVector, set `pgvector_enabled: true` in the online store configuration.
69+
6970
The `vector_len` parameter can be used to specify the length of the vector. The default value is 512.
7071

71-
Then you can use `retrieve_online_documents` to retrieve the top k closest vectors to a query vector.
72+
Please make sure to follow the instructions in the repository, which, as the time of this writing, requires you to
73+
run `CREATE EXTENSION vector;` in the database.
74+
75+
76+
Then you can use `retrieve_online_documents` to retrieve the top k closest vectors to a query vector.
77+
For the Retrieval Augmented Generation (RAG) use-case, you have to embed the query prior to passing the query vector.
7278

7379
{% code title="python" %}
7480
```python

‎sdk/python/feast/feature_store.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1740,6 +1740,7 @@ def retrieve_online_documents(
17401740
feature: str,
17411741
query: Union[str, List[float]],
17421742
top_k: int,
1743+
distance_metric: str,
17431744
) -> OnlineResponse:
17441745
"""
17451746
Retrieves the top k closest document features. Note, embeddings are a subset of features.
@@ -1750,18 +1751,21 @@ def retrieve_online_documents(
17501751
references must have format "feature_view:feature", e.g, "document_fv:document_embeddings".
17511752
query: The query to retrieve the closest document features for.
17521753
top_k: The number of closest document features to retrieve.
1754+
distance_metric: The distance metric to use for retrieval.
17531755
"""
17541756
return self._retrieve_online_documents(
17551757
feature=feature,
17561758
query=query,
17571759
top_k=top_k,
1760+
distance_metric=distance_metric,
17581761
)
17591762

17601763
def _retrieve_online_documents(
17611764
self,
17621765
feature: str,
17631766
query: Union[str, List[float]],
17641767
top_k: int,
1768+
distance_metric: str = "L2",
17651769
):
17661770
if isinstance(query, str):
17671771
raise ValueError(
@@ -1783,6 +1787,7 @@ def _retrieve_online_documents(
17831787
requested_feature,
17841788
query,
17851789
top_k,
1790+
distance_metric,
17861791
)
17871792

17881793
# TODO Refactor to better way of populating result
@@ -2025,6 +2030,7 @@ def _retrieve_from_online_store(
20252030
requested_feature: str,
20262031
query: List[float],
20272032
top_k: int,
2033+
distance_metric: str,
20282034
) -> List[Tuple[Timestamp, "FieldStatus.ValueType", Value, Value, Value]]:
20292035
"""
20302036
Search and return document features from the online document store.
@@ -2035,6 +2041,7 @@ def _retrieve_from_online_store(
20352041
requested_feature=requested_feature,
20362042
query=query,
20372043
top_k=top_k,
2044+
distance_metric=distance_metric,
20382045
)
20392046

20402047
read_row_protos = []

‎sdk/python/feast/infra/online_stores/contrib/postgres.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,13 @@
2121
from feast.repo_config import RepoConfig
2222
from feast.usage import log_exceptions_and_usage
2323

24+
SUPPORTED_DISTANCE_METRICS_DICT = {
25+
"cosine": "<=>",
26+
"L1": "<+>",
27+
"L2": "<->",
28+
"inner_product": "<#>",
29+
}
30+
2431

2532
class PostgreSQLOnlineStoreConfig(PostgreSQLConfig):
2633
type: Literal["postgres"] = "postgres"
@@ -276,6 +283,7 @@ def retrieve_online_documents(
276283
requested_feature: str,
277284
embedding: List[float],
278285
top_k: int,
286+
distance_metric: str = "L2",
279287
) -> List[
280288
Tuple[
281289
Optional[datetime],
@@ -292,6 +300,7 @@ def retrieve_online_documents(
292300
requested_feature: The requested feature as the column to search
293301
embedding: The query embedding to search for
294302
top_k: The number of items to return
303+
distance_metric: The distance metric to use for the search.G
295304
Returns:
296305
List of tuples containing the event timestamp and the document feature
297306
@@ -303,6 +312,12 @@ def retrieve_online_documents(
303312
"pgvector is not enabled in the online store configuration"
304313
)
305314

315+
if distance_metric not in SUPPORTED_DISTANCE_METRICS_DICT:
316+
raise ValueError(
317+
f"Distance metric {distance_metric} is not supported. Supported distance metrics are {SUPPORTED_DISTANCE_METRICS_DICT.keys()}"
318+
)
319+
320+
distance_metric_sql = SUPPORTED_DISTANCE_METRICS_DICT[distance_metric]
306321
# Convert the embedding to a string to be used in postgres vector search
307322
query_embedding_str = f"[{','.join(str(el) for el in embedding)}]"
308323

@@ -327,13 +342,14 @@ def retrieve_online_documents(
327342
feature_name,
328343
value,
329344
vector_value,
330-
vector_value <-> %s as distance,
345+
vector_value {distance_metric_sql} %s as distance,
331346
event_ts FROM {table_name}
332347
WHERE feature_name = {feature_name}
333348
ORDER BY distance
334349
LIMIT {top_k};
335350
"""
336351
).format(
352+
distance_metric_sql=distance_metric_sql,
337353
table_name=sql.Identifier(table_name),
338354
feature_name=sql.Literal(requested_feature),
339355
top_k=sql.Literal(top_k),

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,7 @@ def retrieve_online_documents(
158158
table: The feature view whose feature values should be read.
159159
requested_feature: The name of the feature whose embeddings should be used for retrieval.
160160
embedding: The embeddings to use for retrieval.
161-
top_k: The number of nearest neighbors to retrieve.
161+
top_k: The number of documents to retrieve.
162162
163163
Returns:
164164
object: A list of top k closest documents to the specified embedding. Each item in the list is a tuple

‎sdk/python/feast/infra/passthrough_provider.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -196,12 +196,18 @@ def retrieve_online_documents(
196196
requested_feature: str,
197197
query: List[float],
198198
top_k: int,
199+
distance_metric: str,
199200
) -> List:
200201
set_usage_attribute("provider", self.__class__.__name__)
201202
result = []
202203
if self.online_store:
203204
result = self.online_store.retrieve_online_documents(
204-
config, table, requested_feature, query, top_k
205+
config,
206+
table,
207+
requested_feature,
208+
query,
209+
top_k,
210+
distance_metric,
205211
)
206212
return result
207213

‎sdk/python/feast/infra/provider.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -303,6 +303,7 @@ def retrieve_online_documents(
303303
requested_feature: str,
304304
query: List[float],
305305
top_k: int,
306+
distance_metric: str = "L2",
306307
) -> List[
307308
Tuple[
308309
Optional[datetime],
@@ -312,14 +313,14 @@ def retrieve_online_documents(
312313
]
313314
]:
314315
"""
315-
Searches for the top-k nearest neighbors of the given document in the online document store.
316+
Searches for the top-k most similar documents in the online document store.
316317
317318
Args:
318319
config: The config for the current feature store.
319320
table: The feature view whose embeddings should be searched.
320321
requested_feature: the requested document feature name.
321322
query: The query embedding to search for.
322-
top_k: The number of nearest neighbors to return.
323+
top_k: The number of documents to return.
323324
324325
Returns:
325326
A list of dictionaries, where each dictionary contains the document feature.

‎sdk/python/tests/foo_provider.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ def retrieve_online_documents(
111111
requested_feature: str,
112112
query: List[float],
113113
top_k: int,
114+
distance_metric: str,
114115
) -> List[
115116
Tuple[
116117
Optional[datetime],

‎sdk/python/tests/integration/online_store/test_universal_online.py‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -798,6 +798,25 @@ def test_retrieve_online_documents(environment, fake_document_data):
798798
fs.write_to_online_store("item_embeddings", df)
799799

800800
documents = fs.retrieve_online_documents(
801-
feature="item_embeddings:embedding_float", query=[1.0, 2.0], top_k=2
801+
feature="item_embeddings:embedding_float",
802+
query=[1.0, 2.0],
803+
top_k=2,
804+
distance_metric="L2",
802805
).to_dict()
803806
assert len(documents["embedding_float"]) == 2
807+
808+
documents = fs.retrieve_online_documents(
809+
feature="item_embeddings:embedding_float",
810+
query=[1.0, 2.0],
811+
top_k=2,
812+
distance_metric="L1",
813+
).to_dict()
814+
assert len(documents["embedding_float"]) == 2
815+
816+
with pytest.raises(ValueError):
817+
fs.retrieve_online_documents(
818+
feature="item_embeddings:embedding_float",
819+
query=[1.0, 2.0],
820+
top_k=2,
821+
distance_metric="wrong",
822+
).to_dict()

0 commit comments

Comments
 (0)