feat: Implement and enhance remote document retrieval functionality (#5487)

* Remotely retrieve the docs

Signed-off-by: jyejare <[email protected]>

* Feature server rectified for documents retrival

Signed-off-by: jyejare <[email protected]>

* Updates to remote retrival for response extraction

Signed-off-by: jyejare <[email protected]>

* Remote retrive online doc v2

Signed-off-by: jyejare <[email protected]>

* Unit tests for Remote docuemnts retrival

Signed-off-by: jyejare <[email protected]>

* Fixed checks and comments

Signed-off-by: jyejare <[email protected]>

---------

Signed-off-by: jyejare <[email protected]>

Author: jyejare

sdk/python/feast/feature_server.py

@@ -6,7 +6,7 @@
 import traceback
 from contextlib import asynccontextmanager
 from importlib import resources as importlib_resources
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import pandas as pd
 import psutil
@@ -86,10 +86,18 @@ class MaterializeIncrementalRequest(BaseModel):
 class GetOnlineFeaturesRequest(BaseModel):
     entities: Dict[str, List[Any]]
     feature_service: Optional[str] = None
-    features: Optional[List[str]] = None
+    features: List[str] = []
     full_feature_names: bool = False
-    query_embedding: Optional[List[float]] = None
+
+
+class GetOnlineDocumentsRequest(BaseModel):
+    feature_service: Optional[str] = None
+    features: List[str] = []
+    full_feature_names: bool = False
+    top_k: Optional[int] = None
+    query: Optional[List[float]] = None
     query_string: Optional[str] = None
+    api_version: Optional[int] = 1
 
 
 class ChatMessage(BaseModel):
@@ -110,7 +118,10 @@ class SaveDocumentRequest(BaseModel):
     data: dict
 
 
-def _get_features(request: GetOnlineFeaturesRequest, store: "feast.FeatureStore"):
+def _get_features(
+    request: Union[GetOnlineFeaturesRequest, GetOnlineDocumentsRequest],
+    store: "feast.FeatureStore",
+):
     if request.feature_service:
         feature_service = store.get_feature_service(
             request.feature_service, allow_cache=True
@@ -246,24 +257,26 @@ async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, An
         dependencies=[Depends(inject_user_details)],
     )
     async def retrieve_online_documents(
-        request: GetOnlineFeaturesRequest,
+        request: GetOnlineDocumentsRequest,
     ) -> Dict[str, Any]:
         logger.warning(
             "This endpoint is in alpha and will be moved to /get-online-features when stable."
         )
         # Initialize parameters for FeatureStore.retrieve_online_documents_v2(...) call
         features = await run_in_threadpool(_get_features, request, store)
 
-        read_params = dict(
-            features=features,
-            full_feature_names=request.full_feature_names,
-            query=request.query_embedding,
-            query_string=request.query_string,
-        )
+        read_params = dict(features=features, query=request.query, top_k=request.top_k)
+        if request.api_version == 2 and request.query_string is not None:
+            read_params["query_string"] = request.query_string
 
-        response = await run_in_threadpool(
-            lambda: store.retrieve_online_documents_v2(**read_params)  # type: ignore
-        )
+        if request.api_version == 2:
+            response = await run_in_threadpool(
+                lambda: store.retrieve_online_documents_v2(**read_params)  # type: ignore
+            )
+        else:
+            response = await run_in_threadpool(
+                lambda: store.retrieve_online_documents(**read_params)  # type: ignore
+            )
 
         # Convert the Protobuf object to JSON and return it
         response_dict = await run_in_threadpool(

sdk/python/feast/feature_store.py

@@ -2413,6 +2413,11 @@ def _retrieve_from_online_store_v2(
             output_len=output_len,
         )
 
+        utils._populate_result_rows_from_columnar(
+            online_features_response=online_features_response,
+            data=entity_key_dict,
+        )
+
         return OnlineResponse(online_features_response)
 
     def serve(