updated milvus

franciscojavierarceo · franciscojavierarceo · commit 2b18594e473a · 2024-12-18T08:36:16.000-05:00
Signed-off-by: Francisco Javier Arceo &lt;farceo@redhat.com&gt;
diff --git a/sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py b/sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py
@@ -25,7 +25,13 @@
 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
 from feast.repo_config import FeastConfigBaseModel, RepoConfig
 from feast.type_map import PROTO_VALUE_TO_VALUE_TYPE_MAP
-from feast.types import VALUE_TYPES_TO_FEAST_TYPES, Array, PrimitiveFeastType, ComplexFeastType, ValueType
+from feast.types import (
+    VALUE_TYPES_TO_FEAST_TYPES,
+    Array,
+    ComplexFeastType,
+    PrimitiveFeastType,
+    ValueType,
+)
 from feast.utils import (
     _build_retrieve_online_document_record,
     to_naive_utc,
@@ -46,7 +52,9 @@
     PROTO_VALUE_TO_VALUE_TYPE_MAP["bool_list_val"]: DataType.BINARY_VECTOR,
 }
 
-FEAST_PRIMITIVE_TO_MILVUS_TYPE_MAPPING: Dict[Union[PrimitiveFeastType, Array, ComplexFeastType], DataType] = {}
+FEAST_PRIMITIVE_TO_MILVUS_TYPE_MAPPING: Dict[
+    Union[PrimitiveFeastType, Array, ComplexFeastType], DataType
+] = {}
 
 for value_type, feast_type in VALUE_TYPES_TO_FEAST_TYPES.items():
     if isinstance(feast_type, PrimitiveFeastType):
@@ -267,6 +275,7 @@ def online_read(
                 results.append((None, None))
             else:
                 results.append((res_ts, res_dict))
+
         return results
 
     def update(
diff --git a/sdk/python/tests/integration/feature_repos/universal/online_store/milvus.py b/sdk/python/tests/integration/feature_repos/universal/online_store/milvus.py
@@ -32,4 +32,5 @@ def create_online_store(self) -> Dict[str, Any]:
         }
 
     def teardown(self):
+        # assert 1 == 4
         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
@@ -857,9 +857,7 @@ def assert_feature_service_entity_mapping_correctness(
 
 
 @pytest.mark.integration
-@pytest.mark.universal_online_stores(
-    only=["pgvector", "elasticsearch", "qdrant"]
-)
+@pytest.mark.universal_online_stores(only=["pgvector", "elasticsearch", "qdrant"])
 def test_retrieve_online_documents(vectordb_environment, fake_document_data):
     fs = vectordb_environment.feature_store
     df, data_source = fake_document_data
@@ -904,14 +902,13 @@ def test_retrieve_online_documents2(environment, fake_document_data):
     item_embeddings_feature_view = create_item_embeddings_feature_view(data_source)
     fs.apply([item_embeddings_feature_view, item()])
     fs.write_to_online_store("item_embeddings", df)
-    #
-    # documents = fs.retrieve_online_documents(
-    #     feature="item_embeddings:embedding_float",
-    #     query=[1.0, 2.0],
-    #     top_k=2,
-    #     distance_metric="L2",
-    # ).to_dict()
-    # assert len(documents["embedding_float"]) == 2
+    documents = fs.retrieve_online_documents(
+        feature="item_embeddings:embedding_float",
+        query=[1.0, 2.0],
+        top_k=2,
+        distance_metric="L2",
+    ).to_dict()
+    assert len(documents["embedding_float"]) == 2
     #
     # # assert returned the entity_id
     # assert len(documents["item_id"]) == 2

Original file line number	Diff line number	Diff line change
`@@ -32,4 +32,5 @@ def create_online_store(self) -> Dict[str, Any]:`
`32`	`32`	`}`
`33`	`33`
`34`	`34`	`def teardown(self):`
	`35`	`+ # assert 1 == 4`
`35`	`36`	`self.container.stop()`