# 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. import pytest from zvec.typing import DataType, StatusCode, MetricType, QuantizeType from zvec.model import Collection, Doc, Query from zvec.model.param import ( CollectionOption, InvertIndexParam, HnswIndexParam, FlatIndexParam, IVFIndexParam, DiskAnnIndexParam, HnswQueryParam, IVFQueryParam, DiskAnnQueryParam, ) from zvec.model.schema import FieldSchema, VectorSchema from zvec.extension import RrfReRanker, WeightedReRanker, QwenReRanker from distance_helper import * from zvec import StatusCode from distance_helper import * from fixture_helper import * from doc_helper import * from params_helper import * import time # ==================== helper ==================== def batchdoc_and_check(collection: Collection, multiple_docs, operator="insert"): if operator == "insert": result = collection.insert(multiple_docs) elif operator == "upsert": result = collection.upsert(multiple_docs) elif operator == "update": result = collection.update(multiple_docs) else: logging.error("operator value is error!") assert len(result) == len(multiple_docs) for item in result: assert item.ok(), ( f"result={result},Insert operation failed with code {item.code()}" ) stats = collection.stats assert stats is not None, "Collection stats should not be None" """assert stats.doc_count == len(multiple_docs), ( f"Document count should be {len(multiple_docs)} after insert, but got {stats.doc_count}" )""" doc_ids = [doc.id for doc in multiple_docs] fetched_docs = collection.fetch(doc_ids) assert len(fetched_docs) == len(multiple_docs), ( f"fetched_docs={fetched_docs},Expected {len(multiple_docs)} fetched documents, but got {len(fetched_docs)}" ) for original_doc in multiple_docs: assert original_doc.id in fetched_docs, ( f"Expected document ID {original_doc.id} in fetched documents" ) fetched_doc = fetched_docs[original_doc.id] assert is_doc_equal(fetched_doc, original_doc, collection.schema) assert hasattr(fetched_doc, "score"), "Document should have a score attribute" assert fetched_doc.score == 0.0, ( "Fetch operation should return default score of 0.0" ) def compute_exact_similarity_scores( vectors_a, vectors_b, metric_type=MetricType.IP, DataType=DataType.VECTOR_FP32, QuantizeType=QuantizeType.UNDEFINED, ): similarities = [] for i, vec_a in enumerate(vectors_a): for j, vec_b in enumerate(vectors_b): similarity = distance_recall(vec_a, vec_b, metric_type, DataType) similarities.append((j, similarity)) # For L2,COSINE metric, smaller distances mean higher similarity, so sort in ascending order if ( metric_type in [MetricType.L2] and DataType in [DataType.VECTOR_FP32, DataType.VECTOR_FP16, DataType.VECTOR_INT8] ) or ( metric_type in [MetricType.COSINE] and DataType in [DataType.VECTOR_FP32, DataType.VECTOR_FP16] ): similarities.sort(key=lambda x: x[1], reverse=False) # Ascending order for L2 else: similarities.sort( key=lambda x: x[1], reverse=True ) # Descending order for others # Special handling for COSINE in FP16 to address precision issues if metric_type == MetricType.COSINE and DataType == DataType.VECTOR_FP16: # Clamp values to valid cosine distance range [0, 2] and handle floating point errors similarities = [(idx, max(0.0, min(2.0, score))) for idx, score in similarities] return similarities def get_ground_truth_for_vector_query( collection, query_vector, field_name, all_docs, query_idx, metric_type, k, use_exact_computation=False, ): if use_exact_computation: all_vectors = [doc.vectors[field_name] for doc in all_docs] for d, f in DEFAULT_VECTOR_FIELD_NAME.items(): if field_name == f: DataType = d break similarities = compute_exact_similarity_scores( [query_vector], all_vectors, metric_type, DataType=DataType, QuantizeType=QuantizeType, ) if metric_type == MetricType.COSINE and DataType == DataType.VECTOR_FP16: # Filter out tiny non-zero values that may be caused by precision errors similarities = [ (idx, max(0.0, min(2.0, score))) for idx, score in similarities ] ground_truth_ids_scores = similarities[:k] print("Get the most similar k document IDs k:,ground_truth_ids_scores") print(k, ground_truth_ids_scores) return ground_truth_ids_scores else: full_result = collection.query( Query(field_name=field_name, vector=query_vector), topk=min(len(all_docs), 1024), include_vector=True, ) ground_truth_ids_scores = [ (result.id, result.score) for result in full_result[:k] ] if not ground_truth_ids_scores: ground_truth_ids_scores = [(all_docs[query_idx].id, 0)] return ground_truth_ids_scores def get_ground_truth_map(collection, test_docs, query_vectors_map, metric_type, k): ground_truth_map = {} for field_name, query_vectors in query_vectors_map.items(): ground_truth_map[field_name] = {} # Support per-field metric type: metric_type can be a dict mapping # field_name -> MetricType, or a single MetricType applied to all fields. if isinstance(metric_type, dict): field_metric = metric_type.get(field_name, MetricType.IP) else: field_metric = metric_type for i, query_vector in enumerate(query_vectors): # Get the ground truth for this query relevant_doc_ids_scores = get_ground_truth_for_vector_query( collection, query_vector, field_name, test_docs, i, field_metric, k, True, ) ground_truth_map[field_name][i] = relevant_doc_ids_scores print("ground_truth_map:\n") print(ground_truth_map) return ground_truth_map def calculate_recall_at_k( collection: Collection, test_docs, query_vectors_map, schema, k=1, expected_doc_ids_scores_map=None, tolerance=0.01, ): recall_stats = {} for field_name, query_vectors in query_vectors_map.items(): recall_stats[field_name] = { "relevant_retrieved_count": 0, "total_relevant_count": 0, "retrieved_count": 0, "recall_at_k": 0.0, } for i, query_vector in enumerate(query_vectors): print("Starting %dth query" % i) query_result_list = collection.query( Query(field_name=field_name, vector=query_vector), topk=1024, include_vector=True, ) retrieved_count = len(query_result_list) query_result_ids_scores = [] for word in query_result_list: query_result_ids_scores.append((word.id, word.score)) recall_stats[field_name]["retrieved_count"] += retrieved_count print("expected_doc_ids_scores_map:\n") print(expected_doc_ids_scores_map) if i in (expected_doc_ids_scores_map[field_name]): expected_relevant_ids_scores = expected_doc_ids_scores_map[field_name][ i ] print( "field_name,i,expected_relevant_ids_scores, query_result_ids_scores:\n" ) print( field_name, i, "\n", expected_relevant_ids_scores, "\n", len(query_result_ids_scores), query_result_ids_scores, ) # Update total relevant documents count recall_stats[field_name]["total_relevant_count"] += len( expected_relevant_ids_scores ) relevant_found_count = 0 for ids_scores_except in expected_relevant_ids_scores: for ids_scores_result in query_result_ids_scores[:k]: if int(ids_scores_result[0]) == int(ids_scores_except[0]): relevant_found_count += 1 break elif ( int(ids_scores_result[0]) != int(ids_scores_except[0]) and abs(ids_scores_result[1] - ids_scores_except[1]) <= tolerance ): print("IDs are not equal, but the error is small, tolerance") print( ids_scores_result[0], ids_scores_except[0], ids_scores_result[1], ids_scores_except[1], tolerance, ) relevant_found_count += 1 break else: continue recall_stats[field_name]["relevant_retrieved_count"] += relevant_found_count # Calculate Recall@K if recall_stats[field_name]["total_relevant_count"] > 0: recall_stats[field_name]["recall_at_k"] = ( recall_stats[field_name]["relevant_retrieved_count"] / recall_stats[field_name]["total_relevant_count"] ) return recall_stats class TestRecall: @pytest.mark.parametrize( "full_schema_new", [ (True, True, HnswIndexParam()), (False, True, IVFIndexParam()), (False, True, DiskAnnIndexParam()), (False, True, FlatIndexParam()), # ——ok ( True, True, HnswIndexParam( metric_type=MetricType.IP, m=16, ef_construction=100, ), ), ( True, True, HnswIndexParam( metric_type=MetricType.COSINE, m=24, ef_construction=150, ), ), ( True, True, HnswIndexParam( metric_type=MetricType.L2, m=32, ef_construction=200, ), ), ( False, True, FlatIndexParam( metric_type=MetricType.IP, ), ), ( True, True, FlatIndexParam( metric_type=MetricType.COSINE, ), ), ( True, True, FlatIndexParam( metric_type=MetricType.L2, ), ), ( True, True, IVFIndexParam( metric_type=MetricType.IP, n_list=100, n_iters=10, use_soar=False, ), ), ( True, True, IVFIndexParam( metric_type=MetricType.L2, n_list=200, n_iters=20, use_soar=True, ), ), ( True, True, IVFIndexParam( metric_type=MetricType.COSINE, n_list=150, n_iters=15, use_soar=False, ), ), ( True, True, DiskAnnIndexParam( metric_type=MetricType.IP, max_degree=32, ), ), ( True, True, DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32), ), ], indirect=True, ) @pytest.mark.parametrize("doc_num", [500]) @pytest.mark.parametrize("query_num", [10]) @pytest.mark.parametrize("top_k", [1]) def test_recall_with_single_vector_valid_500( self, full_collection_new: Collection, doc_num, query_num, top_k, full_schema_new, request, ): full_schema_params = request.getfixturevalue("full_schema_new") # Build per-field metric type map so ground truth uses each field's # actual index metric (fields may fall back to HnswIndexParam/IP). field_metric_map = {} for vector_para in full_schema_params.vectors: if vector_para.index_param is not None: field_metric_map[vector_para.name] = vector_para.index_param.metric_type else: field_metric_map[vector_para.name] = MetricType.IP metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP) multiple_docs = [ generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num) ] print("len(multiple_docs):\n") print(len(multiple_docs)) # print(multiple_docs) for i in range(10): if i != 0: pass # print(multiple_docs[i * 1000:1000 * (i + 1)]) batchdoc_and_check( full_collection_new, multiple_docs[i * 1000 : 1000 * (i + 1)], operator="insert", ) stats = full_collection_new.stats assert stats.doc_count == len(multiple_docs) doc_ids = ["0", "1"] fetched_docs = full_collection_new.fetch(doc_ids) print("fetched_docs,multiple_docs") print( fetched_docs[doc_ids[0]].vectors["sparse_vector_fp32_field"], fetched_docs[doc_ids[0]].vectors["sparse_vector_fp16_field"], fetched_docs[doc_ids[1]].vectors["sparse_vector_fp32_field"], fetched_docs[doc_ids[1]].vectors["sparse_vector_fp16_field"], "\n", multiple_docs[0].vectors["sparse_vector_fp32_field"], multiple_docs[0].vectors["sparse_vector_fp32_field"], multiple_docs[1].vectors["sparse_vector_fp32_field"], multiple_docs[1].vectors["sparse_vector_fp16_field"], ) full_collection_new.optimize(option=OptimizeOption()) time.sleep(2) query_vectors_map = {} for field_name in DEFAULT_VECTOR_FIELD_NAME.values(): query_vectors_map[field_name] = [ multiple_docs[i].vectors[field_name] for i in range(query_num) ] # Get ground truth mapping (pass per-field metric map) ground_truth_map = get_ground_truth_map( full_collection_new, multiple_docs, query_vectors_map, field_metric_map, top_k, ) # Validate ground truth mapping structure for field_name in DEFAULT_VECTOR_FIELD_NAME.values(): assert field_name in ground_truth_map field_gt = ground_truth_map[field_name] assert len(field_gt) == query_num for query_idx in range(query_num): assert query_idx in field_gt relevant_ids = field_gt[query_idx] assert isinstance(relevant_ids, list) assert len(relevant_ids) <= top_k # Print ground truth statistics print(f"Ground Truth for Top-{top_k} Retrieval:") for field_name, field_gt in ground_truth_map.items(): print(f" {field_name}:") for query_idx, relevant_ids in field_gt.items(): print( f" Query {query_idx}: {len(relevant_ids)} relevant docs - {relevant_ids[:5]}{'...' if len(relevant_ids) > 5 else ''}" ) # Calculate Recall@K using ground truth recall_at_k_stats = calculate_recall_at_k( full_collection_new, multiple_docs, query_vectors_map, full_schema_new, k=top_k, expected_doc_ids_scores_map=ground_truth_map, tolerance=0.01, ) print("ground_truth_map:\n") print(ground_truth_map) print("(recall_at_k_stats:\n") print(recall_at_k_stats) print("field_metric_map:") print(field_metric_map) # Print Recall@K statistics print(f"Recall@{top_k} using Ground Truth:") for field_name, stats in recall_at_k_stats.items(): print(f" {field_name}:") print( f" Relevant Retrieved: {stats['relevant_retrieved_count']}/{stats['total_relevant_count']}" ) print(f" Recall@{top_k}: {stats['recall_at_k']:.4f}") for k, v in recall_at_k_stats.items(): assert v["recall_at_k"] == 1.0 @pytest.mark.parametrize( "full_schema_new", [ (True, True, HnswIndexParam()), (False, True, IVFIndexParam()), (False, True, FlatIndexParam()), # ——ok ( True, True, HnswIndexParam( metric_type=MetricType.IP, m=16, ef_construction=100, ), ), ( True, True, HnswIndexParam( metric_type=MetricType.COSINE, m=24, ef_construction=150, ), ), # (True, True, HnswIndexParam(metric_type=MetricType.L2, m=32, ef_construction=200, )), ( False, True, FlatIndexParam( metric_type=MetricType.IP, ), ), ( True, True, FlatIndexParam( metric_type=MetricType.COSINE, ), ), # (True, True, FlatIndexParam(metric_type=MetricType.L2, )), ( True, True, IVFIndexParam( metric_type=MetricType.IP, n_list=100, n_iters=10, use_soar=False, ), ), ( True, True, IVFIndexParam( metric_type=MetricType.L2, n_list=200, n_iters=20, use_soar=True, ), ), ( True, True, DiskAnnIndexParam(metric_type=MetricType.IP, max_degree=32), ), ( True, True, DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32), ), ( True, True, DiskAnnIndexParam(metric_type=MetricType.COSINE, max_degree=32), ), ], indirect=True, ) @pytest.mark.parametrize("doc_num", [2000]) @pytest.mark.parametrize("query_num", [2]) @pytest.mark.parametrize("top_k", [1]) @pytest.mark.skip(reason="known bug") def test_recall_with_single_vector_valid_2000( self, full_collection_new: Collection, doc_num, query_num, top_k, full_schema_new, request, ): full_schema_params = request.getfixturevalue("full_schema_new") # Build per-field metric type map so ground truth uses each field's # actual index metric (fields may fall back to HnswIndexParam/IP). field_metric_map = {} for vector_para in full_schema_params.vectors: if vector_para.index_param is not None: field_metric_map[vector_para.name] = vector_para.index_param.metric_type else: field_metric_map[vector_para.name] = MetricType.IP metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP) multiple_docs = [ generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num) ] print("len(multiple_docs):\n") print(len(multiple_docs)) # print(multiple_docs) for i in range(10): if i != 0: pass # print(multiple_docs[i * 1000:1000 * (i + 1)]) batchdoc_and_check( full_collection_new, multiple_docs[i * 1000 : 1000 * (i + 1)], operator="insert", ) stats = full_collection_new.stats assert stats.doc_count == len(multiple_docs) doc_ids = ["0", "1"] fetched_docs = full_collection_new.fetch(doc_ids) print("fetched_docs,multiple_docs") print( fetched_docs[doc_ids[0]].vectors["sparse_vector_fp32_field"], fetched_docs[doc_ids[0]].vectors["sparse_vector_fp16_field"], fetched_docs[doc_ids[1]].vectors["sparse_vector_fp32_field"], fetched_docs[doc_ids[1]].vectors["sparse_vector_fp16_field"], "\n", multiple_docs[0].vectors["sparse_vector_fp32_field"], multiple_docs[0].vectors["sparse_vector_fp32_field"], multiple_docs[1].vectors["sparse_vector_fp32_field"], multiple_docs[1].vectors["sparse_vector_fp16_field"], ) full_collection_new.optimize(option=OptimizeOption()) time.sleep(2) query_vectors_map = {} for field_name in DEFAULT_VECTOR_FIELD_NAME.values(): query_vectors_map[field_name] = [ multiple_docs[i].vectors[field_name] for i in range(query_num) ] # Get ground truth mapping (pass per-field metric map) ground_truth_map = get_ground_truth_map( full_collection_new, multiple_docs, query_vectors_map, field_metric_map, top_k, ) # Validate ground truth mapping structure for field_name in DEFAULT_VECTOR_FIELD_NAME.values(): assert field_name in ground_truth_map field_gt = ground_truth_map[field_name] assert len(field_gt) == query_num for query_idx in range(query_num): assert query_idx in field_gt relevant_ids = field_gt[query_idx] assert isinstance(relevant_ids, list) assert len(relevant_ids) <= top_k # Print ground truth statistics print(f"Ground Truth for Top-{top_k} Retrieval:") for field_name, field_gt in ground_truth_map.items(): print(f" {field_name}:") for query_idx, relevant_ids in field_gt.items(): print( f" Query {query_idx}: {len(relevant_ids)} relevant docs - {relevant_ids[:5]}{'...' if len(relevant_ids) > 5 else ''}" ) # Calculate Recall@K using ground truth recall_at_k_stats = calculate_recall_at_k( full_collection_new, multiple_docs, query_vectors_map, full_schema_new, k=top_k, expected_doc_ids_scores_map=ground_truth_map, tolerance=0.01, ) print("ground_truth_map:\n") print(ground_truth_map) print("(recall_at_k_stats:\n") print(recall_at_k_stats) print("field_metric_map:") print(field_metric_map) # Print Recall@K statistics print(f"Recall@{top_k} using Ground Truth:") for field_name, stats in recall_at_k_stats.items(): print(f" {field_name}:") print( f" Relevant Retrieved: {stats['relevant_retrieved_count']}/{stats['total_relevant_count']}" ) print(f" Recall@{top_k}: {stats['recall_at_k']:.4f}") for k, v in recall_at_k_stats.items(): assert v["recall_at_k"] == 1.0