# Copyright 2025-present the zvec project # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations import pytest import zvec from zvec import ( Collection, CollectionOption, DataType, Doc, FieldSchema, HnswIndexParam, IndexOption, IndexType, InvertIndexParam, LogLevel, LogType, OptimizeOption, StatusCode, Query, VectorSchema, ) from zvec.extension.multi_vector_reranker import ( CallbackReRanker, RrfReRanker, WeightedReRanker, ) # ==================== Common ==================== @pytest.fixture(scope="session") def collection_schema(): return zvec.CollectionSchema( name="test_collection", fields=[ FieldSchema( "id", DataType.INT64, nullable=False, index_param=InvertIndexParam(enable_range_optimization=True), ), FieldSchema( "name", DataType.STRING, nullable=False, index_param=InvertIndexParam() ), FieldSchema("weight", DataType.FLOAT, nullable=True), FieldSchema("height", DataType.INT32, nullable=True), ], vectors=[ VectorSchema( "dense", DataType.VECTOR_FP32, dimension=128, index_param=HnswIndexParam(), ), VectorSchema( "dense2", DataType.VECTOR_FP32, dimension=128, index_param=HnswIndexParam(), ), VectorSchema( "sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam() ), VectorSchema( "sparse2", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam() ), ], ) @pytest.fixture(scope="session") def collection_option(): return CollectionOption(read_only=False, enable_mmap=True) @pytest.fixture def single_doc(): id = 0 return Doc( id=f"{id}", fields={"id": id, "name": "test", "weight": 80.0, "height": id + 140}, vectors={ "dense": [id + 0.1] * 128, "dense2": [id + 0.2] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, }, ) @pytest.fixture def multiple_docs(): return [ Doc( id=f"{id}", fields={"id": id, "name": "test", "weight": 80.0, "height": 210}, vectors={ "dense": [id + 0.1] * 128, "dense2": [id + 0.2] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, }, ) for id in range(1, 101) ] @pytest.fixture(scope="function") def test_collection( tmp_path_factory, collection_schema, collection_option ) -> Collection: """ Function-scoped fixture: creates and opens a collection. Uses tmp_path_factory to ensure shared temp dir per class. """ # Create unique temp directory for this test class temp_dir = tmp_path_factory.mktemp("zvec") collection_path = temp_dir / "test_collection" coll = zvec.create_and_open( path=str(collection_path), schema=collection_schema, option=collection_option ) assert coll is not None, "Failed to create and open collection" assert coll.path == str(collection_path) assert coll.schema.name == collection_schema.name assert list(coll.schema.fields) == list(collection_schema.fields) assert list(coll.schema.vectors) == list(collection_schema.vectors) assert coll.option.read_only == collection_option.read_only assert coll.option.enable_mmap == collection_option.enable_mmap try: yield coll finally: if hasattr(coll, "destroy") and coll is not None: try: coll.destroy() except Exception as e: print(f"Warning: failed to destroy collection: {e}") @pytest.fixture def collection_with_single_doc(test_collection: Collection, single_doc) -> Collection: # Setup: insert single doc assert test_collection.stats.doc_count == 0 result = test_collection.insert(single_doc) assert bool(result) assert result.ok() assert test_collection.stats.doc_count == 1 yield test_collection # Teardown: delete single doc test_collection.delete(single_doc.id) assert test_collection.stats.doc_count == 0 @pytest.fixture def collection_with_multiple_docs( test_collection: Collection, multiple_docs ) -> Collection: # Setup: insert multiple docs assert test_collection.stats.doc_count == 0 result = test_collection.insert(multiple_docs) assert len(result) == len(multiple_docs) for item in result: assert item.ok() assert test_collection.stats.doc_count == len(multiple_docs) yield test_collection # Teardown: delete multiple docs test_collection.delete([doc.id for doc in multiple_docs]) # ==================== Tests ==================== # ---------------------------- # Config Test Case # ---------------------------- class TestConfig: def test_config(self): zvec.init(log_type=LogType.CONSOLE, log_level=LogLevel.ERROR, log_dir="./log") # ---------------------------- # Collection DDL Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionDDL: def test_collection_stats(self, test_collection: Collection): assert test_collection.stats is not None stats = test_collection.stats assert stats.doc_count == 0 assert len(stats.index_completeness) == 4 assert stats.index_completeness["dense"] == 1 assert stats.index_completeness["dense2"] == 1 assert stats.index_completeness["sparse"] == 1 assert stats.index_completeness["sparse2"] == 1 # ---------------------------- # Collection Index DDL Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionIndexDDL: def test_create_index(self, test_collection: Collection): # before create field_schema = test_collection.schema.field("weight") assert field_schema is not None assert field_schema.data_type == DataType.FLOAT assert field_schema.name == "weight" index_param = field_schema.index_param assert index_param is None # create test_collection.create_index( field_name="weight", index_param=InvertIndexParam(), option=IndexOption() ) assert test_collection.schema is not None field_schema = test_collection.schema.field("weight") assert field_schema is not None assert field_schema.data_type == DataType.FLOAT assert field_schema.name == "weight" index_param = field_schema.index_param assert index_param.type == IndexType.INVERT assert index_param.enable_range_optimization is False assert index_param.enable_extended_wildcard is False def test_drop_index(self, test_collection: Collection): # before drop field_schema = test_collection.schema.field("name") assert field_schema is not None assert field_schema.data_type == DataType.STRING assert field_schema.name == "name" index_param = field_schema.index_param assert index_param.type == IndexType.INVERT assert index_param.enable_range_optimization is False assert index_param.enable_extended_wildcard is False # drop test_collection.drop_index("name") field_schema = test_collection.schema.field("name") assert field_schema is not None assert field_schema.data_type == DataType.STRING assert field_schema.name == "name" # without index index_param = field_schema.index_param assert index_param is None def test_create_index_field_is_not_exist(self, test_collection: Collection): with pytest.raises(Exception) as e: test_collection.create_index( field_name="not_exist", index_param=InvertIndexParam(), ) index_param = field_schema.index_param assert index_param.type == IndexType.INVERT assert index_param.enable_range_optimization is False assert index_param.enable_extended_wildcard is False def test_drop_index(self, test_collection: Collection): # before drop field_schema = test_collection.schema.field("name") assert field_schema is not None assert field_schema.data_type == DataType.STRING assert field_schema.name == "name" index_param = field_schema.index_param assert index_param.type == IndexType.INVERT assert index_param.enable_range_optimization is False assert index_param.enable_extended_wildcard is False # drop test_collection.drop_index("name") field_schema = test_collection.schema.field("name") assert field_schema is not None assert field_schema.data_type == DataType.STRING assert field_schema.name == "name" # without index index_param = field_schema.index_param assert index_param is None def test_create_index_field_is_not_exist(self, test_collection: Collection): with pytest.raises(Exception) as e: test_collection.create_index( field_name="not_exist", index_param=InvertIndexParam(), ) # ---------------------------- # Collection Column DDL Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionColumnDDL: def test_create_column(self, test_collection: Collection): # before create column field_schema = test_collection.schema.field("age") assert field_schema is None # create test_collection.add_column(FieldSchema("age", DataType.INT32, nullable=True)) field_schema = test_collection.schema.field("age") assert field_schema is not None assert field_schema.data_type == DataType.INT32 assert field_schema.name == "age" assert field_schema.index_param is None def test_create_column_is_nullable(self, test_collection: Collection): with pytest.raises(ValueError): test_collection.add_column( FieldSchema("age", DataType.INT32, nullable=False) ) def test_drop_column(self, test_collection: Collection): # before drop column field_schema = test_collection.schema.field("id") assert field_schema is not None assert field_schema.data_type == DataType.INT64 assert field_schema.name == "id" index_param = field_schema.index_param assert index_param is not None assert index_param.type == IndexType.INVERT # drop test_collection.drop_column("id") field_schema = test_collection.schema.field("id") assert field_schema is None def test_alert_column_to_rename(self, test_collection: Collection): # before alert column field_schema = test_collection.schema.field("id") assert field_schema is not None assert field_schema.data_type == DataType.INT64 assert field_schema.name == "id" index_param = field_schema.index_param assert index_param is not None assert index_param.type == IndexType.INVERT assert index_param.enable_range_optimization is True assert index_param.enable_extended_wildcard is False # alert rename test_collection.alter_column("id", "doc_id") # validate old column field_schema = test_collection.schema.field("id") assert field_schema is None # validate rename column field_schema = test_collection.schema.field("doc_id") assert field_schema is not None assert field_schema.data_type == DataType.INT64 assert field_schema.name == "doc_id" assert field_schema.nullable is False index_param = field_schema.index_param assert index_param is not None assert index_param.type == IndexType.INVERT assert index_param.enable_range_optimization is True assert index_param.enable_extended_wildcard is False def test_alert_column_to_modify_schema(self, test_collection: Collection): # before alert column field_schema = test_collection.schema.field("id") assert field_schema is not None assert field_schema.data_type == DataType.INT64 assert field_schema.name == "id" index_param = field_schema.index_param assert index_param.type == IndexType.INVERT test_collection.alter_column( old_name="id", field_schema=FieldSchema("doc_id", DataType.UINT64, nullable=True), ) field_schema = test_collection.schema.field("doc_id") assert field_schema is not None assert field_schema.data_type == DataType.UINT64 assert field_schema.name == "doc_id" def test_column_with_other_dtype(self, test_collection: Collection): # only allow number type test_collection.add_column(FieldSchema("age", DataType.INT32, nullable=True)) with pytest.raises(ValueError): test_collection.add_column(FieldSchema("full_name", DataType.STRING)) with pytest.raises(ValueError): test_collection.drop_column("name") with pytest.raises(ValueError): test_collection.alter_column(old_name="name", new_name="full_name") with pytest.raises(ValueError): test_collection.alter_column( old_name="name", field_schema=FieldSchema("full_name", DataType.STRING) ) # ---------------------------- # Collection Optimize Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionOptimize: def test_collection_optimize(self, test_collection: Collection): test_collection.optimize(option=OptimizeOption()) # ---------------------------- # Collection Fetch Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionFetch: def test_collection_fetch( self, collection_with_single_doc: Collection, single_doc: Doc ): result = collection_with_single_doc.fetch(ids=[single_doc.id]) assert bool(result) assert single_doc.id in result.keys() doc = result[single_doc.id] assert doc is not None assert doc.id == single_doc.id assert set(doc.field_names()) == set(single_doc.field_names()) for field_name in doc.field_names(): if field_name in ["dense", "sparse"]: continue assert doc.field(field_name) == single_doc.field(field_name) def test_collection_fetch_contains_nodata_ids( self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc] ): ids = [doc.id for doc in multiple_docs] no_data_key = "x" ids_with_no_data = [no_data_key] + ids result = collection_with_multiple_docs.fetch(ids=ids_with_no_data) assert bool(result) assert len(result) == len(ids) assert no_data_key not in result # ---------------------------- # Collection Insert Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionInsert: def test_collection_insert(self, test_collection, single_doc): result = test_collection.insert(single_doc) assert bool(result) assert result.ok() stats = test_collection.stats assert stats is not None assert stats.doc_count == 1 def test_collection_insert_with_nullable_false_field(self, test_collection): # id, name's nullable == False # weight, height's nullable == True doc = Doc( id="0", fields={ "id": 1, "name": "test", }, vectors={ "dense": [1 + 0.1] * 128, "dense2": [1 + 0.2] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, }, ) result = test_collection.insert(doc) assert bool(result) assert result.ok() stats = test_collection.stats assert stats is not None assert stats.doc_count == 1 def test_collection_insert_without_nullable_false_field(self, test_collection): # id, name's nullable == False # weight, height's nullable == True # without id, name doc = Doc( id="0", vectors={ "dense": [1 + 0.1] * 128, "dense2": [1 + 0.2] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, }, ) with pytest.raises(ValueError) as e: # ValueError: Invalid doc: field[id] is required but not provided test_collection.insert(doc) assert "field[id] is required but not provided" in str(e.value) # without name doc = Doc( id="0", fields={ "id": 1, }, vectors={ "dense": [1 + 0.1] * 128, "dense2": [1 + 0.2] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, }, ) with pytest.raises(ValueError) as e: test_collection.insert(doc) assert "field[name] is required but not provided" in str(e.value) def test_collection_insert_with_nullable_true_field(self, test_collection): # id, name's nullable == False # weight, height's nullable == True doc = Doc( id="0", fields={ "id": 1, "name": "test", }, vectors={ "dense": [1 + 0.1] * 128, "dense2": [1 + 0.2] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}, "sparse2": {4: 1.5, 5: 2.5, 6: 3.5}, }, ) result = test_collection.insert(doc) assert bool(result) assert result.ok() stats = test_collection.stats assert stats is not None assert stats.doc_count == 1 result = test_collection.fetch(ids=[doc.id]) assert doc.id in result ret = result[doc.id] assert ret.field("id") == 1 assert ret.field("name") == "test" assert ret.field("weight") is None assert ret.field("height") is None def test_collection_insert_batch(self, test_collection, multiple_docs): result = test_collection.insert(multiple_docs) assert len(result) == len(multiple_docs) for item in result: assert item.ok() stats = test_collection.stats assert stats is not None assert stats.doc_count == len(multiple_docs) def test_collection_insert_duplicate( self, test_collection, single_doc, multiple_docs ): test_collection.insert(single_doc) result = test_collection.insert(single_doc) assert bool(result) assert result.code() == StatusCode.ALREADY_EXISTS stats = test_collection.stats assert stats is not None assert stats.doc_count == 1 # ---------------------------- # Collection Update Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionUpdate: def test_empty_collection_update( self, test_collection: Collection, single_doc: Doc ): result = test_collection.update(single_doc) assert bool(result) assert result.code() == StatusCode.NOT_FOUND stats = test_collection.stats assert stats is not None assert stats.doc_count == 0 def test_collection_update_with_nullable_false_field( self, collection_with_single_doc: Collection, single_doc: Doc ): # id, name's nullable == False # weight, height's nullable == True # update doc field id doc = Doc( id=single_doc.id, fields={"id": single_doc.field("id") + 1}, ) result = collection_with_single_doc.update(doc) assert bool(result) assert result.ok() stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 1 # fetch result = collection_with_single_doc.fetch(ids=[doc.id]) assert doc.id in result ret = result[doc.id] assert ret.field("id") == doc.field("id") assert ret.field("name") == single_doc.field("name") assert ret.field("weight") == single_doc.field("weight") assert ret.field("height") == single_doc.field("height") def test_collection_update_with_nullable_false_field_is_none( self, collection_with_single_doc: Collection, single_doc: Doc ): # id, name's nullable == False # weight, height's nullable == True # update doc field id doc = Doc( id=single_doc.id, fields={"id": None}, ) with pytest.raises(ValueError) as e: # ValueError: Invalid doc: field[id] is required but its value is null collection_with_single_doc.update(doc) doc = Doc( id=single_doc.id, fields={"id": single_doc.field("id") + 1, "weight": None}, ) result = collection_with_single_doc.update(doc) assert bool(result) assert result.ok() stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 1 ret = collection_with_single_doc.fetch(ids=[doc.id]) assert doc.id in ret ret = ret[doc.id] assert ret.field("id") == doc.field("id") assert ret.field("name") == single_doc.field("name") assert ret.field("weight") is None assert ret.field("height") == single_doc.field("height") def test_collection_update_without_nullable_false_field( self, collection_with_single_doc: Collection, single_doc: Doc ): # id, name's nullable == False # weight, height's nullable == True # update doc field weight doc = Doc( id=single_doc.id, fields={"weight": single_doc.field("weight") + 1}, ) result = collection_with_single_doc.update(doc) assert bool(result) assert result.ok() stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 1 # fetch ret = collection_with_single_doc.fetch(ids=[doc.id]) assert doc.id in ret ret = ret[doc.id] assert ret.field("id") == single_doc.field("id") assert ret.field("name") == single_doc.field("name") assert ret.field("weight") == doc.field("weight") assert ret.field("height") == single_doc.field("height") def test_collection_update_without_nullable_false_field_set_null( self, collection_with_single_doc: Collection, single_doc: Doc ): # id, name's nullable == False # weight, height's nullable == True # update doc field weight is None doc = Doc( id=single_doc.id, fields={"weight": None}, ) result = collection_with_single_doc.update(doc) assert bool(result) assert result.ok() stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 1 # fetch ret = collection_with_single_doc.fetch(ids=[doc.id]) assert doc.id in ret ret = ret[doc.id] assert ret.field("id") == single_doc.field("id") assert ret.field("name") == single_doc.field("name") assert ret.field("weight") is None assert ret.field("height") == single_doc.field("height") def test_empty_collection_update_batch( self, test_collection: Collection, multiple_docs ): result = test_collection.update(multiple_docs) assert len(result) == len(multiple_docs) for item in result: assert item.code() == StatusCode.NOT_FOUND stats = test_collection.stats assert stats is not None assert stats.doc_count == 0 def test_collection_update( self, collection_with_single_doc: Collection, single_doc ): result = collection_with_single_doc.update(single_doc) assert bool(result) == 1 assert result.ok() stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 1 def test_collection_update_batch( self, collection_with_multiple_docs: Collection, multiple_docs ): result = collection_with_multiple_docs.update(multiple_docs) assert len(result) == len(multiple_docs) for item in result: assert item.ok() stats = collection_with_multiple_docs.stats assert stats is not None assert stats.doc_count == len(multiple_docs) # ---------------------------- # Collection Upsert Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionUpsert: def test_empty_collection_upsert(self, test_collection: Collection, single_doc): result = test_collection.upsert(single_doc) assert bool(result) assert result.ok() stats = test_collection.stats assert stats is not None assert stats.doc_count == 1 def test_empty_collection_upsert_batch( self, test_collection: Collection, multiple_docs ): result = test_collection.upsert(multiple_docs) assert len(result) == len(multiple_docs) for item in result: assert item.ok() stats = test_collection.stats assert stats is not None assert stats.doc_count == len(multiple_docs) def test_collection_upsert( self, collection_with_single_doc: Collection, single_doc, multiple_docs ): # doc is existing # upsert => update result = collection_with_single_doc.upsert(single_doc) assert bool(result) assert result.ok() stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 1 def test_collection_upsert_batch( self, collection_with_multiple_docs: Collection, multiple_docs ): # doc is existing # upsert => update result = collection_with_multiple_docs.upsert(multiple_docs) assert len(result) == len(multiple_docs) for item in result: assert item.ok() stats = collection_with_multiple_docs.stats assert stats is not None assert stats.doc_count == len(multiple_docs) # ---------------------------- # Collection Upsert Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionDelete: def test_empty_collection_delete(self, test_collection: Collection, single_doc): result = test_collection.delete(single_doc.id) assert bool(result) assert result.code() == StatusCode.NOT_FOUND def test_empty_collection_delete_batch( self, test_collection: Collection, multiple_docs ): result = test_collection.delete([doc.id for doc in multiple_docs]) assert len(result) == len(multiple_docs) for item in result: assert item.code() == StatusCode.NOT_FOUND def test_collection_delete( self, collection_with_single_doc: Collection, single_doc ): result = collection_with_single_doc.delete(single_doc.id) assert bool(result) assert result.ok() stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 0 result = collection_with_single_doc.insert(single_doc) assert bool(result) assert result.ok() stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 1 def test_collection_delete_batch( self, collection_with_multiple_docs: Collection, multiple_docs ): result = collection_with_multiple_docs.delete([doc.id for doc in multiple_docs]) assert len(result) == len(multiple_docs) for item in result: assert item.ok() stats = collection_with_multiple_docs.stats assert stats is not None assert stats.doc_count == 0 def test_collection_delete_by_filter( self, collection_with_single_doc: Collection, single_doc ): collection_with_single_doc.delete_by_filter( filter=f"height={single_doc.field('height')}" ) stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 0 def test_collection_delete_by_filter_invert_field( self, collection_with_single_doc: Collection, single_doc ): collection_with_single_doc.delete_by_filter( filter=f"id={single_doc.field('id')}" ) stats = collection_with_single_doc.stats assert stats is not None assert stats.doc_count == 0 # ---------------------------- # Collection Upsert Test Case # ---------------------------- @pytest.mark.usefixtures("test_collection") class TestCollectionQuery: def test_empty_collection_query(self, test_collection: Collection): result = test_collection.query() assert len(result) == 0 def test_collection_query(self, collection_with_single_doc: Collection, single_doc): result = collection_with_single_doc.query() assert len(result) == 1 doc = result[0] assert doc.id == single_doc.id assert "dense" not in doc.field_names() assert "sparse" not in doc.field_names() field_without_vector = single_doc.field_names() assert set(doc.field_names()) == set(field_without_vector) for name in field_without_vector: assert doc.field(name) == single_doc.field(name) def test_collection_query_with_include_vector( self, collection_with_single_doc: Collection, single_doc ): result = collection_with_single_doc.query(include_vector=True) assert len(result) == 1 doc = result[0] assert doc.vector("dense") is not None assert doc.vector("sparse") is not None def test_collection_query_with_output_fields( self, collection_with_single_doc: Collection, single_doc ): result = collection_with_single_doc.query(output_fields=["id", "name"]) assert len(result) == 1 doc = result[0] assert doc.id == single_doc.id assert len(doc.field_names()) == 2 assert set(doc.field_names()) == {"id", "name"} def test_collection_query_with_topk( self, collection_with_multiple_docs: Collection ): result = collection_with_multiple_docs.query() assert len(result) == 10 result = collection_with_multiple_docs.query(topk=5) assert len(result) == 5 def test_collection_query_with_range_filter_int_field( self, collection_with_multiple_docs: Collection, multiple_docs ): index = 10 idx = multiple_docs[index].id result = collection_with_multiple_docs.query(filter=f"id>{idx}", topk=100) assert len(result) == len(multiple_docs) - index - 1 result = collection_with_multiple_docs.query(filter=f"id>={idx}", topk=100) assert len(result) == len(multiple_docs) - index result = collection_with_multiple_docs.query(filter=f"id<{idx}", topk=100) assert len(result) == index result = collection_with_multiple_docs.query(filter=f"id<={idx}", topk=100) assert len(result) == index + 1 result = collection_with_multiple_docs.query(filter=f"id={idx}", topk=100) assert len(result) == 1 result = collection_with_multiple_docs.query(filter=f"id!={idx}", topk=100) assert len(result) == len(multiple_docs) - 1 left, right = 10, 90 l_id, r_id = multiple_docs[left].id, multiple_docs[right].id result = collection_with_multiple_docs.query( filter=f"id>{l_id} and id<{r_id}", topk=100 ) assert len(result) == right - left - 1 result = collection_with_multiple_docs.query( filter=f"id>={l_id} and id<{r_id}", topk=100 ) assert len(result) == right - left result = collection_with_multiple_docs.query( filter=f"id>={l_id} and id<={r_id}", topk=100 ) assert len(result) == right - left + 1 result = collection_with_multiple_docs.query( filter=f"id<{l_id} or id>{r_id}", topk=100 ) assert len(result) == len(multiple_docs) - (right - left) - 1 result = collection_with_multiple_docs.query( filter=f"id<={l_id} or id>{r_id}", topk=100 ) assert len(result) == len(multiple_docs) - (right - left) result = collection_with_multiple_docs.query( filter=f"id<={l_id} or id>={r_id}", topk=100 ) assert len(result) == len(multiple_docs) - (right - left) + 1 result = collection_with_multiple_docs.query(filter="id in (1)", topk=100) assert len(result) == 1 def test_collection_query_with_filter_not_in( self, collection_with_multiple_docs: Collection, multiple_docs ): result = collection_with_multiple_docs.query(filter="id not in (1)", topk=100) assert len(result) == len(multiple_docs) - 1 def test_collection_with_error_query_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): query = Query( field_name="dense", vector=multiple_docs[0].vector("dense"), param=[1, 2, 3] ) with pytest.raises(TypeError): result = collection_with_multiple_docs.query( query, filter="id in (1)", topk=100 ) def test_collection_query_by_id( self, collection_with_multiple_docs: Collection, multiple_docs ): result = collection_with_multiple_docs.query( Query(field_name="dense", id=multiple_docs[0].id) ) assert len(result) == 10 def test_collection_query_by_dense_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): result = collection_with_multiple_docs.query( Query(field_name="dense", vector=multiple_docs[0].vector("dense")), topk=10, ) assert len(result) > 0 assert len(result) <= 10 def test_collection_query_by_sparse_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): result = collection_with_multiple_docs.query( Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), topk=10, ) assert len(result) > 0 assert len(result) <= 10 def test_collection_query_by_dense_vector_with_filter( self, collection_with_multiple_docs: Collection, multiple_docs ): result = collection_with_multiple_docs.query( Query(field_name="dense", vector=multiple_docs[0].vector("dense")), topk=10, filter="id > 50", ) assert len(result) > 0 assert len(result) <= 10 for doc in result: assert int(doc.id) > 50 def test_collection_query_by_sparse_vector_with_filter( self, collection_with_multiple_docs: Collection, multiple_docs ): result = collection_with_multiple_docs.query( Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), topk=10, filter="id > 50", ) assert len(result) > 0 assert len(result) <= 10 for doc in result: assert int(doc.id) > 50 def test_collection_query_with_rrf_reranker_by_multi_dense_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): """Test multi-vector query with RRF reranker on multiple dense vectors.""" reranker = RrfReRanker(rank_constant=60) result = collection_with_multiple_docs.query( [ Query(field_name="dense", vector=multiple_docs[0].vector("dense")), Query(field_name="dense2", vector=multiple_docs[0].vector("dense2")), ], topk=10, reranker=reranker, ) assert len(result) > 0 assert len(result) <= 10 # Results should have RRF-fused scores for doc in result: assert hasattr(doc, "score") def test_collection_query_with_rrf_reranker_by_multi_sparse_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): """Test multi-vector query with RRF reranker on multiple sparse vectors.""" reranker = RrfReRanker(rank_constant=60) result = collection_with_multiple_docs.query( [ Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), Query( field_name="sparse2", vector=multiple_docs[0].vector("sparse2"), ), ], topk=10, reranker=reranker, ) assert len(result) > 0 assert len(result) <= 10 def test_collection_query_with_rrf_reranker_by_hybrid_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): """Test multi-vector query with RRF reranker combining dense + sparse.""" reranker = RrfReRanker(rank_constant=60) result = collection_with_multiple_docs.query( [ Query(field_name="dense", vector=multiple_docs[0].vector("dense")), Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), ], topk=10, reranker=reranker, ) assert len(result) > 0 assert len(result) <= 10 def test_collection_query_with_weighted_reranker_by_multi_dense_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): """Test multi-vector query with Weighted reranker on multiple dense vectors.""" weights = [0.6, 0.4] reranker = WeightedReRanker(weights=weights) result = collection_with_multiple_docs.query( [ Query(field_name="dense", vector=multiple_docs[0].vector("dense")), Query(field_name="dense2", vector=multiple_docs[0].vector("dense2")), ], topk=10, reranker=reranker, ) assert len(result) > 0 assert len(result) <= 10 def test_collection_query_with_weighted_reranker_by_multi_sparse_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): """Test multi-vector query with Weighted reranker on multiple sparse vectors.""" weights = [0.6, 0.4] reranker = WeightedReRanker(weights=weights) result = collection_with_multiple_docs.query( [ Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), Query( field_name="sparse2", vector=multiple_docs[0].vector("sparse2"), ), ], topk=10, reranker=reranker, ) assert len(result) > 0 assert len(result) <= 10 def test_collection_query_with_weighted_reranker_by_hybrid_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): """Test multi-vector query with Weighted reranker combining dense + sparse.""" weights = [0.7, 0.3] reranker = WeightedReRanker(weights=weights) result = collection_with_multiple_docs.query( [ Query(field_name="dense", vector=multiple_docs[0].vector("dense")), Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), ], topk=10, reranker=reranker, ) assert len(result) > 0 assert len(result) <= 10 def test_collection_query_with_callback_reranker_by_multi_dense_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): """Test multi-vector query with CallbackReRanker (Python callback via C++).""" callback_invoked = [] def my_rerank_callback(query_results, fields, topn): callback_invoked.append(True) all_docs = [] for docs in query_results: all_docs.extend(docs) seen = set() unique_docs = [] for doc in all_docs: if doc.pk() not in seen: seen.add(doc.pk()) unique_docs.append(doc) unique_docs.sort(key=lambda d: d.score(), reverse=True) return unique_docs[:topn] reranker = CallbackReRanker(callback=my_rerank_callback) result = collection_with_multiple_docs.query( [ Query(field_name="dense", vector=multiple_docs[0].vector("dense")), Query(field_name="dense2", vector=multiple_docs[0].vector("dense2")), ], topk=10, reranker=reranker, ) assert len(callback_invoked) == 1 assert len(result) > 0 assert len(result) <= 10 def test_collection_query_with_callback_reranker_by_hybrid_vector( self, collection_with_multiple_docs: Collection, multiple_docs ): """Test multi-vector query with CallbackReRanker combining dense + sparse.""" def my_rerank_callback(query_results, fields, topn): all_docs = [] for docs in query_results: all_docs.extend(docs) seen = set() unique_docs = [] for doc in all_docs: if doc.pk() not in seen: seen.add(doc.pk()) unique_docs.append(doc) unique_docs.sort(key=lambda d: d.score(), reverse=True) return unique_docs[:topn] reranker = CallbackReRanker(callback=my_rerank_callback) result = collection_with_multiple_docs.query( [ Query(field_name="dense", vector=multiple_docs[0].vector("dense")), Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")), ], topk=5, reranker=reranker, ) assert len(result) > 0 assert len(result) <= 5