test: Fix for integration tests (feast-dev#5126)

* chore: Fixing Milvus integration test configuration

Signed-off-by: Francisco Javier Arceo <[email protected]>

* chore: Fix Milvus integration tests

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

* Make sure all entities have values

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

* Fixed test_retrieve_online_milvus_documents

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

---------

Signed-off-by: Francisco Javier Arceo <[email protected]>
Signed-off-by: ntkathole <[email protected]>
Co-authored-by: Francisco Javier Arceo <[email protected]>

Author: ntkathole

pyproject.toml

@@ -135,7 +135,7 @@ ci = [
     "pytest-mock==1.10.4",
     "pytest-env",
     "Sphinx>4.0.0,<7",
-    "testcontainers==4.8.2",
+    "testcontainers==4.9.0",
     "python-keycloak==4.2.2",
     "pre-commit<3.3.2",
     "assertpy==1.1",

sdk/python/feast/cli.py

@@ -884,6 +884,7 @@ def materialize_incremental_command(ctx: click.Context, end_ts: str, views: List
             "hazelcast",
             "ikv",
             "couchbase",
+            "milvus",
         ],
         case_sensitive=False,
     ),

sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py

@@ -88,7 +88,7 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
     """
 
     type: Literal["milvus"] = "milvus"
-    path: Optional[StrictStr] = "data/online_store.db"
+    path: Optional[StrictStr] = "online_store.db"
     host: Optional[StrictStr] = "localhost"
     port: Optional[int] = 19530
     index_type: Optional[str] = "FLAT"
@@ -245,6 +245,8 @@ def online_write_batch(
         collection = self._get_or_create_collection(config, table)
         vector_cols = [f.name for f in table.features if f.vector_index]
         entity_batch_to_insert = []
+        unique_entities: dict[str, dict[str, Any]] = {}
+        required_fields = {field["name"] for field in collection["fields"]}
         for entity_key, values_dict, timestamp, created_ts in data:
             # need to construct the composite primary key also need to handle the fact that entities are a list
             entity_key_str = serialize_entity_key(
@@ -278,12 +280,22 @@ def online_write_batch(
                 "created_ts": created_ts_int,
             }
             single_entity_record.update(values_dict)
-            entity_batch_to_insert.append(single_entity_record)
+            # Ensure all required fields exist, setting missing ones to empty strings
+            for field in required_fields:
+                if field not in single_entity_record:
+                    single_entity_record[field] = ""
+            # Store only the latest event timestamp per entity
+            if (
+                entity_key_str not in unique_entities
+                or unique_entities[entity_key_str]["event_ts"] < timestamp_int
+            ):
+                unique_entities[entity_key_str] = single_entity_record
 
             if progress:
                 progress(1)
 
-        self.client.insert(
+        entity_batch_to_insert = list(unique_entities.values())
+        self.client.upsert(
             collection_name=collection["collection_name"],
             data=entity_batch_to_insert,
         )
@@ -551,11 +563,16 @@ def retrieve_online_documents_v2(
                         field, PrimitiveFeastType.INVALID
                     ) in [
                         PrimitiveFeastType.STRING,
+                        PrimitiveFeastType.BYTES,
+                    ]:
+                        res[field] = ValueProto(string_val=str(field_value))
+                    elif entity_name_feast_primitive_type_map.get(
+                        field, PrimitiveFeastType.INVALID
+                    ) in [
                         PrimitiveFeastType.INT64,
                         PrimitiveFeastType.INT32,
-                        PrimitiveFeastType.BYTES,
                     ]:
-                        res[field] = ValueProto(string_val=field_value)
+                        res[field] = ValueProto(int64_val=int(field_value))
                     elif field == composite_key_name:
                         pass
                     elif isinstance(field_value, bytes):

sdk/python/feast/infra/online_stores/milvus_online_store/milvus_repo_configuration.py

@@ -1,6 +1,7 @@
 from tests.integration.feature_repos.integration_test_repo_config import (
     IntegrationTestRepoConfig,
 )
+from tests.integration.feature_repos.repo_configuration import MILVUS_CONFIG
 from tests.integration.feature_repos.universal.online_store.milvus import (
     MilvusOnlineStoreCreator,
 )
@@ -10,3 +11,5 @@
         online_store="milvus", online_store_creator=MilvusOnlineStoreCreator
     ),
 ]
+
+AVAILABLE_ONLINE_STORES = {"milvus": (MILVUS_CONFIG, MilvusOnlineStoreCreator)}

sdk/python/feast/repo_config.py

@@ -61,6 +61,7 @@
     "feast.infra.online_stores.contrib.elasticsearch.ElasticSearchOnlineStore": "feast.infra.online_stores.elasticsearch_online_store.ElasticSearchOnlineStore",
     "feast.infra.online_stores.contrib.singlestore_online_store.singlestore.SingleStoreOnlineStore": "feast.infra.online_stores.singlestore_online_store.singlestore.SingleStoreOnlineStore",
     "feast.infra.online_stores.contrib.qdrant.QdrantOnlineStore": "feast.infra.online_stores.cqdrant.QdrantOnlineStore",
+    "feast.infra.online_stores.contrib.milvus.MilvusOnlineStore": "feast.infra.online_stores.milvus.MilvusOnlineStore",
 }
 
 ONLINE_STORE_CLASS_FOR_TYPE = {

sdk/python/feast/templates/milvus/.gitignore

@@ -0,0 +1,45 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*.pyo
+*.pyd
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+*.egg-info/
+dist/
+build/
+.venv
+
+# Pytest
+.cache
+*.cover
+*.log
+.coverage
+nosetests.xml
+coverage.xml
+*.hypothesis/
+*.pytest_cache/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IDEs and Editors
+.vscode/
+.idea/
+*.swp
+*.swo
+*.sublime-workspace
+*.sublime-project
+
+# OS generated files
+.DS_Store
+Thumbs.db

sdk/python/feast/templates/milvus/README.md

@@ -0,0 +1,29 @@
+# Feast Quickstart
+If you haven't already, check out the quickstart guide on Feast's website (http://docs.feast.dev/quickstart), which 
+uses this repo. A quick view of what's in this repository's `feature_repo/` directory:
+
+* `data/` contains raw demo parquet data
+* `feature_repo/example_repo.py` contains demo feature definitions
+* `feature_repo/feature_store.yaml` contains a demo setup configuring where data sources are
+* `feature_repo/test_workflow.py` showcases how to run all key Feast commands, including defining, retrieving, and pushing features. 
+
+You can run the overall workflow with `python test_workflow.py`.
+
+## To move from this into a more production ready workflow:
+> See more details in [Running Feast in production](https://docs.feast.dev/how-to-guides/running-feast-in-production)
+
+1. First: you should start with a different Feast template, which delegates to a more scalable offline store. 
+   - For example, running `feast init -t gcp`
+   or `feast init -t aws` or `feast init -t snowflake`. 
+   - You can see your options if you run `feast init --help`.
+2. `feature_store.yaml` points to a local file as a registry. You'll want to setup a remote file (e.g. in S3/GCS) or a 
+SQL registry. See [registry docs](https://docs.feast.dev/getting-started/concepts/registry) for more details. 
+3. This example uses a file [offline store](https://docs.feast.dev/getting-started/components/offline-store) 
+   to generate training data. It does not scale. We recommend instead using a data warehouse such as BigQuery, 
+   Snowflake, Redshift. There is experimental support for Spark as well.
+4. Setup CI/CD + dev vs staging vs prod environments to automatically update the registry as you change Feast feature definitions. See [docs](https://docs.feast.dev/how-to-guides/running-feast-in-production#1.-automatically-deploying-changes-to-your-feature-definitions).
+5. (optional) Regularly scheduled materialization to power low latency feature retrieval (e.g. via Airflow). See [Batch data ingestion](https://docs.feast.dev/getting-started/concepts/data-ingestion#batch-data-ingestion)
+for more details.
+6. (optional) Deploy feature server instances with `feast serve` to expose endpoints to retrieve online features.
+   - See [Python feature server](https://docs.feast.dev/reference/feature-servers/python-feature-server) for details.
+   - Use cases can also directly call the Feast client to fetch features as per [Feature retrieval](https://docs.feast.dev/getting-started/concepts/feature-retrieval)

sdk/python/feast/templates/milvus/__init__.py

@@ -0,0 +1,37 @@
+from feast.file_utils import replace_str_in_file
+
+
+def bootstrap():
+    # Bootstrap() will automatically be called from the init_repo() during `feast init`
+
+    import pathlib
+    from datetime import datetime, timedelta
+
+    from feast.driver_test_data import create_driver_hourly_stats_df
+
+    repo_path = pathlib.Path(__file__).parent.absolute() / "feature_repo"
+    project_name = pathlib.Path(__file__).parent.absolute().name
+    data_path = repo_path / "data"
+    data_path.mkdir(exist_ok=True)
+
+    end_date = datetime.now().replace(microsecond=0, second=0, minute=0)
+    start_date = end_date - timedelta(days=15)
+
+    driver_entities = [1001, 1002, 1003, 1004, 1005]
+    driver_df = create_driver_hourly_stats_df(driver_entities, start_date, end_date)
+
+    driver_stats_path = data_path / "driver_stats.parquet"
+    driver_df.to_parquet(path=str(driver_stats_path), allow_truncated_timestamps=True)
+
+    example_py_file = repo_path / "example_repo.py"
+    replace_str_in_file(example_py_file, "%PROJECT_NAME%", str(project_name))
+    replace_str_in_file(
+        example_py_file, "%PARQUET_PATH%", str(driver_stats_path.relative_to(repo_path))
+    )
+    replace_str_in_file(
+        example_py_file, "%LOGGING_PATH%", str(data_path.relative_to(repo_path))
+    )
+
+
+if __name__ == "__main__":
+    bootstrap()