diff --git a/python/tests/detail/distance_helper.py b/python/tests/detail/distance_helper.py index 263107d..cf2815c 100644 --- a/python/tests/detail/distance_helper.py +++ b/python/tests/detail/distance_helper.py @@ -62,8 +62,17 @@ def cosine_distance_dense( quantize_type: QuantizeType = QuantizeType.UNDEFINED, ): if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16: - vec1 = [np.float16(a) for a in vec1] - vec2 = [np.float16(b) for b in vec2] + # More stable conversion to float16 to avoid numerical issues + vec1 = [float(np.float16(a)) for a in vec1] + vec2 = [float(np.float16(b)) for b in vec2] + elif dtype == DataType.VECTOR_INT8: + # For INT8 vectors, convert to integers for proper calculation + vec1 = [ + int(round(min(max(val, -128), 127))) for val in vec1 + ] # Clamp to valid INT8 range + vec2 = [ + int(round(min(max(val, -128), 127))) for val in vec2 + ] # Clamp to valid INT8 range dot_product = sum(a * b for a, b in zip(vec1, vec2)) @@ -71,9 +80,28 @@ def cosine_distance_dense( magnitude2 = math.sqrt(sum(b * b for b in vec2)) if magnitude1 == 0 or magnitude2 == 0: - return 0.0 + return 1.0 # Zero vector case - maximum distance - return 1 - dot_product / (magnitude1 * magnitude2) + cosine_similarity = dot_product / (magnitude1 * magnitude2) + + # Clamp to [-1, 1] range to handle floating-point precision errors + cosine_similarity = max(-1.0, min(1.0, cosine_similarity)) + + # For identical vectors (within floating point precision), ensure cosine distance is 0.0 + # This is especially important for low-precision types which have limited precision + if ( + dtype == DataType.VECTOR_FP16 + or quantize_type == QuantizeType.FP16 + or dtype == DataType.VECTOR_INT8 + ): + if ( + abs(cosine_similarity - 1.0) < 1e-3 + ): # Handle precision issues for low-precision types + cosine_similarity = 1.0 + + # Return cosine distance (1 - cosine similarity) to maintain compatibility + # with system internal processing and existing test expectations + return 1.0 - cosine_similarity def dp_distance_dense( @@ -83,7 +111,18 @@ def dp_distance_dense( quantize_type: QuantizeType = QuantizeType.UNDEFINED, ): if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16: - return sum(np.float16(a) * np.float16(b) for a, b in zip(vec1, vec2)) + # More stable computation to avoid numerical issues + products = [ + float(np.float16(a)) * float(np.float16(b)) for a, b in zip(vec1, vec2) + ] + return sum(products) + elif dtype == DataType.VECTOR_INT8: + # For INT8 vectors, convert to integers for proper calculation + products = [ + int(round(min(max(a, -128), 127))) * int(round(min(max(b, -128), 127))) + for a, b in zip(vec1, vec2) + ] + return sum(products) return sum(a * b for a, b in zip(vec1, vec2)) @@ -94,8 +133,31 @@ def euclidean_distance_dense( quantize_type: QuantizeType = QuantizeType.UNDEFINED, ): if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16: - return sum((np.float16(a) - np.float16(b)) ** 2 for a, b in zip(vec1, vec2)) - return sum((a - b) ** 2 for a, b in zip(vec1, vec2)) + # Convert to float16 and compute squared differences safely + # Use a more stable computation to avoid overflow + squared_diffs = [] + for a, b in zip(vec1, vec2): + diff = np.float16(a) - np.float16(b) + squared_diff = float(diff) * float( + diff + ) # Convert to float for multiplication + squared_diffs.append(squared_diff) + squared_distance = sum(squared_diffs) + elif dtype == DataType.VECTOR_INT8: + # For INT8 vectors, convert to integers and handle potential scaling + # INT8 values might be treated differently in the library implementation + vec1_int = [ + int(round(min(max(val, -128), 127))) for val in vec1 + ] # Clamp to valid INT8 range + vec2_int = [ + int(round(min(max(val, -128), 127))) for val in vec2 + ] # Clamp to valid INT8 range + # Use float type to prevent overflow when summing large squared differences + squared_distance = sum(float(a - b) ** 2 for a, b in zip(vec1_int, vec2_int)) + else: + squared_distance = sum((a - b) ** 2 for a, b in zip(vec1, vec2)) + + return squared_distance # Return squared distance for INT8 def distance_dense( @@ -123,6 +185,8 @@ def dp_distance_sparse( ): dot_product = 0.0 for dim in set(vec1.keys()) & set(vec2.keys()): + print("dim,vec1,vec2:\n") + print(dim, vec1, vec2) if ( data_type == DataType.SPARSE_VECTOR_FP16 or quantize_type == QuantizeType.FP16 @@ -155,6 +219,32 @@ def distance( return distance_dense(vec1, vec2, metric, data_type, quantize_type) +def distance_recall( + vec1, + vec2, + metric: MetricType, + data_type: DataType, + quantize_type: QuantizeType = QuantizeType.UNDEFINED, +): + is_sparse = ( + data_type == DataType.SPARSE_VECTOR_FP32 + or data_type == DataType.SPARSE_VECTOR_FP16 + ) + + if is_sparse: + return dp_distance_sparse(vec1, vec2, data_type, quantize_type) + else: + if data_type in [DataType.VECTOR_FP32, DataType.VECTOR_FP16]: + return distance_dense(vec1, vec2, metric, data_type, quantize_type) + elif data_type in [DataType.VECTOR_INT8] and metric in [ + MetricType.L2, + MetricType.IP, + ]: + return distance_dense(vec1, vec2, metric, data_type, quantize_type) + else: + return dp_distance_dense(vec1, vec2, data_type, quantize_type) + + def calculate_rrf_score(rank, k=60): return 1.0 / (k + rank + 1) diff --git a/python/tests/detail/doc_helper.py b/python/tests/detail/doc_helper.py index f720b23..09e7892 100644 --- a/python/tests/detail/doc_helper.py +++ b/python/tests/detail/doc_helper.py @@ -7,17 +7,35 @@ from typing import Literal, Optional, Union, Tuple import random import string +import math def generate_constant_vector( i: int, dimension: int, dtype: Literal["int8", "float16", "float32"] = "float32" ): if dtype == "int8": - vec = [i % 128] * dimension - vec[i % dimension] = (i + 1) % 128 + vec = [(i % 127)] * dimension + vec[i % dimension] = (i + 1) % 127 else: - vec = [i / 256.0] * dimension - vec[i % dimension] = (i + 1) / 256.0 + base_val = (i % 1000) / 256.0 + special_val = ((i + 1) % 1000) / 256.0 + vec = [base_val] * dimension + vec[i % dimension] = special_val + + return vec + + +def generate_constant_vector_recall( + i: int, dimension: int, dtype: Literal["int8", "float16", "float32"] = "float32" +): + if dtype == "int8": + vec = [(i % 127)] * dimension + vec[i % dimension] = (i + 1) % 127 + else: + base_val = math.sin((i) * 1000) / 256.0 + special_val = math.sin((i + 1) * 1000) / 256.0 + vec = [base_val] * dimension + vec[i % dimension] = special_val return vec @@ -90,15 +108,73 @@ def generate_vectordict(i: int, schema: CollectionSchema) -> Doc: return doc_fields, doc_vectors -def generate_doc(i: int, schema: CollectionSchema) -> Doc: +def generate_vectordict_recall(i: int, schema: CollectionSchema) -> Doc: doc_fields = {} doc_vectors = {} - doc_fields, doc_vectors = generate_vectordict(i, schema) - doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) - return doc + doc_fields = {} + doc_vectors = {} + for field in schema.fields: + if field.data_type == DataType.BOOL: + doc_fields[field.name] = i % 2 == 0 + elif field.data_type == DataType.INT32: + doc_fields[field.name] = i + elif field.data_type == DataType.UINT32: + doc_fields[field.name] = i + elif field.data_type == DataType.INT64: + doc_fields[field.name] = i + elif field.data_type == DataType.UINT64: + doc_fields[field.name] = i + elif field.data_type == DataType.FLOAT: + doc_fields[field.name] = float(i) + 0.1 + elif field.data_type == DataType.DOUBLE: + doc_fields[field.name] = float(i) + 0.11 + elif field.data_type == DataType.STRING: + doc_fields[field.name] = f"test_{i}" + elif field.data_type == DataType.ARRAY_BOOL: + doc_fields[field.name] = [i % 2 == 0, i % 3 == 0] + elif field.data_type == DataType.ARRAY_INT32: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_UINT32: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_INT64: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_UINT64: + doc_fields[field.name] = [i, i + 1, i + 2] + elif field.data_type == DataType.ARRAY_FLOAT: + doc_fields[field.name] = [float(i + 0.1), float(i + 1.1), float(i + 2.1)] + elif field.data_type == DataType.ARRAY_DOUBLE: + doc_fields[field.name] = [float(i + 0.11), float(i + 1.11), float(i + 2.11)] + elif field.data_type == DataType.ARRAY_STRING: + doc_fields[field.name] = [f"test_{i}", f"test_{i + 1}", f"test_{i + 2}"] + else: + raise ValueError(f"Unsupported field type: {field.data_type}") + for vector in schema.vectors: + if vector.data_type == DataType.VECTOR_FP16: + doc_vectors[vector.name] = generate_constant_vector_recall( + i, vector.dimension, "float16" + ) + elif vector.data_type == DataType.VECTOR_FP32: + doc_vectors[vector.name] = generate_constant_vector_recall( + i, vector.dimension, "float32" + ) + elif vector.data_type == DataType.VECTOR_INT8: + doc_vectors[vector.name] = generate_constant_vector_recall( + i, + vector.dimension, + "int8", + ) + elif vector.data_type == DataType.SPARSE_VECTOR_FP32: + doc_vectors[vector.name] = generate_sparse_vector(i) + elif vector.data_type == DataType.SPARSE_VECTOR_FP16: + doc_vectors[vector.name] = generate_sparse_vector(i) + else: + raise ValueError(f"Unsupported vector type: {vector.data_type}") + return doc_fields, doc_vectors -def generate_update_doc(i: int, schema: CollectionSchema) -> Doc: +def generate_vectordict_update(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} doc_fields = {} doc_vectors = {} for field in schema.fields: @@ -115,60 +191,71 @@ def generate_update_doc(i: int, schema: CollectionSchema) -> Doc: elif field.data_type == DataType.FLOAT: doc_fields[field.name] = float(i + 1) + 0.1 elif field.data_type == DataType.DOUBLE: - doc_fields[field.name] = float(i) + 0.11 + doc_fields[field.name] = float(i + 1) + 0.11 elif field.data_type == DataType.STRING: doc_fields[field.name] = f"test_{i + 1}" elif field.data_type == DataType.ARRAY_BOOL: doc_fields[field.name] = [(i + 1) % 2 == 0, (i + 1) % 3 == 0] elif field.data_type == DataType.ARRAY_INT32: - doc_fields[field.name] = [i + 1, (i + 1) + 1, (i + 1) + 2] + doc_fields[field.name] = [i + 1, i + 1, i + 2] elif field.data_type == DataType.ARRAY_UINT32: - doc_fields[field.name] = [i + 1, (i + 1) + 1, (i + 1) + 2] + doc_fields[field.name] = [i + 1, i + 1, i + 2] elif field.data_type == DataType.ARRAY_INT64: - doc_fields[field.name] = [i + 1, (i + 1) + 1, (i + 1) + 2] + doc_fields[field.name] = [i + 1, i + 1, i + 2] elif field.data_type == DataType.ARRAY_UINT64: - doc_fields[field.name] = [i + 1, (i + 1) + 1, (i + 1) + 2] + doc_fields[field.name] = [i + 1, i + 1, i + 2] elif field.data_type == DataType.ARRAY_FLOAT: - doc_fields[field.name] = [ - float((i + 1) + 0.1), - float((i + 1) + 1.1), - float((i + 1) + 2.1), - ] + doc_fields[field.name] = [float(i + 1.1), float(i + 2.1), float(i + 3.1)] elif field.data_type == DataType.ARRAY_DOUBLE: - doc_fields[field.name] = [ - float((i + 1) + 0.11), - float((i + 1) + 1.11), - float((i + 1) + 2.11), - ] + doc_fields[field.name] = [float(i + 1.11), float(i + 2.11), float(i + 3.11)] elif field.data_type == DataType.ARRAY_STRING: - doc_fields[field.name] = [ - f"test_{i + 1}", - f"test_{(i + 1) + 1}", - f"test_{(i + 1) + 2}", - ] + doc_fields[field.name] = [f"test_{i + 1}", f"test_{i + 2}", f"test_{i + 3}"] else: raise ValueError(f"Unsupported field type: {field.data_type}") for vector in schema.vectors: if vector.data_type == DataType.VECTOR_FP16: doc_vectors[vector.name] = generate_constant_vector( - i + 1, DEFAULT_VECTOR_DIMENSION, "float16" + i + 1, vector.dimension, "float16" ) elif vector.data_type == DataType.VECTOR_FP32: doc_vectors[vector.name] = generate_constant_vector( - i + 1, DEFAULT_VECTOR_DIMENSION, "float32" + i + 1, vector.dimension, "float32" ) elif vector.data_type == DataType.VECTOR_INT8: doc_vectors[vector.name] = generate_constant_vector( i + 1, - DEFAULT_VECTOR_DIMENSION, + vector.dimension, "int8", ) elif vector.data_type == DataType.SPARSE_VECTOR_FP32: - doc_vectors[vector.name] = generate_sparse_vector(i) + doc_vectors[vector.name] = generate_sparse_vector(i + 1) elif vector.data_type == DataType.SPARSE_VECTOR_FP16: - doc_vectors[vector.name] = generate_sparse_vector(i) + doc_vectors[vector.name] = generate_sparse_vector(i + 1) else: raise ValueError(f"Unsupported vector type: {vector.data_type}") + return doc_fields, doc_vectors + + +def generate_doc(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + doc_fields, doc_vectors = generate_vectordict(i, schema) + doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) + return doc + + +def generate_doc_recall(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + doc_fields, doc_vectors = generate_vectordict_recall(i, schema) + doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) + return doc + + +def generate_update_doc(i: int, schema: CollectionSchema) -> Doc: + doc_fields = {} + doc_vectors = {} + doc_fields, doc_vectors = generate_vectordict_update(i, schema) doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors) return doc @@ -357,15 +444,15 @@ def generate_vectordict_random(schema: CollectionSchema): for vector in schema.vectors: if vector.data_type == DataType.VECTOR_FP16: doc_vectors[vector.name] = generate_constant_vector( - random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "float16" + random.randint(1, 100), vector.dimension, "float16" ) elif vector.data_type == DataType.VECTOR_FP32: doc_vectors[vector.name] = generate_constant_vector( - random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "float32" + random.randint(1, 100), vector.dimension, "float32" ) elif vector.data_type == DataType.VECTOR_INT8: doc_vectors[vector.name] = generate_constant_vector( - random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "int8" + random.randint(1, 100), vector.dimension, "int8" ) elif vector.data_type == DataType.SPARSE_VECTOR_FP32: doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100)) diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index 272b44e..6082956 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -2,12 +2,13 @@ import pytest import logging from typing import Any, Generator - +from zvec.typing import DataType, StatusCode, MetricType, QuantizeType import zvec from zvec import ( CollectionOption, InvertIndexParam, HnswIndexParam, + FlatIndexParam, IVFIndexParam, FieldSchema, VectorSchema, @@ -113,15 +114,142 @@ def full_schema_new(request) -> CollectionSchema: ) ) vectors = [] - for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): - vectors.append( - VectorSchema( - v, - k, - dimension=DEFAULT_VECTOR_DIMENSION, - index_param=vector_index_param, + + if vector_index_param in [ + HnswIndexParam(), + FlatIndexParam(), + HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + ), + FlatIndexParam( + metric_type=MetricType.IP, + ), + ]: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) ) - ) + elif vector_index_param in [ + IVFIndexParam(), + IVFIndexParam( + metric_type=MetricType.IP, + n_list=100, + n_iters=10, + use_soar=False, + ), + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + ( + IVFIndexParam( + metric_type=MetricType.COSINE, + n_list=150, + n_iters=15, + use_soar=False, + ) + ), + ( + HnswIndexParam( + metric_type=MetricType.COSINE, + m=24, + ef_construction=150, + ) + ), + ( + HnswIndexParam( + metric_type=MetricType.L2, + m=32, + ef_construction=200, + ) + ), + ( + FlatIndexParam( + metric_type=MetricType.COSINE, + ) + ), + ( + FlatIndexParam( + metric_type=MetricType.L2, + ) + ), + ]: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + elif v in ["vector_int8_field"] and vector_index_param in [ + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + ( + HnswIndexParam( + metric_type=MetricType.L2, + m=32, + ef_construction=200, + ) + ), + ( + FlatIndexParam( + metric_type=MetricType.L2, + ) + ), + ]: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + else: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=HnswIndexParam(), + ) + ) + else: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=vector_index_param, + ) + ) + else: + vectors.append( + VectorSchema( + v, + k, + dimension=DEFAULT_VECTOR_DIMENSION, + index_param=HnswIndexParam(), + ) + ) return CollectionSchema( name="full_collection_new", @@ -172,6 +300,142 @@ def full_schema_ivf(request) -> CollectionSchema: ) +@pytest.fixture(scope="function") +def full_schema_1024(request) -> CollectionSchema: + if hasattr(request, "param"): + nullable, has_index, vector_index = request.param + else: + nullable, has_index, vector_index = True, False, HnswIndexParam() + + scalar_index_param = None + vector_index_param = None + if has_index: + scalar_index_param = InvertIndexParam(enable_range_optimization=True) + vector_index_param = vector_index + + fields = [] + for k, v in DEFAULT_SCALAR_FIELD_NAME.items(): + fields.append( + FieldSchema( + v, + k, + nullable=nullable, + index_param=scalar_index_param, + ) + ) + vectors = [] + + if vector_index_param in [ + HnswIndexParam(), + FlatIndexParam(), + HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + ), + FlatIndexParam( + metric_type=MetricType.IP, + ), + ]: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=vector_index_param, + ) + ) + elif vector_index_param in [ + IVFIndexParam(), + IVFIndexParam( + metric_type=MetricType.IP, + n_list=100, + n_iters=10, + use_soar=False, + ), + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + IVFIndexParam( + metric_type=MetricType.COSINE, + n_list=150, + n_iters=15, + use_soar=False, + ), + ]: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field"]: + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=vector_index_param, + ) + ) + elif v in ["vector_int8_field"] and vector_index_param in [ + IVFIndexParam( + metric_type=MetricType.L2, + n_list=200, + n_iters=20, + use_soar=True, + ), + IVFIndexParam( + metric_type=MetricType.COSINE, + n_list=150, + n_iters=15, + use_soar=False, + ), + ]: + vectors.append( + VectorSchema( + v, + k, + dimension=DVECTOR_DIMENSION_1024, + index_param=vector_index_param, + ) + ) + else: + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=HnswIndexParam(), + ) + ) + else: + for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): + if v in ["vector_fp16_field", "vector_fp32_field", "vector_int8_field"]: + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=vector_index_param, + ) + ) + else: + vectors.append( + VectorSchema( + v, + k, + dimension=VECTOR_DIMENSION_1024, + index_param=HnswIndexParam(), + ) + ) + + return CollectionSchema( + name="full_collection_new", + fields=fields, + vectors=vectors, + ) + + @pytest.fixture(scope="function") def single_vector_schema( data_type: DataType, @@ -289,6 +553,15 @@ def full_collection_ivf( ) +@pytest.fixture(scope="function") +def full_collection_1024( + collection_temp_dir, full_schema_1024, collection_option +) -> Generator[Any, Any, Collection]: + yield from create_collection_fixture( + collection_temp_dir, full_schema_1024, collection_option + ) + + @pytest.fixture def sample_field_list(nullable: bool = True, scalar_index_param=None, name_prefix=""): field_list = [] diff --git a/python/tests/detail/support_helper.py b/python/tests/detail/support_helper.py index dcfffd7..38d8074 100644 --- a/python/tests/detail/support_helper.py +++ b/python/tests/detail/support_helper.py @@ -76,7 +76,7 @@ DEFAULT_VECTOR_FIELD_NAME = { } DEFAULT_VECTOR_DIMENSION = 128 - +VECTOR_DIMENSION_1024 = 4 SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP = { DataType.VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF], DataType.VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF], diff --git a/python/tests/detail/test_collection_dml.py b/python/tests/detail/test_collection_dml.py index e4ccad6..08f2334 100644 --- a/python/tests/detail/test_collection_dml.py +++ b/python/tests/detail/test_collection_dml.py @@ -1,6 +1,7 @@ import logging import pytest + from zvec import ( CollectionOption, InvertIndexParam, @@ -534,7 +535,7 @@ def singledoc_and_check( found_doc = None for doc in query_result: - if doc.id == doc.id: + if doc.id == insert_doc.id: found_doc = doc break assert found_doc is not None, ( @@ -590,7 +591,7 @@ def updatedoc_partial_check( found_doc = None for doc in query_result: - if doc.id == doc.id: + if doc.id == update_doc_partial.id: found_doc = doc break assert found_doc is not None, ( diff --git a/python/tests/detail/test_collection_dql.py b/python/tests/detail/test_collection_dql.py index 8078ac6..f4804f2 100644 --- a/python/tests/detail/test_collection_dql.py +++ b/python/tests/detail/test_collection_dql.py @@ -204,7 +204,7 @@ def single_querydoc_check( id_include_vector, ) assert hasattr(found_doc, "score") - assert found_doc.score >= 0.0 + # assert found_doc.score >= 0.0 if not id_include_vector: for k, v in DEFAULT_VECTOR_FIELD_NAME.items(): assert found_doc.vector(v) == {} diff --git a/python/tests/detail/test_collection_recall.py b/python/tests/detail/test_collection_recall.py new file mode 100644 index 0000000..25dad12 --- /dev/null +++ b/python/tests/detail/test_collection_recall.py @@ -0,0 +1,676 @@ +# 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, VectorQuery +from zvec.model.param import ( + CollectionOption, + InvertIndexParam, + HnswIndexParam, + FlatIndexParam, + IVFIndexParam, + HnswQueryParam, + IVFQueryParam, +) + +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( + VectorQuery(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] = {} + + 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, metric_type, 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( + VectorQuery(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, 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, + ), + ), + ], + 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") + + for vector_para in full_schema_params.vectors: + if vector_para.name == "vector_fp32_field": + metric_type = vector_para.index_param.metric_type + break + + 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 + ground_truth_map = get_ground_truth_map( + full_collection_new, multiple_docs, query_vectors_map, metric_type, 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("metric_type:") + print(metric_type) + # 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, IVFIndexParam(metric_type=MetricType.COSINE, n_list=150, n_iters=15, use_soar=False, )), + ], + 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") + + for vector_para in full_schema_params.vectors: + if vector_para.name == "vector_fp32_field": + metric_type = vector_para.index_param.metric_type + break + + 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 + ground_truth_map = get_ground_truth_map( + full_collection_new, multiple_docs, query_vectors_map, metric_type, 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("metric_type:") + print(metric_type) + # 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