Removing the singular feature parameter for doc retrieval

jyejare · jyejare · commit d9446d1d62e4 · 2025-03-26T23:45:41.000+05:30
Signed-off-by: jyejare &lt;jyejare@redhat.com&gt;
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -1831,19 +1831,15 @@ async def get_online_features_async(
 
     def retrieve_online_documents(
         self,
-        feature: Optional[str],
         query: Union[str, List[float]],
         top_k: int,
-        features: Optional[List[str]] = None,
+        features: List[str],
         distance_metric: Optional[str] = "L2",
     ) -> OnlineResponse:
         """
         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_embeddings".
             features: The list of features that should be retrieved from the online store.
             query: The query to retrieve the closest document features for.
             top_k: The number of closest document features to retrieve.
@@ -1853,38 +1849,27 @@ def retrieve_online_documents(
             raise ValueError(
                 "Using embedding functionality is not supported for document retrieval. Please embed the query before calling retrieve_online_documents."
             )
-        feature_list: List[str] = (
-            features
-            if features is not None
-            else ([feature] if feature is not None else [])
-        )
 
         (
             available_feature_views,
             _,
         ) = utils._get_feature_views_to_use(
             registry=self._registry,
             project=self.project,
-            features=feature_list,
+            features=features,
             allow_cache=True,
             hide_dummy_entity=False,
         )
-        if feature_list:
-            feature_view_set = set()
-            for _feature in feature_list:
-                feature_view_name = _feature.split(":")[0]
-                feature_view = self.get_feature_view(feature_view_name)
-                feature_view_set.add(feature_view.name)
-            if len(feature_view_set) > 1:
-                raise ValueError(
-                    "Document retrieval only supports a single feature view."
-                )
-            requested_feature = (
-                feature.split(":")[1] if isinstance(feature, str) else feature
-            )
-            requested_features = [
-                f.split(":")[1] for f in feature_list if isinstance(f, str) and ":" in f
-            ]
+        feature_view_set = set()
+        for _feature in features:
+            feature_view_name = _feature.split(":")[0]
+            feature_view = self.get_feature_view(feature_view_name)
+            feature_view_set.add(feature_view.name)
+        if len(feature_view_set) > 1:
+            raise ValueError("Document retrieval only supports a single feature view.")
+        requested_features = [
+            f.split(":")[1] for f in features if isinstance(f, str) and ":" in f
+        ]
         requested_feature_view_name = list(feature_view_set)[0]
         for feature_view in available_feature_views:
             if feature_view.name == requested_feature_view_name:
@@ -1899,15 +1884,20 @@ def retrieve_online_documents(
         document_features = self._retrieve_from_online_store(
             provider,
             requested_feature_view,
-            requested_feature,
             requested_features,
             query,
             top_k,
             distance_metric,
         )
+
         # TODO currently not return the vector value since it is same as feature value, if embedding is supported,
         # the feature value can be raw text before embedded
-        entity_key_vals = [feature[1] for feature in document_features]
+        def _doc_feature(x):
+            return [feature[x] for feature in document_features]
+
+        entity_key_vals, document_feature_vals, document_feature_distance_vals = map(
+            _doc_feature, (1, 4, 5)
+        )
         join_key_values: Dict[str, List[ValueProto]] = {}
         for entity_key_val in entity_key_vals:
             if entity_key_val is not None:
@@ -1917,11 +1907,7 @@ def retrieve_online_documents(
                     if join_key not in join_key_values:
                         join_key_values[join_key] = []
                     join_key_values[join_key].append(entity_value)
-
-        document_feature_vals = [feature[4] for feature in document_features]
-        document_feature_distance_vals = [feature[5] for feature in document_features]
         online_features_response = GetOnlineFeaturesResponse(results=[])
-        requested_feature = requested_feature or requested_features[0]
         if vector_field_metadata := _get_feature_view_vector_field_metadata(
             requested_feature_view
         ):
@@ -1931,7 +1917,7 @@ def retrieve_online_documents(
             vector_field_name: document_feature_vals,
             "distance": document_feature_distance_vals,
         }
-        _requested_features = [_feature.split(":")[-1] for _feature in feature_list]
+        _requested_features = [_feature.split(":")[-1] for _feature in features]
         requested_features_data = {
             _feature: data[_feature]
             for _feature in _requested_features
@@ -2016,7 +2002,6 @@ def _retrieve_from_online_store(
         self,
         provider: Provider,
         table: FeatureView,
-        requested_feature: Optional[str],
         requested_features: Optional[List[str]],
         query: List[float],
         top_k: int,
@@ -2036,7 +2021,6 @@ def _retrieve_from_online_store(
         documents = provider.retrieve_online_documents(
             config=self.config,
             table=table,
-            requested_feature=requested_feature,
             requested_features=requested_features,
             query=query,
             top_k=top_k,
diff --git a/sdk/python/feast/infra/online_stores/elasticsearch_online_store/elasticsearch.py b/sdk/python/feast/infra/online_stores/elasticsearch_online_store/elasticsearch.py
@@ -213,8 +213,7 @@ def retrieve_online_documents(
         self,
         config: RepoConfig,
         table: FeatureView,
-        requested_feature: Optional[str],
-        requested_features: Optional[List[str]],
+        requested_features: List[str],
         embedding: List[float],
         top_k: int,
         *args,
diff --git a/sdk/python/feast/infra/online_stores/faiss_online_store.py b/sdk/python/feast/infra/online_stores/faiss_online_store.py
@@ -176,8 +176,7 @@ def retrieve_online_documents(
         self,
         config: RepoConfig,
         table: FeatureView,
-        requested_feature: Optional[str],
-        requested_featres: Optional[List[str]],
+        requested_featres: List[str],
         embedding: List[float],
         top_k: int,
         distance_metric: Optional[str] = None,
diff --git a/sdk/python/feast/infra/online_stores/online_store.py b/sdk/python/feast/infra/online_stores/online_store.py
@@ -392,8 +392,7 @@ def retrieve_online_documents(
         self,
         config: RepoConfig,
         table: FeatureView,
-        requested_feature: Optional[str],
-        requested_features: Optional[List[str]],
+        requested_features: List[str],
         embedding: List[float],
         top_k: int,
         distance_metric: Optional[str] = None,
@@ -413,7 +412,6 @@ def retrieve_online_documents(
             distance_metric: distance metric to use for retrieval.
             config: The config for the current feature store.
             table: The feature view whose feature values should be read.
-            requested_feature: The name of the feature whose embeddings should be used for retrieval.
             requested_features: The list of features whose embeddings should be used for retrieval.
             embedding: The embeddings to use for retrieval.
             top_k: The number of documents to retrieve.
@@ -423,10 +421,8 @@ def retrieve_online_documents(
             where the first item is the event timestamp for the row, and the second item is a dict of feature
             name to embeddings.
         """
-        if not requested_feature and not requested_features:
-            raise ValueError(
-                "Either requested_feature or requested_features must be specified"
-            )
+        if not requested_features:
+            raise ValueError("Requested_features must be specified")
         raise NotImplementedError(
             f"Online store {self.__class__.__name__} does not support online retrieval"
         )
diff --git a/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py b/sdk/python/feast/infra/online_stores/postgres_online_store/postgres.py
@@ -354,7 +354,6 @@ def retrieve_online_documents(
         self,
         config: RepoConfig,
         table: FeatureView,
-        requested_feature: Optional[str],
         requested_features: Optional[List[str]],
         embedding: List[float],
         top_k: int,
@@ -373,7 +372,6 @@ def retrieve_online_documents(
         Args:
             config: Feast configuration object
             table: FeatureView object as the table to search
-            requested_feature: The requested feature as the column to search
             requested_features: The list of features whose embeddings should be used for retrieval.
             embedding: The query embedding to search for
             top_k: The number of items to return
@@ -425,15 +423,13 @@ def retrieve_online_documents(
                         vector_value,
                         vector_value {distance_metric_sql} %s::vector as distance,
                         event_ts FROM {table_name}
-                    WHERE feature_name = {feature_name}
                     ORDER BY distance
                     LIMIT {top_k};
                     """
                 ).format(
                     distance_metric_sql=sql.SQL(distance_metric_sql),
                     table_name=sql.Identifier(table_name),
                     feature_names=required_feature_names,
-                    feature_name=sql.Literal(requested_feature),
                     top_k=sql.Literal(top_k),
                 ),
                 (embedding,),
diff --git a/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py b/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py
@@ -248,8 +248,7 @@ def retrieve_online_documents(
         self,
         config: RepoConfig,
         table: FeatureView,
-        requested_feature: Optional[str],
-        requested_features: Optional[List[str]],
+        requested_features: List[str],
         embedding: List[float],
         top_k: int,
         distance_metric: Optional[str] = "cosine",
diff --git a/sdk/python/feast/infra/online_stores/sqlite.py b/sdk/python/feast/infra/online_stores/sqlite.py
@@ -322,8 +322,7 @@ def retrieve_online_documents(
         self,
         config: RepoConfig,
         table: FeatureView,
-        requested_feature: Optional[str],
-        requested_featuers: Optional[List[str]],
+        requested_featuers: List[str],
         embedding: List[float],
         top_k: int,
         distance_metric: Optional[str] = None,
@@ -341,7 +340,7 @@ def retrieve_online_documents(
         Args:
             config: Feast configuration object
             table: FeatureView object as the table to search
-            requested_feature: The requested feature as the column to search
+            requested_features: The list of requested features to retrieve
             embedding: The query embedding to search for
             top_k: The number of items to return
         Returns:
diff --git a/sdk/python/feast/infra/passthrough_provider.py b/sdk/python/feast/infra/passthrough_provider.py
@@ -294,7 +294,6 @@ def retrieve_online_documents(
         self,
         config: RepoConfig,
         table: FeatureView,
-        requested_feature: Optional[str],
         requested_features: Optional[List[str]],
         query: List[float],
         top_k: int,
@@ -305,7 +304,6 @@ def retrieve_online_documents(
             result = self.online_store.retrieve_online_documents(
                 config,
                 table,
-                requested_feature,
                 requested_features,
                 query,
                 top_k,
diff --git a/sdk/python/feast/infra/provider.py b/sdk/python/feast/infra/provider.py
@@ -419,7 +419,6 @@ def retrieve_online_documents(
         self,
         config: RepoConfig,
         table: FeatureView,
-        requested_feature: Optional[str],
         requested_features: Optional[List[str]],
         query: List[float],
         top_k: int,
@@ -440,7 +439,6 @@ def retrieve_online_documents(
             distance_metric: distance metric to use for the search.
             config: The config for the current feature store.
             table: The feature view whose embeddings should be searched.
-            requested_feature: the requested document feature name.
             requested_features: the requested document feature names.
             query: The query embedding to search for.
             top_k: The number of documents to return.
