Requested only features with values

jyejare · jyejare · commit d7f185f7a4ea · 2025-03-24T19:42:49.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
@@ -1869,39 +1869,31 @@ def retrieve_online_documents(
             allow_cache=True,
             hide_dummy_entity=False,
         )
-        if features:
+        if feature_list:
             feature_view_set = set()
-            for feature in features:
-                feature_view_name = feature.split(":")[0]
+            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 = None
-            requested_features = [
-                f.split(":")[1] for f in features if isinstance(f, str) and ":" in f
-            ]
-        else:
             requested_feature = (
                 feature.split(":")[1] if isinstance(feature, str) else feature
             )
-            requested_features = [requested_feature] if requested_feature else []
-
-        requested_feature_view_name = (
-            feature.split(":")[0] if feature else list(feature_view_set)[0]
-        )
+            requested_features = [
+                f.split(":")[1] for f in feature_list 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:
                 requested_feature_view = feature_view
                 break
-        if not requested_feature_view:
+        else:
             raise ValueError(
                 f"Feature view {requested_feature_view} not found in the registry."
             )
-        # Unnecessary code here a its overriding the requested_feature_view set from the above for loop
-        # requested_feature_view = available_feature_views[0]
 
         provider = self._get_provider()
         document_features = self._retrieve_from_online_store(
@@ -1945,7 +1937,6 @@ def retrieve_online_documents(
             for _feature in _requested_features
             if _feature in data
         }
-        # TODO currently the requested 'features' list is not functioning as expected
         utils._populate_result_rows_from_columnar(
             online_features_response=online_features_response,
             data=requested_features_data,
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
@@ -394,6 +394,11 @@ def retrieve_online_documents(
                 f"Distance metric {distance_metric} is not supported. Supported distance metrics are {SUPPORTED_DISTANCE_METRICS_DICT.keys()}"
             )
 
+        if requested_features:
+            required_feature_names = ", ".join(
+                [feature for feature in requested_features]
+            )
+
         distance_metric_sql = SUPPORTED_DISTANCE_METRICS_DICT[distance_metric]
 
         result: List[
@@ -415,7 +420,7 @@ def retrieve_online_documents(
                     """
                     SELECT
                         entity_key,
-                        feature_name,
+                        {feature_names},
                         value,
                         vector_value,
                         vector_value {distance_metric_sql} %s::vector as distance,
@@ -427,6 +432,7 @@ def retrieve_online_documents(
                 ).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),
                 ),
diff --git a/sdk/python/tests/unit/online_store/test_online_retrieval.py b/sdk/python/tests/unit/online_store/test_online_retrieval.py
@@ -718,7 +718,10 @@ def test_sqlite_get_online_documents() -> None:
             vector_length,
         )
         result = store.retrieve_online_documents(
-            feature="document_embeddings:Embeddings", query=query_embedding, top_k=3
+            feature="document_embeddings:item_id",
+            query=query_embedding,
+            top_k=3,
+            features=["document_embeddings:Embeddings", "document_embeddings:distance"],
         ).to_dict()
 
         assert "Embeddings" in result
