From 5eb7471b08068f4581ab3df991c439ee42e8f4c3 Mon Sep 17 00:00:00 2001 From: Alex Korbonits Date: Mon, 20 Apr 2026 11:06:13 -0700 Subject: [PATCH] test: expand MilvusOnlineStore integration test coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds three cases to test_retrieve_online_milvus_documents that exercise code paths adjacent to recent MilvusOnlineStore bug fixes (#6275): 1. Empty-store query before any writes — verifies the v2 retrieval path returns 0 rows instead of raising when the collection exists but is empty. 2. Oversized top_k (top_k=5 on a 3-row dataset) — verifies the hit-parsing loop handles a result set smaller than the requested top_k and returns all available rows. 3. Cosine-metric variant via a second FeatureView with vector_search_metric="COSINE" — verifies the cosine index path end to end (write, index creation, retrieval). No new imports; all types were already imported in the module. Co-Authored-By: Claude Opus 4.7 Signed-off-by: Alex Korbonits --- .../online_store/test_universal_online.py | 68 +++++++++++++++++-- 1 file changed, 63 insertions(+), 5 deletions(-) diff --git a/sdk/python/tests/integration/online_store/test_universal_online.py b/sdk/python/tests/integration/online_store/test_universal_online.py index 0c27585139e..f72ac64586c 100644 --- a/sdk/python/tests/integration/online_store/test_universal_online.py +++ b/sdk/python/tests/integration/online_store/test_universal_online.py @@ -921,13 +921,27 @@ def test_retrieve_online_milvus_documents(environment, fake_document_data): df, data_source = fake_document_data item_embeddings_feature_view = create_item_embeddings_feature_view(data_source) fs.apply([item_embeddings_feature_view, item()]) + + features = [ + "item_embeddings:embedding_float", + "item_embeddings:item_id", + "item_embeddings:string_feature", + ] + + # Empty-store query: collection exists but has no rows yet. + empty = fs.retrieve_online_documents_v2( + features=features, + query=[1.0, 2.0], + top_k=2, + distance_metric="L2", + ).to_dict() + assert len(empty["embedding_float"]) == 0 + assert len(empty["item_id"]) == 0 + fs.write_to_online_store("item_embeddings", df) + documents = fs.retrieve_online_documents_v2( - features=[ - "item_embeddings:embedding_float", - "item_embeddings:item_id", - "item_embeddings:string_feature", - ], + features=features, query=[1.0, 2.0], top_k=2, distance_metric="L2", @@ -948,6 +962,50 @@ def test_retrieve_online_milvus_documents(environment, fake_document_data): f"Integration test: embedding {i} has {len(embedding)} dimensions, expected {query_dim}" ) + # Oversized top_k: dataset has 3 rows, request 5 -> expect 3 back. + all_docs = fs.retrieve_online_documents_v2( + features=features, + query=[1.0, 2.0], + top_k=5, + distance_metric="L2", + ).to_dict() + assert len(all_docs["embedding_float"]) == 3 + assert sorted(all_docs["item_id"]) == [1, 2, 3] + + # Cosine-metric variant: separate FV so the Milvus collection is created + # with COSINE as its index metric. + cosine_fv = FeatureView( + name="item_embeddings_cosine", + entities=[item()], + schema=[ + Field( + name="embedding_float", + dtype=Array(Float32), + vector_index=True, + vector_search_metric="COSINE", + ), + Field(name="string_feature", dtype=String), + Field(name="float_feature", dtype=Float32), + ], + source=data_source, + ttl=timedelta(hours=2), + ) + fs.apply([cosine_fv]) + fs.write_to_online_store("item_embeddings_cosine", df) + + cosine_docs = fs.retrieve_online_documents_v2( + features=[ + "item_embeddings_cosine:embedding_float", + "item_embeddings_cosine:item_id", + "item_embeddings_cosine:string_feature", + ], + query=[1.0, 2.0], + top_k=2, + distance_metric="COSINE", + ).to_dict() + assert len(cosine_docs["embedding_float"]) == 2 + assert len(cosine_docs["item_id"]) == 2 + @pytest.mark.integration @pytest.mark.universal_online_stores(only=["milvus"])