add test

HaoXuAI · HaoXuAI · commit ba39f93e6fd5 · 2024-04-10T22:26:36.000-07:00
diff --git a/Makefile b/Makefile
@@ -200,7 +200,7 @@ test-python-universal-postgres-offline:
 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 \
+		PYTEST_PLUGINS=sdk.python.tests.integration.feature_repos.universal.online_store.postgres \
 		python -m pytest -n 8 --integration \
  			-k "not test_universal_cli and \
  				not test_go_feature_server and \
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -1698,12 +1698,12 @@ def retrieve_online_documents(
         top_k: int,
     ) -> OnlineResponse:
         """
-        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".
+                references must have format "feature_view:feature", e.g, "document_fv:document_embeddings".
             query: The query to retrieve the closest document features for.
             top_k: The number of closest document features to retrieve.
         """
diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py
@@ -267,7 +267,7 @@ def retrieve_online_documents(
         requested_feature: str,
         embedding: List[float],
         top_k: int,
-    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
+    ) -> List[Tuple[Optional[datetime], Optional[ValueProto]]]:
         """
 
         Args:
diff --git a/sdk/python/feast/infra/passthrough_provider.py b/sdk/python/feast/infra/passthrough_provider.py
@@ -47,7 +47,6 @@ def __init__(self, config: RepoConfig):
         self.repo_config = config
         self._offline_store = None
         self._online_store = None
-        self._document_store = None
         self._batch_engine: Optional[BatchMaterializationEngine] = None
 
     @property
diff --git a/sdk/python/tests/conftest.py b/sdk/python/tests/conftest.py
@@ -405,3 +405,15 @@ def fake_ingest_data():
         "created": [pd.Timestamp(datetime.utcnow()).round("ms")],
     }
     return pd.DataFrame(data)
+
+
+@pytest.fixture
+def fake_ingest_document_data():
+    """Fake document data to ingest into the feature store"""
+    data = {
+        "driver_id": [1],
+        "doc": [4, 5],
+        "event_timestamp": [pd.Timestamp(datetime.utcnow()).round("ms")],
+        "created": [pd.Timestamp(datetime.utcnow()).round("ms")],
+    }
+    return pd.DataFrame(data)
diff --git a/sdk/python/tests/integration/feature_repos/universal/online_store/postgres.py b/sdk/python/tests/integration/feature_repos/universal/online_store/postgres.py
@@ -0,0 +1,33 @@
+from typing import Dict
+
+from testcontainers.postgres import PostgresContainer
+
+from tests.integration.feature_repos.universal.online_store_creator import (
+    OnlineStoreCreator,
+)
+
+
+class PostgresOnlieStoreCreator(OnlineStoreCreator):
+    def __init__(self, project_name: str, **kwargs):
+        super().__init__(project_name)
+        self.container = (
+            PostgresContainer("postgres:latest", platform="linux/amd64")
+            .with_exposed_ports(5432)
+            .with_env("POSTGRES_USER", "root")
+            .with_env("POSTGRES_PASSWORD", "test")
+            .with_env("POSTGRES_DB", "test")
+        )
+
+    def create_online_store(self) -> Dict[str, str]:
+        self.container.start()
+        exposed_port = self.container.get_exposed_port(5432)
+        return {
+            "type": "postgres",
+            "user": "root",
+            "password": "test",
+            "database": "test",
+            "port": exposed_port,
+        }
+
+    def teardown(self):
+        self.container.stop()
diff --git a/sdk/python/tests/integration/online_store/test_universal_online.py b/sdk/python/tests/integration/online_store/test_universal_online.py
@@ -785,3 +785,28 @@ def assert_feature_service_entity_mapping_correctness(
                 entity_rows=entity_rows,
                 full_feature_names=full_feature_names,
             )
+
+
+@pytest.mark.integration
+@pytest.mark.universal_online_stores(only=["postgres"])
+def test_retrieve_online_documents(
+     environment, universal_data_sources, fake_ingest_document_data
+):
+    fs = environment.feature_store
+    entities, datasets, data_sources = universal_data_sources
+    driver_hourly_stats = create_driver_hourly_stats_feature_view(data_sources.driver)
+    driver_entity = driver()
+
+    # Register Feature View and Entity
+    fs.apply([driver_hourly_stats, driver_entity])
+
+    # directly ingest data into the Online Store
+    fs.write_to_online_store("document_fv", fake_ingest_document_data)
+
+    # retrieve the online documents
+    documents = fs.retrieve_online_documents(
+        feature="document_fv:doc",
+        query="[1, 2]",
+        top_k=5
+    )
+    assert len(documents) == 2
