From d8eafbada18c8daa017bf7536ccd00c4b3312387 Mon Sep 17 00:00:00 2001 From: ntkathole Date: Sat, 5 Apr 2025 20:08:32 +0530 Subject: [PATCH] fix: Fixed integration tests for qdrant and milvus Signed-off-by: ntkathole --- sdk/python/feast/feature_store.py | 4 ---- .../infra/online_stores/qdrant_online_store/qdrant.py | 8 ++++++-- .../integration/feature_repos/repo_configuration.py | 10 +++++++--- .../integration/online_store/test_universal_online.py | 2 +- 4 files changed, 14 insertions(+), 10 deletions(-) diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 413ddff4d82..1578a91574e 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -2020,10 +2020,6 @@ def _retrieve_from_online_store( """ Search and return document features from the online document store. """ - vector_field_metadata = _get_feature_view_vector_field_metadata(table) - if vector_field_metadata: - distance_metric = vector_field_metadata.vector_search_metric - documents = provider.retrieve_online_documents( config=self.config, table=table, diff --git a/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py b/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py index 2a2fafdab1c..840458b5471 100644 --- a/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py +++ b/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py @@ -115,7 +115,11 @@ def online_write_batch( encoded_value = base64.b64encode(value.SerializeToString()).decode( "utf-8" ) - vector_val = json.loads(get_list_val_str(value)) + vector_val = get_list_val_str(value) + if vector_val: + vector = {config.online_store.vector_name: json.loads(vector_val)} + else: + vector = {} points.append( models.PointStruct( id=uuid.uuid4().hex, @@ -126,7 +130,7 @@ def online_write_batch( "timestamp": timestamp, "created_ts": created_ts, }, - vector={config.online_store.vector_name: vector_val}, + vector=vector, ) ) diff --git a/sdk/python/tests/integration/feature_repos/repo_configuration.py b/sdk/python/tests/integration/feature_repos/repo_configuration.py index c24976ffcfc..0b6d8b85bfb 100644 --- a/sdk/python/tests/integration/feature_repos/repo_configuration.py +++ b/sdk/python/tests/integration/feature_repos/repo_configuration.py @@ -565,9 +565,13 @@ def construct_test_environment( cache_ttl_seconds=1, ) - if isinstance( - test_repo_config.online_store, dict - ) and test_repo_config.online_store.get("type") in ["milvus", "pgvector", "qdrant"]: + online_store = ( + test_repo_config.online_store.get("type") + if isinstance(test_repo_config.online_store, dict) + else test_repo_config.online_store + ) + + if online_store in ["milvus", "pgvector", "qdrant"]: entity_key_serialization_version = 3 environment_params = { 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 3445cd27aa3..b563f00bfd1 100644 --- a/sdk/python/tests/integration/online_store/test_universal_online.py +++ b/sdk/python/tests/integration/online_store/test_universal_online.py @@ -870,7 +870,7 @@ def test_retrieve_online_documents(environment, fake_document_data): fs.write_to_online_store("item_embeddings", df) documents = fs.retrieve_online_documents( - features=["item_embeddings:embedding_float"], + features=["item_embeddings:embedding_float", "item_embeddings:item_id"], query=[1.0, 2.0], top_k=2, distance_metric="L2",