feat: Enable Vector database and retrieve_online_documents API - #4061
Conversation
franciscojavierarceo
left a comment
There was a problem hiding this comment.
Overall this looks great, I left some comments in the doc and think the only thing that's worth doing is us aligning on some of the naming conventions to be more similar to the industry and to what Feast is already doing.
Great work and can't wait to get this out! 🚀
| top_k: int, | ||
| ) -> OnlineResponse: | ||
| """ | ||
| Retrieves the top k cloeses document features. |
There was a problem hiding this comment.
| Retrieves the top k cloeses document features. | |
| Retrieves the top k closest document features. |
| def get_top_k_document_features( | ||
| self, | ||
| feature: str, | ||
| document: Union[str, np.ndarray], |
There was a problem hiding this comment.
is document the user input? often referred to as the query?
| top_k: int, | ||
| ) -> OnlineResponse: | ||
| """ | ||
| Retrieves the top k closest document features. |
There was a problem hiding this comment.
| Retrieves the top k closest document features. | |
| Retrieves the top k closest document features. Note, embeddings are a subset of features. |
| Args: | ||
| feature: The list of document features that should be retrieved from the online document store. These features can be | ||
| specified either as a list of string document feature references or as a feature service. String feature | ||
| references must have format "feature_view:feature", e.g, "document_fv:document_embedding_feature". |
There was a problem hiding this comment.
| references must have format "feature_view:feature", e.g, "document_fv:document_embedding_feature". | |
| references must have format "feature_view:feature", e.g, "document_fv:document_embeddings". |
| return OnlineResponse(online_features_response) | ||
|
|
||
| @log_exceptions_and_usage | ||
| def retrieve_online_documents( |
There was a problem hiding this comment.
There's probably something to be said about having a configurable distance metric to let the user choose which way to get the top_k
There was a problem hiding this comment.
yeah, there are a bunch of different algorithms/configs for Postgresql to retrieve the documents. We can support it in the future after this PR
| """ | ||
|
|
||
| # Convert the embedding to a string to be used in postgres vector search | ||
| query_embedding_str = f"'[{','.join(str(el) for el in embedding)}]'" |
There was a problem hiding this comment.
Is this the best serialization we can do? This feels pretty brittle but I get it.
| self.repo_config = config | ||
| self._offline_store = None | ||
| self._online_store = None | ||
| self._document_store = None |
There was a problem hiding this comment.
This is no longer necessary
There was a problem hiding this comment.
Thanks for addressing the issue. One question, though: what if we want to continue using Redis or any other online store for usual features, and use PG vector solely for embedding and search? Do we have the option to use the online store and the document store in the feature_store.yaml, both?
There was a problem hiding this comment.
I think currently Feast doesn't support multiple online store. but that would be a good feature to add.
There was a problem hiding this comment.
That could get complicated but agreed it'd be good to add. I could imagine a Redis + another DB layer would be an obvious one.
There was a problem hiding this comment.
Thank you again, guys, for these amazing features. Yes, having multiple online stores will make it easier to use the right database layer for the appropriate use case! 🙌
| requested_feature: str, | ||
| embedding: List[float], | ||
| top_k: int, | ||
| ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]: |
There was a problem hiding this comment.
Do we need Optional[Dict[str, ValueProto]]? If only a single feature can be searched, wouldn't Optional[ValueProto] be sufficient?
|
|
||
| @pytest.mark.integration | ||
| @pytest.mark.universal_online_stores(only=["postgres"]) | ||
| def test_retrieve_online_documents( |
There was a problem hiding this comment.
Will you be outputting the cosine similarity as well? That would be useful possibly for debugging. Would be good to be able to test that the engine computes it...maybe not doable though.
There was a problem hiding this comment.
Should be possible. Somehow just the integration test doesn't startup the Postgres container. And I'm debugging it.
thanks! |
|
@franciscojavierarceo @tokoko fyi this is the result of the API. there are still some TODOs I've added into the code directly. Will add to the document as well. Also feels like having a layer to abstract the 
|
| test-python-universal-postgres-online: | ||
| PYTHONPATH='.' \ | ||
| FULL_REPO_CONFIGS_MODULE=sdk.python.feast.infra.online_stores.contrib.postgres_repo_configuration \ | ||
| PYTEST_PLUGINS=sdk.python.feast.infra.offline_stores.contrib.postgres_offline_store.tests \ |
There was a problem hiding this comment.
Looking at this now, was this the right choice?
What this PR does / why we need it:
RFC: https://docs.google.com/document/d/18IWzLEA9i2lDWnbfbwXnMCg3StlqaLVI-uRpQjr_Vos/edit#heading=h.9gaqqtox9jg6
Which issue(s) this PR fixes:
Fixes #
#3965