chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,392 @@
|
||||
import logging
|
||||
import math
|
||||
import numpy as np
|
||||
|
||||
from zvec import (
|
||||
MetricType,
|
||||
DataType,
|
||||
QuantizeType,
|
||||
Doc,
|
||||
CollectionSchema,
|
||||
FieldSchema,
|
||||
VectorSchema,
|
||||
)
|
||||
|
||||
from typing import Dict
|
||||
|
||||
|
||||
def is_float_equal(actual, expected, rel_tol=1e-5, abs_tol=1e-8):
|
||||
if actual is None and expected is None:
|
||||
return True
|
||||
return math.isclose(actual, expected, rel_tol=rel_tol, abs_tol=abs_tol)
|
||||
|
||||
|
||||
def is_dense_vector_equal(vec1, vec2, rtol=1e-5, atol=1e-8):
|
||||
"""Compare two dense vectors with tolerance."""
|
||||
return np.allclose(vec1, vec2, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
def is_sparse_vector_equal(vec1, vec2, rtol=1e-5, atol=1e-8):
|
||||
"""Compare two sparse vectors with tolerance."""
|
||||
# Check if they have the same keys
|
||||
if set(vec1.keys()) != set(vec2.keys()):
|
||||
return False
|
||||
|
||||
# Check if all values are close
|
||||
for key in vec1:
|
||||
if not math.isclose(vec1[key], vec2[key], rel_tol=rtol, abs_tol=atol):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def is_float_array_equal(arr1, arr2, rtol=1e-5, atol=1e-8):
|
||||
"""Compare two float arrays with tolerance."""
|
||||
return np.allclose(arr1, arr2, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
def is_double_array_equal(arr1, arr2, rtol=1e-9, atol=1e-12):
|
||||
"""Compare two double arrays with tolerance."""
|
||||
return np.allclose(arr1, arr2, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
def is_int_array_equal(arr1, arr2):
|
||||
"""Compare two integer arrays with exact equality."""
|
||||
return np.array_equal(arr1, arr2)
|
||||
|
||||
|
||||
def cosine_distance_dense(
|
||||
vec1,
|
||||
vec2,
|
||||
dtype: DataType = DataType.VECTOR_FP32,
|
||||
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
|
||||
):
|
||||
if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16:
|
||||
# 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))
|
||||
|
||||
magnitude1 = math.sqrt(sum(a * a for a in vec1))
|
||||
magnitude2 = math.sqrt(sum(b * b for b in vec2))
|
||||
|
||||
if magnitude1 == 0 or magnitude2 == 0:
|
||||
return 1.0 # Zero vector case - maximum distance
|
||||
|
||||
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(
|
||||
vec1,
|
||||
vec2,
|
||||
dtype: DataType = DataType.VECTOR_FP32,
|
||||
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
|
||||
):
|
||||
if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16:
|
||||
# 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))
|
||||
|
||||
|
||||
def euclidean_distance_dense(
|
||||
vec1,
|
||||
vec2,
|
||||
dtype: DataType = DataType.VECTOR_FP32,
|
||||
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
|
||||
):
|
||||
if dtype == DataType.VECTOR_FP16 or quantize_type == QuantizeType.FP16:
|
||||
# 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(
|
||||
vec1,
|
||||
vec2,
|
||||
metric: MetricType,
|
||||
data_type: DataType = DataType.VECTOR_FP32,
|
||||
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
|
||||
):
|
||||
if metric == MetricType.COSINE:
|
||||
return cosine_distance_dense(vec1, vec2, data_type, quantize_type)
|
||||
elif metric == MetricType.L2:
|
||||
return euclidean_distance_dense(vec1, vec2, data_type, quantize_type)
|
||||
elif metric == MetricType.IP:
|
||||
return dp_distance_dense(vec1, vec2, data_type, quantize_type)
|
||||
else:
|
||||
raise ValueError("Unsupported metric type")
|
||||
|
||||
|
||||
def dp_distance_sparse(
|
||||
vec1,
|
||||
vec2,
|
||||
data_type: DataType = DataType.SPARSE_VECTOR_FP32,
|
||||
quantize_type: QuantizeType = QuantizeType.UNDEFINED,
|
||||
):
|
||||
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
|
||||
):
|
||||
vec1[dim] = np.float16(vec1[dim])
|
||||
vec2[dim] = np.float16(vec2[dim])
|
||||
dot_product += vec1[dim] * vec2[dim]
|
||||
return dot_product
|
||||
|
||||
|
||||
def distance(
|
||||
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:
|
||||
if metric != MetricType.IP:
|
||||
raise ValueError("Unsupported metric type for sparse vectors")
|
||||
|
||||
if is_sparse:
|
||||
return dp_distance_sparse(vec1, vec2, data_type, quantize_type)
|
||||
else:
|
||||
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)
|
||||
|
||||
|
||||
def calculate_multi_vector_rrf_scores(query_results: Dict[str, Doc], k=60):
|
||||
rrf_scores = {}
|
||||
|
||||
for vector_name, docs in query_results.items():
|
||||
for rank, doc in enumerate(docs):
|
||||
doc_id = doc.id
|
||||
rrf_score = calculate_rrf_score(rank, k)
|
||||
if doc_id in rrf_scores:
|
||||
rrf_scores[doc_id] += rrf_score
|
||||
else:
|
||||
rrf_scores[doc_id] = rrf_score
|
||||
|
||||
return rrf_scores
|
||||
|
||||
|
||||
def calculate_multi_vector_weighted_scores(
|
||||
query_results: Dict[str, Doc], weights: Dict[str, float], metric: MetricType
|
||||
):
|
||||
def _normalize_score(score: float, metric: MetricType) -> float:
|
||||
if metric == MetricType.L2:
|
||||
return 1.0 - 2 * math.atan(score) / math.pi
|
||||
if metric == MetricType.IP:
|
||||
return 0.5 + math.atan(score) / math.pi
|
||||
if metric == MetricType.COSINE:
|
||||
return 1.0 - score / 2.0
|
||||
raise ValueError("Unsupported metric type")
|
||||
|
||||
weighted_scores = {}
|
||||
|
||||
for vector_name, docs in query_results.items():
|
||||
weight = weights.get(vector_name, 1.0)
|
||||
|
||||
for doc in docs:
|
||||
doc_id = doc.id
|
||||
weighted_score = (_normalize_score(doc.score, metric)) * weight
|
||||
if doc_id in weighted_scores:
|
||||
weighted_scores[doc_id] += weighted_score
|
||||
else:
|
||||
weighted_scores[doc_id] = weighted_score
|
||||
|
||||
return weighted_scores
|
||||
|
||||
|
||||
def is_field_equal(field1, field2, schema: FieldSchema) -> bool:
|
||||
if field1 is None and field2 is None:
|
||||
return True
|
||||
if field1 is None or field2 is None:
|
||||
return False
|
||||
|
||||
if schema.data_type == DataType.ARRAY_FLOAT:
|
||||
return is_float_array_equal(field1, field2)
|
||||
elif schema.data_type == DataType.ARRAY_DOUBLE:
|
||||
return is_double_array_equal(field1, field2)
|
||||
elif schema.data_type in [
|
||||
DataType.ARRAY_INT32,
|
||||
DataType.ARRAY_INT64,
|
||||
DataType.ARRAY_BOOL,
|
||||
DataType.ARRAY_STRING,
|
||||
DataType.ARRAY_UINT32,
|
||||
DataType.ARRAY_UINT64,
|
||||
DataType.ARRAY_INT64,
|
||||
]:
|
||||
return is_int_array_equal(field1, field2)
|
||||
elif schema.data_type in [DataType.FLOAT, DataType.DOUBLE]:
|
||||
return is_float_equal(field1, field2)
|
||||
|
||||
return field1 == field2
|
||||
|
||||
|
||||
def is_vector_equal(vec1, vec2, schema: VectorSchema) -> bool:
|
||||
if (
|
||||
schema.data_type == DataType.SPARSE_VECTOR_FP16
|
||||
or schema.data_type == DataType.VECTOR_FP16
|
||||
):
|
||||
# skip fp16 vector equal
|
||||
return True
|
||||
|
||||
is_sparse = (
|
||||
schema.data_type == DataType.SPARSE_VECTOR_FP32
|
||||
or schema.data_type == DataType.SPARSE_VECTOR_FP16
|
||||
)
|
||||
|
||||
if is_sparse:
|
||||
return is_sparse_vector_equal(vec1, vec2)
|
||||
else:
|
||||
return is_dense_vector_equal(vec1, vec2)
|
||||
|
||||
|
||||
def is_doc_equal(
|
||||
doc1: Doc,
|
||||
doc2: Doc,
|
||||
schema: CollectionSchema,
|
||||
except_score: bool = True,
|
||||
include_vector: bool = True,
|
||||
):
|
||||
if doc1.id != doc2.id:
|
||||
logging.error("doc ids are not equal")
|
||||
return False
|
||||
|
||||
reduce_field_names = set(doc1.field_names() + doc2.field_names())
|
||||
reduce_vector_names = set(doc1.vector_names() + doc2.vector_names())
|
||||
|
||||
is_doc1_fields_empty = doc1.fields is None or doc1.fields == {}
|
||||
is_doc2_fields_empty = doc2.fields is None or doc2.fields == {}
|
||||
|
||||
if is_doc1_fields_empty or is_doc2_fields_empty:
|
||||
if is_doc1_fields_empty != is_doc2_fields_empty:
|
||||
return False
|
||||
else:
|
||||
for field_name in reduce_field_names:
|
||||
field_schema = schema.field(field_name)
|
||||
if field_schema is None:
|
||||
return False
|
||||
if is_field_equal(
|
||||
doc1.field(field_name), doc2.field(field_name), field_schema
|
||||
):
|
||||
continue
|
||||
else:
|
||||
logging.error(f"{field_name} are not equal")
|
||||
return False
|
||||
|
||||
if include_vector:
|
||||
is_doc1_vectors_empty = doc1.vectors is None or doc1.vectors == {}
|
||||
is_doc2_vectors_empty = doc2.vectors is None or doc2.vectors == {}
|
||||
|
||||
if is_doc1_vectors_empty or is_doc2_vectors_empty:
|
||||
if is_doc1_fields_empty != is_doc2_vectors_empty:
|
||||
return False
|
||||
else:
|
||||
for vector_name in reduce_vector_names:
|
||||
vector_schema = schema.vector(vector_name)
|
||||
if vector_schema is None:
|
||||
return False
|
||||
if is_vector_equal(
|
||||
doc1.vector(vector_name), doc2.vector(vector_name), vector_schema
|
||||
):
|
||||
continue
|
||||
else:
|
||||
return False
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,464 @@
|
||||
from zvec import CollectionSchema, Doc
|
||||
|
||||
from support_helper import *
|
||||
|
||||
import numpy as np
|
||||
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 % 127)] * dimension
|
||||
vec[i % dimension] = (i + 1) % 127
|
||||
else:
|
||||
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
|
||||
|
||||
|
||||
def generate_sparse_vector(i: int):
|
||||
return {i: i + 0.1}
|
||||
|
||||
|
||||
def generate_vectordict(i: int, schema: CollectionSchema) -> Doc:
|
||||
doc_fields = {}
|
||||
doc_vectors = {}
|
||||
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(
|
||||
i, vector.dimension, "float16"
|
||||
)
|
||||
elif vector.data_type == DataType.VECTOR_FP32:
|
||||
doc_vectors[vector.name] = generate_constant_vector(
|
||||
i, vector.dimension, "float32"
|
||||
)
|
||||
elif vector.data_type == DataType.VECTOR_INT8:
|
||||
doc_vectors[vector.name] = generate_constant_vector(
|
||||
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_vectordict_recall(i: int, schema: CollectionSchema) -> Doc:
|
||||
doc_fields = {}
|
||||
doc_vectors = {}
|
||||
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_vectordict_update(i: int, schema: CollectionSchema) -> Doc:
|
||||
doc_fields = {}
|
||||
doc_vectors = {}
|
||||
doc_fields = {}
|
||||
doc_vectors = {}
|
||||
for field in schema.fields:
|
||||
if field.data_type == DataType.BOOL:
|
||||
doc_fields[field.name] = (i + 1) % 2 == 0
|
||||
elif field.data_type == DataType.INT32:
|
||||
doc_fields[field.name] = i + 1
|
||||
elif field.data_type == DataType.UINT32:
|
||||
doc_fields[field.name] = i + 1
|
||||
elif field.data_type == DataType.INT64:
|
||||
doc_fields[field.name] = i + 1
|
||||
elif field.data_type == DataType.UINT64:
|
||||
doc_fields[field.name] = i + 1
|
||||
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 + 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, i + 2]
|
||||
elif field.data_type == DataType.ARRAY_UINT32:
|
||||
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, i + 2]
|
||||
elif field.data_type == DataType.ARRAY_UINT64:
|
||||
doc_fields[field.name] = [i + 1, i + 1, i + 2]
|
||||
elif field.data_type == DataType.ARRAY_FLOAT:
|
||||
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.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 + 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, vector.dimension, "float16"
|
||||
)
|
||||
elif vector.data_type == DataType.VECTOR_FP32:
|
||||
doc_vectors[vector.name] = generate_constant_vector(
|
||||
i + 1, vector.dimension, "float32"
|
||||
)
|
||||
elif vector.data_type == DataType.VECTOR_INT8:
|
||||
doc_vectors[vector.name] = generate_constant_vector(
|
||||
i + 1,
|
||||
vector.dimension,
|
||||
"int8",
|
||||
)
|
||||
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
|
||||
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 + 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
|
||||
|
||||
|
||||
def generate_doc_random(i, schema: CollectionSchema) -> Doc:
|
||||
doc_fields = {}
|
||||
doc_vectors = {}
|
||||
|
||||
random.seed(i)
|
||||
|
||||
for field in schema.fields:
|
||||
if field.data_type == DataType.BOOL:
|
||||
doc_fields[field.name] = random.choice([True, False])
|
||||
elif field.data_type == DataType.INT32:
|
||||
doc_fields[field.name] = random.randint(-2147483648, 2147483647)
|
||||
elif field.data_type == DataType.UINT32:
|
||||
doc_fields[field.name] = random.randint(0, 4294967295)
|
||||
elif field.data_type == DataType.INT64:
|
||||
doc_fields[field.name] = random.randint(
|
||||
-9223372036854775808, 9223372036854775807
|
||||
)
|
||||
elif field.data_type == DataType.UINT64:
|
||||
doc_fields[field.name] = random.randint(0, 18446744073709551615)
|
||||
elif field.data_type == DataType.FLOAT:
|
||||
doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38)
|
||||
elif field.data_type == DataType.DOUBLE:
|
||||
doc_fields[field.name] = random.uniform(
|
||||
-1.7976931348623157e308, 1.7976931348623157e308
|
||||
)
|
||||
elif field.data_type == DataType.STRING:
|
||||
length = random.randint(1, 999)
|
||||
doc_fields[field.name] = "".join(
|
||||
random.choices(string.ascii_letters + string.digits, k=length)
|
||||
)
|
||||
elif field.data_type == DataType.ARRAY_BOOL:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.choice([True, False]) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_INT32:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.randint(-2147483648, 2147483647) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_UINT32:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.randint(0, 4294967295) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_INT64:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.randint(-9223372036854775808, 9223372036854775807)
|
||||
for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_UINT64:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.randint(0, 18446744073709551615) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_FLOAT:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_DOUBLE:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.uniform(-1.7976931348623157e308, 1.7976931348623157e308)
|
||||
for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_STRING:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
"".join(
|
||||
random.choices(
|
||||
string.ascii_letters + string.digits, k=random.randint(1, 100)
|
||||
)
|
||||
)
|
||||
for _ in range(array_length)
|
||||
]
|
||||
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(
|
||||
random.randint(1, 100), DEFAULT_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"
|
||||
)
|
||||
elif vector.data_type == DataType.VECTOR_INT8:
|
||||
doc_vectors[vector.name] = generate_constant_vector(
|
||||
random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "int8"
|
||||
)
|
||||
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
|
||||
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
|
||||
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
|
||||
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
|
||||
else:
|
||||
raise ValueError(f"Unsupported vector type: {vector.data_type}")
|
||||
|
||||
doc = Doc(id=i, fields=doc_fields, vectors=doc_vectors)
|
||||
return doc
|
||||
|
||||
|
||||
def generate_vectordict_random(schema: CollectionSchema):
|
||||
doc_fields = {}
|
||||
doc_vectors = {}
|
||||
for field in schema.fields:
|
||||
if field.data_type == DataType.BOOL:
|
||||
doc_fields[field.name] = random.choice([True, False])
|
||||
elif field.data_type == DataType.INT32:
|
||||
doc_fields[field.name] = random.randint(-2147483648, 2147483647)
|
||||
elif field.data_type == DataType.UINT32:
|
||||
doc_fields[field.name] = random.randint(0, 4294967295)
|
||||
elif field.data_type == DataType.INT64:
|
||||
doc_fields[field.name] = random.randint(
|
||||
-9223372036854775808, 9223372036854775807
|
||||
)
|
||||
elif field.data_type == DataType.UINT64:
|
||||
doc_fields[field.name] = random.randint(0, 18446744073709551615)
|
||||
elif field.data_type == DataType.FLOAT:
|
||||
doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38)
|
||||
elif field.data_type == DataType.DOUBLE:
|
||||
doc_fields[field.name] = random.uniform(
|
||||
-1.7976931348623157e308, 1.7976931348623157e308
|
||||
)
|
||||
elif field.data_type == DataType.STRING:
|
||||
length = random.randint(1, 999)
|
||||
doc_fields[field.name] = "".join(
|
||||
random.choices(string.ascii_letters + string.digits, k=length)
|
||||
)
|
||||
elif field.data_type == DataType.ARRAY_BOOL:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.choice([True, False]) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_INT32:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.randint(-2147483648, 2147483647) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_UINT32:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.randint(0, 4294967295) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_INT64:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.randint(-9223372036854775808, 9223372036854775807)
|
||||
for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_UINT64:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.randint(0, 18446744073709551615) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_FLOAT:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_DOUBLE:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
random.uniform(-1.7976931348623157e308, 1.7976931348623157e308)
|
||||
for _ in range(array_length)
|
||||
]
|
||||
elif field.data_type == DataType.ARRAY_STRING:
|
||||
array_length = random.randint(0, 10)
|
||||
doc_fields[field.name] = [
|
||||
"".join(
|
||||
random.choices(
|
||||
string.ascii_letters + string.digits, k=random.randint(1, 100)
|
||||
)
|
||||
)
|
||||
for _ in range(array_length)
|
||||
]
|
||||
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(
|
||||
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), vector.dimension, "float32"
|
||||
)
|
||||
elif vector.data_type == DataType.VECTOR_INT8:
|
||||
doc_vectors[vector.name] = generate_constant_vector(
|
||||
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))
|
||||
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
|
||||
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
|
||||
else:
|
||||
raise ValueError(f"Unsupported vector type: {vector.data_type}")
|
||||
|
||||
return doc_fields, doc_vectors
|
||||
@@ -0,0 +1,652 @@
|
||||
import pytest
|
||||
import logging
|
||||
import platform
|
||||
|
||||
DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in (
|
||||
"x86_64",
|
||||
"AMD64",
|
||||
"i686",
|
||||
"i386",
|
||||
)
|
||||
|
||||
from typing import Any, Generator
|
||||
from zvec.typing import DataType, StatusCode, MetricType, QuantizeType
|
||||
import zvec
|
||||
|
||||
|
||||
# Cache the DiskAnn plugin preload status so we pay the load cost once per
|
||||
# test session. The plugin normally auto-loads on first DiskAnn use, but we
|
||||
# preload it explicitly here so a missing libaio / misplaced plugin .so
|
||||
# surfaces as a clear pytest skip instead of a confusing
|
||||
# "Create vector column indexer failed" deep inside the collection code path.
|
||||
_DISKANN_PRELOAD_REASON: str | None = None
|
||||
_DISKANN_PRELOAD_DONE: bool = False
|
||||
|
||||
|
||||
def _ensure_diskann_runtime_or_reason() -> str | None:
|
||||
"""Preload the DiskAnn plugin and return None on success or a human-readable
|
||||
skip reason on failure. Idempotent across calls."""
|
||||
global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON
|
||||
if _DISKANN_PRELOAD_DONE:
|
||||
return _DISKANN_PRELOAD_REASON
|
||||
_DISKANN_PRELOAD_DONE = True
|
||||
|
||||
if not DISKANN_SUPPORTED:
|
||||
_DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64"
|
||||
return _DISKANN_PRELOAD_REASON
|
||||
|
||||
if not zvec.is_libaio_available():
|
||||
_DISKANN_PRELOAD_REASON = (
|
||||
"libaio is not available on this host; DiskAnn cannot run. "
|
||||
"Install libaio1 (or libaio1t64 on Ubuntu 24.04+) and retry."
|
||||
)
|
||||
return _DISKANN_PRELOAD_REASON
|
||||
|
||||
status = zvec.load_diskann_plugin()
|
||||
if status != zvec.DISKANN_PLUGIN_OK:
|
||||
_DISKANN_PRELOAD_REASON = (
|
||||
f"Failed to load DiskAnn plugin (status={status}); "
|
||||
"check that libzvec_diskann_plugin.so is installed alongside "
|
||||
"_zvec.so in the Python site-packages directory."
|
||||
)
|
||||
return _DISKANN_PRELOAD_REASON
|
||||
|
||||
_DISKANN_PRELOAD_REASON = None
|
||||
return None
|
||||
|
||||
|
||||
from zvec import (
|
||||
CollectionOption,
|
||||
InvertIndexParam,
|
||||
HnswIndexParam,
|
||||
FlatIndexParam,
|
||||
IVFIndexParam,
|
||||
FieldSchema,
|
||||
VectorSchema,
|
||||
CollectionSchema,
|
||||
Collection,
|
||||
Doc,
|
||||
Query,
|
||||
)
|
||||
|
||||
from support_helper import *
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def basic_schema(collection_name="test_collection") -> CollectionSchema:
|
||||
return CollectionSchema(
|
||||
name=collection_name if len(collection_name) > 0 else "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),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
),
|
||||
VectorSchema(
|
||||
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def full_schema(
|
||||
nullable: bool = False,
|
||||
has_index: bool = False,
|
||||
) -> CollectionSchema:
|
||||
scalar_index_param = None
|
||||
vector_index_param = None
|
||||
if has_index:
|
||||
scalar_index_param = InvertIndexParam(enable_range_optimization=True)
|
||||
vector_index_param = HnswIndexParam()
|
||||
|
||||
fields = []
|
||||
for k, v in DEFAULT_SCALAR_FIELD_NAME.items():
|
||||
fields.append(
|
||||
FieldSchema(
|
||||
v,
|
||||
k,
|
||||
nullable=nullable,
|
||||
index_param=scalar_index_param,
|
||||
)
|
||||
)
|
||||
vetors = []
|
||||
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
||||
vetors.append(
|
||||
VectorSchema(
|
||||
v,
|
||||
k,
|
||||
dimension=DEFAULT_VECTOR_DIMENSION,
|
||||
index_param=vector_index_param,
|
||||
)
|
||||
)
|
||||
|
||||
return CollectionSchema(
|
||||
name="full_collection",
|
||||
fields=fields,
|
||||
vectors=vetors,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def full_schema_new(request) -> CollectionSchema:
|
||||
if hasattr(request, "param"):
|
||||
nullable, has_index, vector_index = request.param
|
||||
else:
|
||||
nullable, has_index, vector_index = True, False, HnswIndexParam()
|
||||
|
||||
# Skip DiskAnn tests on unsupported platforms or when the runtime cannot
|
||||
# be brought up (missing libaio, plugin .so not installed, etc.).
|
||||
from zvec.model.param import DiskAnnIndexParam
|
||||
|
||||
if isinstance(vector_index, DiskAnnIndexParam):
|
||||
skip_reason = _ensure_diskann_runtime_or_reason()
|
||||
if skip_reason is not None:
|
||||
pytest.skip(skip_reason)
|
||||
|
||||
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=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",
|
||||
fields=fields,
|
||||
vectors=vectors,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def full_schema_ivf(request) -> CollectionSchema:
|
||||
if hasattr(request, "param"):
|
||||
nullable, has_index, vector_index = request.param
|
||||
else:
|
||||
nullable, has_index, vector_index = True, False, IVFIndexParam()
|
||||
|
||||
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 = []
|
||||
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,
|
||||
)
|
||||
)
|
||||
|
||||
return CollectionSchema(
|
||||
name="full_collection_ivf",
|
||||
fields=fields,
|
||||
vectors=vectors,
|
||||
)
|
||||
|
||||
|
||||
@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,
|
||||
) -> CollectionSchema:
|
||||
vector_schema = [
|
||||
VectorSchema(
|
||||
DEFAULT_VECTOR_FIELD_NAME[data_type],
|
||||
data_type,
|
||||
DEFAULT_VECTOR_DIMENSION,
|
||||
)
|
||||
]
|
||||
|
||||
return CollectionSchema(
|
||||
name="full_collection",
|
||||
vectors=vector_schema,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def single_vector_schema_with_index_param(
|
||||
data_type: DataType, index_param
|
||||
) -> CollectionSchema:
|
||||
vector_schema = [
|
||||
VectorSchema(
|
||||
DEFAULT_VECTOR_FIELD_NAME[data_type],
|
||||
data_type,
|
||||
DEFAULT_VECTOR_DIMENSION,
|
||||
index_param,
|
||||
)
|
||||
]
|
||||
|
||||
return CollectionSchema(
|
||||
name="full_collection",
|
||||
vectors=vector_schema,
|
||||
)
|
||||
|
||||
|
||||
def create_collection_fixture(
|
||||
collection_temp_dir, schema: CollectionSchema, collection_option: CollectionOption
|
||||
) -> Generator[Any, Any, Collection]:
|
||||
"""Common helper function to create and manage collection fixtures."""
|
||||
coll = zvec.create_and_open(
|
||||
path=str(collection_temp_dir),
|
||||
schema=schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert coll is not None, "Failed to create and open collection"
|
||||
assert coll.path == str(collection_temp_dir)
|
||||
assert coll.schema.name == schema.name
|
||||
assert list(coll.schema.fields) == list(schema.fields)
|
||||
assert list(coll.schema.vectors) == list(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:
|
||||
logging.warning(f"Warning: failed to destroy collection: {e}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def basic_collection(
|
||||
collection_temp_dir, basic_schema, collection_option
|
||||
) -> Generator[Any, Any, Collection]:
|
||||
yield from create_collection_fixture(
|
||||
collection_temp_dir, basic_schema, collection_option
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def collection_option():
|
||||
return CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def collection_temp_dir(tmp_path_factory):
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "test_collection_path"
|
||||
return str(collection_path)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def full_collection(
|
||||
collection_temp_dir,
|
||||
full_schema,
|
||||
collection_option,
|
||||
nullable: bool = True,
|
||||
has_index: bool = False,
|
||||
) -> Generator[Any, Any, Collection]:
|
||||
yield from create_collection_fixture(
|
||||
collection_temp_dir, full_schema, collection_option
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def full_collection_new(
|
||||
collection_temp_dir, full_schema_new, collection_option
|
||||
) -> Generator[Any, Any, Collection]:
|
||||
yield from create_collection_fixture(
|
||||
collection_temp_dir, full_schema_new, collection_option
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def full_collection_ivf(
|
||||
collection_temp_dir, full_schema_ivf, collection_option
|
||||
) -> Generator[Any, Any, Collection]:
|
||||
yield from create_collection_fixture(
|
||||
collection_temp_dir, full_schema_ivf, collection_option
|
||||
)
|
||||
|
||||
|
||||
@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 = []
|
||||
for k, v in DEFAULT_SCALAR_FIELD_NAME.items():
|
||||
field_list.append(
|
||||
FieldSchema(
|
||||
f"{name_prefix}_{v}" if len(name_prefix) > 0 else v,
|
||||
k,
|
||||
nullable=nullable,
|
||||
index_param=scalar_index_param,
|
||||
)
|
||||
)
|
||||
return field_list
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_vector_list(vector_index_param=None, name_prefix=""):
|
||||
vector_list = []
|
||||
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
|
||||
vector_list.append(
|
||||
VectorSchema(
|
||||
f"{name_prefix}_{v}" if len(name_prefix) > 0 else v,
|
||||
k,
|
||||
dimension=DEFAULT_VECTOR_DIMENSION,
|
||||
index_param=vector_index_param,
|
||||
)
|
||||
)
|
||||
return vector_list
|
||||
@@ -0,0 +1,209 @@
|
||||
from zvec import (
|
||||
CollectionOption,
|
||||
IndexOption,
|
||||
OptimizeOption,
|
||||
InvertIndexParam,
|
||||
HnswIndexParam,
|
||||
IVFIndexParam,
|
||||
FlatIndexParam,
|
||||
AlterColumnOption,
|
||||
AddColumnOption,
|
||||
DataType,
|
||||
MetricType,
|
||||
QuantizeType,
|
||||
)
|
||||
|
||||
|
||||
VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP = {
|
||||
DataType.VECTOR_FP32: [
|
||||
HnswIndexParam(),
|
||||
HnswIndexParam(
|
||||
metric_type=MetricType.IP,
|
||||
m=16,
|
||||
ef_construction=100,
|
||||
quantize_type=QuantizeType.INT8,
|
||||
),
|
||||
HnswIndexParam(
|
||||
metric_type=MetricType.COSINE,
|
||||
m=24,
|
||||
ef_construction=150,
|
||||
quantize_type=QuantizeType.INT4,
|
||||
),
|
||||
HnswIndexParam(
|
||||
metric_type=MetricType.L2,
|
||||
m=32,
|
||||
ef_construction=200,
|
||||
quantize_type=QuantizeType.FP16,
|
||||
),
|
||||
FlatIndexParam(),
|
||||
FlatIndexParam(metric_type=MetricType.IP, quantize_type=QuantizeType.INT4),
|
||||
FlatIndexParam(metric_type=MetricType.L2, quantize_type=QuantizeType.INT8),
|
||||
FlatIndexParam(metric_type=MetricType.COSINE, quantize_type=QuantizeType.FP16),
|
||||
IVFIndexParam(),
|
||||
IVFIndexParam(
|
||||
metric_type=MetricType.IP,
|
||||
quantize_type=QuantizeType.INT4,
|
||||
n_list=100,
|
||||
n_iters=10,
|
||||
use_soar=False,
|
||||
),
|
||||
IVFIndexParam(
|
||||
metric_type=MetricType.L2,
|
||||
quantize_type=QuantizeType.INT8,
|
||||
n_list=200,
|
||||
n_iters=20,
|
||||
use_soar=True,
|
||||
),
|
||||
IVFIndexParam(
|
||||
metric_type=MetricType.COSINE,
|
||||
quantize_type=QuantizeType.FP16,
|
||||
n_list=150,
|
||||
n_iters=15,
|
||||
use_soar=False,
|
||||
),
|
||||
],
|
||||
DataType.VECTOR_FP16: [
|
||||
HnswIndexParam(),
|
||||
FlatIndexParam(),
|
||||
# IVFIndexParam(),
|
||||
],
|
||||
DataType.VECTOR_INT8: [
|
||||
HnswIndexParam(),
|
||||
FlatIndexParam(),
|
||||
# IVFIndexParam(),
|
||||
],
|
||||
DataType.SPARSE_VECTOR_FP32: [
|
||||
HnswIndexParam(),
|
||||
FlatIndexParam(),
|
||||
HnswIndexParam(
|
||||
metric_type=MetricType.IP,
|
||||
m=16,
|
||||
ef_construction=100,
|
||||
quantize_type=QuantizeType.FP16,
|
||||
),
|
||||
],
|
||||
DataType.SPARSE_VECTOR_FP16: [
|
||||
HnswIndexParam(),
|
||||
FlatIndexParam(),
|
||||
HnswIndexParam(
|
||||
metric_type=MetricType.IP,
|
||||
m=16,
|
||||
ef_construction=100,
|
||||
),
|
||||
],
|
||||
}
|
||||
|
||||
VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS = [
|
||||
(data_type, param)
|
||||
for data_type, params in VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP.items()
|
||||
for param in params
|
||||
]
|
||||
|
||||
INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP = {
|
||||
DataType.VECTOR_FP32: [
|
||||
InvertIndexParam(),
|
||||
],
|
||||
DataType.VECTOR_FP16: [
|
||||
InvertIndexParam(),
|
||||
],
|
||||
DataType.VECTOR_INT8: [
|
||||
InvertIndexParam(),
|
||||
],
|
||||
DataType.SPARSE_VECTOR_FP32: [
|
||||
HnswIndexParam(metric_type=MetricType.L2),
|
||||
FlatIndexParam(metric_type=MetricType.COSINE),
|
||||
IVFIndexParam(),
|
||||
InvertIndexParam(),
|
||||
],
|
||||
DataType.SPARSE_VECTOR_FP16: [
|
||||
HnswIndexParam(metric_type=MetricType.L2),
|
||||
FlatIndexParam(metric_type=MetricType.COSINE),
|
||||
IVFIndexParam(),
|
||||
InvertIndexParam(),
|
||||
],
|
||||
}
|
||||
|
||||
INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS = [
|
||||
(data_type, param)
|
||||
for data_type, params in INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP.items()
|
||||
for param in params
|
||||
]
|
||||
|
||||
COLLECTION_NAME_MAX_LENGTH = 64
|
||||
|
||||
COLLECTION_NAME_VALID_LIST = [
|
||||
"col",
|
||||
"C0llECTION",
|
||||
"Collection1",
|
||||
"collection_2",
|
||||
"123collection-",
|
||||
"a" * COLLECTION_NAME_MAX_LENGTH,
|
||||
]
|
||||
|
||||
COLLECTION_NAME_INVALID_LIST = [
|
||||
"l",
|
||||
"1C",
|
||||
"",
|
||||
" ",
|
||||
None,
|
||||
"abcdefghijklmnopqrstuvwxzy123456abcdefghijklmnopqrstuvwxzy1234561",
|
||||
"test/",
|
||||
"!@#$%^&*()test",
|
||||
]
|
||||
|
||||
FIELD_NAME_VALID_LIST = [
|
||||
"1",
|
||||
"12",
|
||||
"col",
|
||||
"ID",
|
||||
"name1",
|
||||
"Weigt_12-",
|
||||
"123age",
|
||||
"name_with_underscores",
|
||||
"123numeric_start",
|
||||
"name-with-dashes",
|
||||
]
|
||||
|
||||
FIELD_NAME_INVALID_LIST = [
|
||||
"",
|
||||
" ",
|
||||
None,
|
||||
"abcdefghijklmnopqrstuvwxzy1234561",
|
||||
"test/",
|
||||
"!@#$%^&*()test",
|
||||
"name@with#special$chars",
|
||||
"name with spaces",
|
||||
]
|
||||
|
||||
FIELD_LIST_MAX_LENGTH = 1024
|
||||
VECTOR_LIST_MAX_LENGTH = 5
|
||||
DENSE_VECTOR_MAX_DIMENSION = 20000
|
||||
SPARSE_VECTOR_MAX_DIMENSION = 4096
|
||||
|
||||
FIELD_VECTOR_LIST_DIMENSION_VALID_LIST = [
|
||||
# field_list_len, vector_list_len, dimension
|
||||
(1, 1, 1),
|
||||
(2, 2, 512),
|
||||
(512, 3, 1024),
|
||||
(1024, 4, 20000),
|
||||
]
|
||||
|
||||
FIELD_VECTOR_LIST_DIMENSION_INVALID_LIST = [
|
||||
# field_list_len, vector_list_len, dimension
|
||||
(1, 1, 0),
|
||||
(1, 1, -1),
|
||||
(1, 1, "1"),
|
||||
(1, 1, 20001),
|
||||
]
|
||||
|
||||
|
||||
INCOMPATIBLE_CONSTRUCTOR_ERROR_MSG = "incompatible constructor arguments"
|
||||
SCHEMA_VALIDATE_ERROR_MSG = "schema validate failed"
|
||||
CREATE_READ_ONLY_ERROR_MSG = "Unable to create collection with read-only mode"
|
||||
INCOMPATIBLE_FUNCTION_ERROR_MSG = "incompatible function arguments"
|
||||
INVALID_PATH_ERROR_MSG = "path validate failed"
|
||||
INDEX_NON_EXISTENT_COLUMN_ERROR_MSG = "not found in schema"
|
||||
ACCESS_DESTROYED_COLLECTION_ERROR_MSG = "is already destroyed"
|
||||
COLLECTION_PATH_NOT_EXIST_ERROR_MSG = "not exist"
|
||||
NOT_SUPPORT_ADD_COLUMN_ERROR_MSG = "Only support basic numeric data type"
|
||||
NOT_EXIST_COLUMN_TO_DROP_ERROR_MSG = "Column not exists"
|
||||
@@ -0,0 +1,126 @@
|
||||
from zvec import (
|
||||
CollectionOption,
|
||||
IndexOption,
|
||||
OptimizeOption,
|
||||
InvertIndexParam,
|
||||
HnswIndexParam,
|
||||
IVFIndexParam,
|
||||
FlatIndexParam,
|
||||
DataType,
|
||||
IndexType,
|
||||
QuantizeType,
|
||||
)
|
||||
|
||||
SUPPORT_SCALAR_DATA_TYPES = [
|
||||
DataType.BOOL,
|
||||
DataType.FLOAT,
|
||||
DataType.DOUBLE,
|
||||
DataType.INT32,
|
||||
DataType.INT64,
|
||||
DataType.UINT32,
|
||||
DataType.UINT64,
|
||||
DataType.STRING,
|
||||
DataType.ARRAY_BOOL,
|
||||
DataType.ARRAY_FLOAT,
|
||||
DataType.ARRAY_DOUBLE,
|
||||
DataType.ARRAY_INT32,
|
||||
DataType.ARRAY_INT64,
|
||||
DataType.ARRAY_UINT32,
|
||||
DataType.ARRAY_UINT64,
|
||||
DataType.ARRAY_STRING,
|
||||
]
|
||||
|
||||
DEFAULT_SCALAR_FIELD_NAME = {
|
||||
DataType.BOOL: "bool_field",
|
||||
DataType.FLOAT: "float_field",
|
||||
DataType.DOUBLE: "double_field",
|
||||
DataType.INT32: "int32_field",
|
||||
DataType.INT64: "int64_field",
|
||||
DataType.UINT32: "uint32_field",
|
||||
DataType.UINT64: "uint64_field",
|
||||
DataType.STRING: "string_field",
|
||||
DataType.ARRAY_BOOL: "array_bool_field",
|
||||
DataType.ARRAY_FLOAT: "array_float_field",
|
||||
DataType.ARRAY_DOUBLE: "array_double_field",
|
||||
DataType.ARRAY_INT32: "array_int32_field",
|
||||
DataType.ARRAY_INT64: "array_int64_field",
|
||||
DataType.ARRAY_UINT32: "array_uint32_field",
|
||||
DataType.ARRAY_UINT64: "array_uint64_field",
|
||||
DataType.ARRAY_STRING: "array_string_field",
|
||||
}
|
||||
|
||||
SUPPORT_SCALAR_INDEX_TYPES = [
|
||||
IndexType.INVERT,
|
||||
]
|
||||
|
||||
SUPPORT_VECTOR_DATA_TYPES = [
|
||||
DataType.VECTOR_FP16,
|
||||
DataType.VECTOR_FP32,
|
||||
DataType.VECTOR_INT8,
|
||||
DataType.SPARSE_VECTOR_FP32,
|
||||
DataType.SPARSE_VECTOR_FP16,
|
||||
]
|
||||
|
||||
SUPPORT_VECTOR_INDEX_TYPES = [
|
||||
IndexType.FLAT,
|
||||
IndexType.HNSW,
|
||||
IndexType.IVF,
|
||||
]
|
||||
|
||||
DEFAULT_VECTOR_FIELD_NAME = {
|
||||
DataType.VECTOR_FP16: "vector_fp16_field",
|
||||
DataType.VECTOR_FP32: "vector_fp32_field",
|
||||
DataType.VECTOR_INT8: "vector_int8_field",
|
||||
DataType.SPARSE_VECTOR_FP32: "sparse_vector_fp32_field",
|
||||
DataType.SPARSE_VECTOR_FP16: "sparse_vector_fp16_field",
|
||||
}
|
||||
|
||||
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],
|
||||
DataType.VECTOR_INT8: [IndexType.FLAT, IndexType.HNSW],
|
||||
DataType.SPARSE_VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW],
|
||||
DataType.SPARSE_VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW],
|
||||
}
|
||||
|
||||
SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP_PARAMS = [
|
||||
(data_type, index_type)
|
||||
for data_type, index_types in SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP.items()
|
||||
for index_type in index_types
|
||||
]
|
||||
|
||||
DEFAULT_INDEX_PARAMS = {
|
||||
IndexType.FLAT: FlatIndexParam(),
|
||||
IndexType.HNSW: HnswIndexParam(),
|
||||
IndexType.IVF: IVFIndexParam(),
|
||||
IndexType.INVERT: InvertIndexParam(),
|
||||
}
|
||||
|
||||
SUPPORT_VECTOR_DATA_TYPE_QUANT_MAP = {
|
||||
DataType.VECTOR_FP32: [QuantizeType.FP16, QuantizeType.INT8, QuantizeType.INT4],
|
||||
DataType.SPARSE_VECTOR_FP32: [QuantizeType.FP16],
|
||||
}
|
||||
|
||||
SUPPORT_ADD_COLUMN_DATA_TYPE = [
|
||||
DataType.INT32,
|
||||
DataType.UINT32,
|
||||
DataType.INT64,
|
||||
DataType.UINT64,
|
||||
DataType.FLOAT,
|
||||
DataType.DOUBLE,
|
||||
]
|
||||
|
||||
NOT_SUPPORT_ADD_COLUMN_DATA_TYPE = [
|
||||
DataType.BOOL,
|
||||
DataType.STRING,
|
||||
DataType.ARRAY_BOOL,
|
||||
DataType.ARRAY_INT32,
|
||||
DataType.ARRAY_INT64,
|
||||
DataType.ARRAY_UINT32,
|
||||
DataType.ARRAY_UINT64,
|
||||
DataType.ARRAY_FLOAT,
|
||||
DataType.ARRAY_DOUBLE,
|
||||
DataType.ARRAY_STRING,
|
||||
]
|
||||
@@ -0,0 +1,429 @@
|
||||
# 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 logging
|
||||
import pytest
|
||||
import threading
|
||||
import numpy as np
|
||||
import zvec
|
||||
|
||||
from zvec import (
|
||||
CollectionOption,
|
||||
InvertIndexParam,
|
||||
HnswIndexParam,
|
||||
Collection,
|
||||
Doc,
|
||||
DataType,
|
||||
FieldSchema,
|
||||
VectorSchema,
|
||||
)
|
||||
|
||||
|
||||
class TestCollectionConcurrency:
|
||||
@pytest.fixture(scope="function")
|
||||
def test_collection(self, tmp_path_factory):
|
||||
"""Fixture to create a test collection"""
|
||||
collection_schema = 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),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
),
|
||||
VectorSchema(
|
||||
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
collection_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
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"
|
||||
|
||||
yield coll
|
||||
|
||||
# Clean up
|
||||
if hasattr(coll, "destroy") and coll is not None:
|
||||
try:
|
||||
coll.destroy()
|
||||
except Exception as e:
|
||||
print(f"Warning: failed to destroy collection: {e}")
|
||||
|
||||
def test_concurrent_read_write(self, test_collection: Collection):
|
||||
results = []
|
||||
|
||||
def insert_docs(thread_id):
|
||||
try:
|
||||
docs = [
|
||||
Doc(
|
||||
id=f"{thread_id}_{i}",
|
||||
fields={
|
||||
"id": int(f"{thread_id}{i}"),
|
||||
"name": f"thread_{thread_id}_doc_{i}",
|
||||
"weight": float(i),
|
||||
},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: float(i), 2: float(i * 2)},
|
||||
},
|
||||
)
|
||||
for i in range(5)
|
||||
]
|
||||
|
||||
result = test_collection.insert(docs)
|
||||
results.append((thread_id, "insert", len(result)))
|
||||
except Exception as e:
|
||||
results.append((thread_id, "insert_exception", str(e)))
|
||||
|
||||
def query_docs(thread_id):
|
||||
try:
|
||||
result = test_collection.query(filter="id > 0", topk=10)
|
||||
results.append((thread_id, "query", len(result)))
|
||||
except Exception as e:
|
||||
results.append((thread_id, "query_exception", str(e)))
|
||||
|
||||
# Create threads for concurrent operations
|
||||
threads = []
|
||||
|
||||
# Start insert threads
|
||||
for i in range(3):
|
||||
thread = threading.Thread(target=insert_docs, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Start query threads
|
||||
for i in range(3):
|
||||
thread = threading.Thread(target=query_docs, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Wait for all threads to complete
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# Analyze results
|
||||
insert_results = [r for r in results if r[1] == "insert"]
|
||||
query_results = [r for r in results if r[1] == "query"]
|
||||
|
||||
logging.info(
|
||||
f"Concurrent read/write results - Inserts: {len(insert_results)}, Queries: {len(query_results)}"
|
||||
)
|
||||
|
||||
# At least some operations should succeed
|
||||
assert len(insert_results) + len(query_results) > 0
|
||||
|
||||
def test_concurrent_query(self, test_collection: Collection):
|
||||
# First insert some data
|
||||
docs = [
|
||||
Doc(
|
||||
id=f"{i}",
|
||||
fields={"id": i, "name": f"test_{i}", "weight": float(i)},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: float(i), 2: float(i * 2)},
|
||||
},
|
||||
)
|
||||
for i in range(20)
|
||||
]
|
||||
|
||||
insert_result = test_collection.insert(docs)
|
||||
assert len(insert_result) == 20
|
||||
|
||||
results = []
|
||||
|
||||
def query_operation(thread_id):
|
||||
"""Perform query operation from a thread"""
|
||||
try:
|
||||
result = test_collection.query(filter=f"id > {thread_id}", topk=5)
|
||||
results.append((thread_id, "query", len(result)))
|
||||
except Exception as e:
|
||||
results.append((thread_id, "query_exception", str(e)))
|
||||
|
||||
# Create multiple threads for concurrent queries
|
||||
threads = []
|
||||
for i in range(5):
|
||||
thread = threading.Thread(target=query_operation, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Wait for all threads to complete
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# Analyze results
|
||||
query_results = [r for r in results if r[1] == "query"]
|
||||
logging.info(f"Concurrent query results - Queries: {len(query_results)}")
|
||||
|
||||
# All query operations should succeed
|
||||
assert len(query_results) == 5
|
||||
|
||||
def test_concurrent_modifications(self, test_collection: Collection):
|
||||
# First insert some data
|
||||
docs = [
|
||||
Doc(
|
||||
id=f"{i}",
|
||||
fields={"id": i, "name": f"test_{i}", "weight": float(i)},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: float(i), 2: float(i * 2)},
|
||||
},
|
||||
)
|
||||
for i in range(10)
|
||||
]
|
||||
|
||||
insert_result = test_collection.insert(docs)
|
||||
assert len(insert_result) == 10
|
||||
|
||||
results = []
|
||||
|
||||
def update_operation(thread_id):
|
||||
"""Perform update operation from a thread"""
|
||||
try:
|
||||
# Each thread updates different documents
|
||||
update_docs = [
|
||||
Doc(
|
||||
id=f"{i}",
|
||||
fields={
|
||||
"id": i,
|
||||
"name": f"updated_by_thread_{thread_id}",
|
||||
"weight": float(i + thread_id),
|
||||
},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: float(i) + 0.5, 2: float(i * 2) + 0.5},
|
||||
},
|
||||
)
|
||||
for i in range(thread_id * 2, thread_id * 2 + 2)
|
||||
]
|
||||
|
||||
result = test_collection.update(update_docs)
|
||||
results.append((thread_id, "update", len(result)))
|
||||
except Exception as e:
|
||||
results.append((thread_id, "update_exception", str(e)))
|
||||
|
||||
def delete_operation(thread_id):
|
||||
"""Perform delete operation from a thread"""
|
||||
try:
|
||||
# Each thread deletes different documents
|
||||
delete_ids = [f"{thread_id * 2 + 2}", f"{thread_id * 2 + 3}"]
|
||||
result = test_collection.delete(delete_ids)
|
||||
results.append((thread_id, "delete", len(result)))
|
||||
except Exception as e:
|
||||
results.append((thread_id, "delete_exception", str(e)))
|
||||
|
||||
# Create threads for concurrent operations
|
||||
threads = []
|
||||
|
||||
# Start update threads
|
||||
for i in range(3):
|
||||
thread = threading.Thread(target=update_operation, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Start delete threads
|
||||
for i in range(2):
|
||||
thread = threading.Thread(target=delete_operation, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Wait for all threads to complete
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# Analyze results
|
||||
update_results = [r for r in results if r[1] == "update"]
|
||||
delete_results = [r for r in results if r[1] == "delete"]
|
||||
|
||||
logging.info(
|
||||
f"Concurrent modification results - Updates: {len(update_results)}, Deletes: {len(delete_results)}"
|
||||
)
|
||||
|
||||
# At least some operations should succeed
|
||||
assert len(update_results) + len(delete_results) > 0
|
||||
|
||||
def test_read_write_locking(self, test_collection: Collection):
|
||||
# Perform operations that should be thread-safe
|
||||
docs = [
|
||||
Doc(
|
||||
id=f"{i}",
|
||||
fields={"id": i, "name": f"test_{i}", "weight": float(i)},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: float(i), 2: float(i * 2)},
|
||||
},
|
||||
)
|
||||
for i in range(5)
|
||||
]
|
||||
|
||||
# Insert data
|
||||
insert_result = test_collection.insert(docs)
|
||||
assert len(insert_result) == 5
|
||||
|
||||
# Concurrent operations should not cause data corruption
|
||||
results = []
|
||||
|
||||
def mixed_operation(thread_id):
|
||||
"""Perform mixed operations from a thread"""
|
||||
try:
|
||||
# Mix of read and write operations
|
||||
if thread_id % 2 == 0:
|
||||
# Read operation
|
||||
result = test_collection.fetch([f"{thread_id % 5}"])
|
||||
results.append((thread_id, "read", len(result)))
|
||||
else:
|
||||
# Write operation
|
||||
doc = Doc(
|
||||
id=f"{thread_id % 5}",
|
||||
fields={
|
||||
"id": thread_id % 5,
|
||||
"name": f"mixed_op_{thread_id}",
|
||||
"weight": float(thread_id),
|
||||
},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: float(thread_id), 2: float(thread_id * 2)},
|
||||
},
|
||||
)
|
||||
result = test_collection.upsert(doc)
|
||||
results.append((thread_id, "write", len(result)))
|
||||
except Exception as e:
|
||||
results.append((thread_id, "exception", str(e)))
|
||||
|
||||
# Create multiple threads
|
||||
threads = []
|
||||
for i in range(10):
|
||||
thread = threading.Thread(target=mixed_operation, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Wait for all threads to complete
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# Verify that the collection is still in a consistent state
|
||||
final_result = test_collection.query()
|
||||
assert len(final_result) >= 0 # Should not crash or return corrupted data
|
||||
|
||||
def test_race_condition_detection(self, test_collection: Collection):
|
||||
# Insert initial data
|
||||
docs = [
|
||||
Doc(
|
||||
id=f"{i}",
|
||||
fields={"id": i, "name": f"initial_{i}", "weight": float(i)},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: float(i), 2: float(i * 2)},
|
||||
},
|
||||
)
|
||||
for i in range(10)
|
||||
]
|
||||
|
||||
insert_result = test_collection.insert(docs)
|
||||
assert len(insert_result) == 10
|
||||
|
||||
# Perform many rapid concurrent operations
|
||||
operation_count = 100
|
||||
results = []
|
||||
|
||||
def rapid_operation(op_id):
|
||||
"""Perform rapid operations"""
|
||||
try:
|
||||
# Alternate between different types of operations
|
||||
if op_id % 4 == 0:
|
||||
# Insert
|
||||
doc = Doc(
|
||||
id=f"rapid_{op_id}",
|
||||
fields={
|
||||
"id": op_id,
|
||||
"name": f"rapid_{op_id}",
|
||||
"weight": float(op_id),
|
||||
},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: float(op_id), 2: float(op_id * 2)},
|
||||
},
|
||||
)
|
||||
result = test_collection.insert(doc)
|
||||
results.append(("insert", len(result)))
|
||||
elif op_id % 4 == 1:
|
||||
# Update
|
||||
doc = Doc(
|
||||
id=f"{op_id % 10}",
|
||||
fields={
|
||||
"id": op_id % 10,
|
||||
"name": f"rapid_update_{op_id}",
|
||||
"weight": float(op_id),
|
||||
},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: float(op_id), 2: float(op_id * 2)},
|
||||
},
|
||||
)
|
||||
result = test_collection.update(doc)
|
||||
results.append(("update", len(result)))
|
||||
elif op_id % 4 == 2:
|
||||
# Query
|
||||
result = test_collection.query(filter=f"id > {op_id % 5}", topk=3)
|
||||
results.append(("query", len(result)))
|
||||
else:
|
||||
# Fetch
|
||||
result = test_collection.fetch([f"{op_id % 10}"])
|
||||
results.append(("fetch", len(result)))
|
||||
except Exception as e:
|
||||
results.append(("exception", str(e)))
|
||||
|
||||
# Create many threads for rapid concurrent operations
|
||||
threads = []
|
||||
for i in range(operation_count):
|
||||
thread = threading.Thread(target=rapid_operation, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Wait for all threads to complete
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# Verify collection is still functional
|
||||
final_query = test_collection.query()
|
||||
assert len(final_query) >= 0 # Should not be corrupted
|
||||
|
||||
logging.info(
|
||||
f"Rapid concurrent operations completed - Total operations: {len(results)}"
|
||||
)
|
||||
@@ -0,0 +1,791 @@
|
||||
# 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 threading
|
||||
import os
|
||||
|
||||
from distance_helper import *
|
||||
from fixture_helper import *
|
||||
from doc_helper import *
|
||||
from params_helper import *
|
||||
|
||||
|
||||
def check_collection_info(
|
||||
coll: Collection, schema: CollectionSchema, option: CollectionOption, path: str
|
||||
):
|
||||
assert coll is not None, "Failed to create and open collection"
|
||||
assert coll.path == path
|
||||
assert coll.schema.name == schema.name
|
||||
assert list(coll.schema.fields) == list(schema.fields)
|
||||
assert list(coll.schema.vectors) == list(schema.vectors)
|
||||
assert coll.option.read_only == option.read_only
|
||||
assert coll.option.enable_mmap == option.enable_mmap
|
||||
|
||||
|
||||
def check_collection_basic(coll: Collection, optimize: bool = False):
|
||||
schema = coll.schema
|
||||
|
||||
docs = [generate_doc(i, schema) for i in range(10)]
|
||||
|
||||
results = coll.insert(docs=docs)
|
||||
assert len(results) == len(docs)
|
||||
for result in results:
|
||||
assert result.ok()
|
||||
|
||||
assert coll.stats.doc_count == len(docs)
|
||||
|
||||
def check_fetch_query():
|
||||
results = coll.fetch([str(i) for i in range(len(docs))])
|
||||
assert len(results) == len(docs)
|
||||
for i in range(len(docs)):
|
||||
assert str(i) in results
|
||||
|
||||
results = coll.query()
|
||||
assert len(results) == len(docs)
|
||||
|
||||
check_fetch_query()
|
||||
|
||||
if optimize:
|
||||
coll.optimize()
|
||||
check_fetch_query()
|
||||
|
||||
|
||||
def check_collection_full(coll: Collection):
|
||||
test_doc = generate_doc(1, coll.schema)
|
||||
|
||||
insert_result = coll.insert(test_doc)
|
||||
assert insert_result.ok()
|
||||
|
||||
stats = coll.stats
|
||||
assert stats.doc_count == 1
|
||||
|
||||
fetched_docs = coll.fetch(ids=["1"])
|
||||
assert len(fetched_docs) == 1
|
||||
assert "1" in fetched_docs
|
||||
assert fetched_docs["1"] is not None
|
||||
assert is_doc_equal(fetched_docs["1"], test_doc, coll.schema)
|
||||
|
||||
query_result = coll.query()
|
||||
assert len(query_result) == 1
|
||||
|
||||
updated_doc = Doc(
|
||||
id="1",
|
||||
fields={"int32_field": 1},
|
||||
vectors={"vector_fp32_field": [0.2] * 128},
|
||||
)
|
||||
update_result = coll.update(updated_doc)
|
||||
assert update_result.ok()
|
||||
|
||||
upserted_doc = generate_doc(1, coll.schema)
|
||||
upsert_result = coll.upsert(upserted_doc)
|
||||
assert upsert_result.ok()
|
||||
|
||||
# 8. Delete document
|
||||
delete_result = coll.delete("1")
|
||||
assert delete_result.ok()
|
||||
|
||||
# Verify document was deleted
|
||||
stats = coll.stats
|
||||
assert stats.doc_count == 0
|
||||
|
||||
|
||||
valid_collection_options = [
|
||||
# (read_only, enable_mmap)
|
||||
(False, True),
|
||||
(False, False),
|
||||
]
|
||||
invalid_collection_options = [
|
||||
# (read_only, enable_mmap)
|
||||
(True, True),
|
||||
(True, False),
|
||||
]
|
||||
duplicate_names_test = [
|
||||
("field1", "field1", "vector1", "vector2"),
|
||||
("field1", "field2", "vector1", "vector1"),
|
||||
(
|
||||
"shared_name1",
|
||||
"shared_name2",
|
||||
"shared_name1",
|
||||
"shared_name2",
|
||||
),
|
||||
]
|
||||
long_names = [
|
||||
"a" * 100, # 100 characters
|
||||
"b" * 200, # 200 characters
|
||||
]
|
||||
|
||||
valid_path_list = [
|
||||
"/tmp/nonexistent/directory/test_collection",
|
||||
"test/collection/with/slashes",
|
||||
"test/collection/with/slashes/哈哈",
|
||||
]
|
||||
invalid_path_list = [
|
||||
"invalid\0path",
|
||||
"",
|
||||
]
|
||||
|
||||
|
||||
class TestCreateAndOpen:
|
||||
@pytest.mark.parametrize("collection_name", COLLECTION_NAME_VALID_LIST)
|
||||
def test_valid_collection_name(
|
||||
self,
|
||||
collection_temp_dir,
|
||||
collection_name,
|
||||
collection_option,
|
||||
sample_field_list,
|
||||
sample_vector_list,
|
||||
):
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name=collection_name,
|
||||
fields=sample_field_list,
|
||||
vectors=sample_vector_list,
|
||||
)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
check_collection_info(
|
||||
coll, collection_schema, collection_option, collection_temp_dir
|
||||
)
|
||||
check_collection_basic(coll)
|
||||
|
||||
coll.destroy()
|
||||
|
||||
@pytest.mark.parametrize("collection_name", COLLECTION_NAME_INVALID_LIST)
|
||||
def test_invalid_collection_name(
|
||||
self,
|
||||
collection_temp_dir,
|
||||
collection_name,
|
||||
collection_option,
|
||||
sample_field_list,
|
||||
sample_vector_list,
|
||||
):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name=collection_name,
|
||||
fields=sample_field_list,
|
||||
vectors=sample_vector_list,
|
||||
)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
|
||||
@pytest.mark.parametrize("name_prefix", FIELD_NAME_VALID_LIST)
|
||||
def test_valid_field_vector_name(
|
||||
self,
|
||||
collection_temp_dir,
|
||||
collection_option,
|
||||
name_prefix,
|
||||
sample_field_list,
|
||||
sample_vector_list,
|
||||
):
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="test_collection",
|
||||
fields=sample_field_list,
|
||||
vectors=sample_vector_list,
|
||||
)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
check_collection_info(
|
||||
coll, collection_schema, collection_option, collection_temp_dir
|
||||
)
|
||||
check_collection_basic(coll)
|
||||
|
||||
coll.destroy()
|
||||
|
||||
@pytest.mark.parametrize("field_name", FIELD_NAME_INVALID_LIST)
|
||||
def test_invalid_field_name(
|
||||
self, collection_temp_dir, collection_option, field_name
|
||||
):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
field_list = [FieldSchema(field_name, DataType.STRING)]
|
||||
vector_list = [
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
]
|
||||
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="collection_name", fields=field_list, vectors=vector_list
|
||||
)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
|
||||
@pytest.mark.parametrize("vector_name", FIELD_NAME_INVALID_LIST)
|
||||
def test_invalid_vector_name(
|
||||
self, collection_temp_dir, collection_option, vector_name
|
||||
):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
field_list = [
|
||||
FieldSchema(
|
||||
"id",
|
||||
DataType.INT64,
|
||||
nullable=False,
|
||||
index_param=InvertIndexParam(enable_range_optimization=True),
|
||||
)
|
||||
]
|
||||
vector_list = [
|
||||
VectorSchema(vector_name, DataType.VECTOR_FP32, dimension=128)
|
||||
]
|
||||
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="collection_name", fields=field_list, vectors=vector_list
|
||||
)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field_list_len,vector_list_len,dimension",
|
||||
FIELD_VECTOR_LIST_DIMENSION_VALID_LIST,
|
||||
)
|
||||
def test_valid_field_vector_size_dimension(
|
||||
self,
|
||||
collection_temp_dir,
|
||||
collection_option,
|
||||
field_list_len,
|
||||
vector_list_len,
|
||||
dimension,
|
||||
):
|
||||
field_list = []
|
||||
vector_list = []
|
||||
for i in range(0, field_list_len):
|
||||
field_list.append(
|
||||
FieldSchema("id_" + str(i), DataType.INT64, nullable=True)
|
||||
)
|
||||
|
||||
for i in range(0, vector_list_len):
|
||||
vector_list.append(
|
||||
VectorSchema(
|
||||
"dense_vector_" + str(i),
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=dimension,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
)
|
||||
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="test_dense_vector_list", fields=field_list, vectors=vector_list
|
||||
)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
check_collection_info(
|
||||
coll, collection_schema, collection_option, collection_temp_dir
|
||||
)
|
||||
check_collection_basic(coll)
|
||||
|
||||
coll.destroy()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field_list_len,vector_list_len,dimension",
|
||||
FIELD_VECTOR_LIST_DIMENSION_INVALID_LIST,
|
||||
)
|
||||
def test_invalid_field_vector_size_dimension(
|
||||
self,
|
||||
collection_temp_dir,
|
||||
collection_option,
|
||||
vector_list_len,
|
||||
field_list_len,
|
||||
dimension,
|
||||
):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
field_list = []
|
||||
vector_list = []
|
||||
for i in range(0, field_list_len):
|
||||
field_list.append(
|
||||
FieldSchema(
|
||||
"id_" + str(i),
|
||||
DataType.INT64,
|
||||
nullable=False,
|
||||
)
|
||||
)
|
||||
|
||||
for i in range(0, vector_list_len):
|
||||
vector_list.append(
|
||||
VectorSchema(
|
||||
"dense_vector_" + str(i),
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=dimension,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
)
|
||||
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="test_dense_vector_list", fields=field_list, vectors=vector_list
|
||||
)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
|
||||
def test_valid_single_vector_field_construction(
|
||||
self, collection_temp_dir, collection_option
|
||||
):
|
||||
field = FieldSchema(
|
||||
"id",
|
||||
DataType.INT64,
|
||||
nullable=True,
|
||||
index_param=InvertIndexParam(enable_range_optimization=True),
|
||||
)
|
||||
|
||||
vector = VectorSchema(
|
||||
"dense_vector",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="test_single_dense_vector_non_list",
|
||||
fields=field,
|
||||
vectors=vector, # Non-list form
|
||||
)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
check_collection_info(
|
||||
coll, collection_schema, collection_option, collection_temp_dir
|
||||
)
|
||||
check_collection_basic(coll)
|
||||
coll.destroy()
|
||||
|
||||
def test_collection_concurrent_create(
|
||||
self, collection_temp_dir, basic_schema, collection_option
|
||||
):
|
||||
results = []
|
||||
errors = []
|
||||
lock = threading.Lock()
|
||||
|
||||
# Function to be executed by each thread
|
||||
def create_collection_thread(thread_id):
|
||||
try:
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=basic_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
with lock:
|
||||
results.append((thread_id, coll))
|
||||
except Exception as e:
|
||||
with lock:
|
||||
errors.append((thread_id, str(e)))
|
||||
|
||||
threads = []
|
||||
for i in range(5):
|
||||
thread = threading.Thread(target=create_collection_thread, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
assert len(results) == 1, (
|
||||
f"Expected exactly one successful creation, but got {len(results)}"
|
||||
)
|
||||
assert len(errors) == 4, (
|
||||
f"Expected exactly four failures, but got {len(errors)}"
|
||||
)
|
||||
|
||||
successful_thread_id, successful_collection = results[0]
|
||||
assert successful_collection is not None, (
|
||||
"Successful creation should return a valid collection"
|
||||
)
|
||||
assert successful_collection.path == collection_temp_dir, (
|
||||
"Collection path mismatch"
|
||||
)
|
||||
|
||||
def test_create_open_loop(
|
||||
self, collection_temp_dir, collection_option, full_schema
|
||||
):
|
||||
for cycle in range(10):
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=full_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
assert coll is not None, (
|
||||
f"Failed to create and open collection in cycle {cycle}"
|
||||
)
|
||||
assert coll.path == collection_temp_dir, (
|
||||
f"Collection path mismatch in cycle {cycle}"
|
||||
)
|
||||
|
||||
del coll
|
||||
|
||||
reopened_coll = zvec.open(
|
||||
path=collection_temp_dir, option=collection_option
|
||||
)
|
||||
assert reopened_coll is not None, (
|
||||
f"Failed to reopen collection in cycle {cycle}"
|
||||
)
|
||||
assert reopened_coll.path == collection_temp_dir, (
|
||||
f"Reopened collection path mismatch in cycle {cycle}"
|
||||
)
|
||||
|
||||
check_collection_full(reopened_coll)
|
||||
|
||||
reopened_coll.destroy()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_type, index_param", VALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS
|
||||
)
|
||||
def test_valid_vector_index_params(
|
||||
self,
|
||||
data_type,
|
||||
index_param,
|
||||
single_vector_schema_with_index_param,
|
||||
collection_temp_dir,
|
||||
collection_option,
|
||||
):
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=single_vector_schema_with_index_param,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
check_collection_info(
|
||||
coll,
|
||||
single_vector_schema_with_index_param,
|
||||
collection_option,
|
||||
collection_temp_dir,
|
||||
)
|
||||
|
||||
check_collection_basic(coll, True)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_type, index_param", INVALID_VECTOR_DATA_TYPE_INDEX_PARAM_MAP_PARAMS
|
||||
)
|
||||
def test_invalid_vector_index_params(
|
||||
self,
|
||||
data_type,
|
||||
index_param,
|
||||
single_vector_schema_with_index_param,
|
||||
collection_temp_dir,
|
||||
collection_option,
|
||||
):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=single_vector_schema_with_index_param,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
|
||||
def test_open_concurrent_same_path(self, tmp_path_factory, collection_option):
|
||||
"""Test concurrent opening of the same collection path.
|
||||
|
||||
- Multi-threading concurrency: 5 threads simultaneously open the same collection
|
||||
- Result verification: Verify that only one can open successfully, others must fail
|
||||
"""
|
||||
# Create a temporary directory and path for the collection
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "concurrent_open_test_collection"
|
||||
|
||||
# First, create a collection that we'll try to open concurrently
|
||||
field_list = [
|
||||
FieldSchema(
|
||||
"id",
|
||||
DataType.INT64,
|
||||
nullable=False,
|
||||
index_param=InvertIndexParam(enable_range_optimization=True),
|
||||
),
|
||||
FieldSchema(
|
||||
"name", DataType.STRING, nullable=False, index_param=InvertIndexParam()
|
||||
),
|
||||
]
|
||||
|
||||
vector_list = [
|
||||
VectorSchema(
|
||||
"dense_vector",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
]
|
||||
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="concurrent_open_test_collection",
|
||||
fields=field_list,
|
||||
vectors=vector_list,
|
||||
)
|
||||
|
||||
# Create the collection first
|
||||
coll = zvec.create_and_open(
|
||||
path=str(collection_path),
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
# Close the collection so we can test opening it
|
||||
if hasattr(coll, "close") and coll is not None:
|
||||
coll.close()
|
||||
|
||||
# Shared variables to collect results from threads
|
||||
results = []
|
||||
errors = []
|
||||
|
||||
# Lock for thread-safe operations
|
||||
lock = threading.Lock()
|
||||
# Clean up the created collection reference
|
||||
del coll
|
||||
|
||||
# Function to be executed by each thread
|
||||
def open_collection_thread(thread_id):
|
||||
try:
|
||||
reopened_coll = zvec.open(
|
||||
path=str(collection_path), option=collection_option
|
||||
)
|
||||
with lock:
|
||||
results.append((thread_id, reopened_coll))
|
||||
# Clean up the collection if opened successfully
|
||||
if hasattr(reopened_coll, "close") and reopened_coll is not None:
|
||||
reopened_coll.close()
|
||||
except Exception as e:
|
||||
with lock:
|
||||
errors.append((thread_id, str(e)))
|
||||
|
||||
# Create and start 5 threads
|
||||
threads = []
|
||||
for i in range(5):
|
||||
thread = threading.Thread(target=open_collection_thread, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Wait for all threads to complete
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# Verify results:
|
||||
# 1. Only one open should succeed (exactly one collection in results)
|
||||
# 2. Others should fail (4 errors in errors)
|
||||
assert len(results) == 1, (
|
||||
f"Expected exactly one successful open, but got {len(results)}"
|
||||
)
|
||||
assert len(errors) == 4, (
|
||||
f"Expected exactly four failures, but got {len(errors)}"
|
||||
)
|
||||
|
||||
# Additional verification: check that the successful open has a valid collection
|
||||
successful_thread_id, successful_collection = results[0]
|
||||
assert successful_collection is not None, (
|
||||
"Successful open should return a valid collection"
|
||||
)
|
||||
assert successful_collection.path == str(collection_path), (
|
||||
"Collection path mismatch"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("read_only,enable_mmap", valid_collection_options)
|
||||
def test_valid_option(
|
||||
self, collection_temp_dir, basic_schema, read_only, enable_mmap
|
||||
):
|
||||
option = CollectionOption(read_only=read_only, enable_mmap=enable_mmap)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=basic_schema,
|
||||
option=option,
|
||||
)
|
||||
|
||||
check_collection_info(coll, basic_schema, option, collection_temp_dir)
|
||||
check_collection_basic(coll)
|
||||
|
||||
coll.destroy()
|
||||
|
||||
def test_valid_none_option(self, collection_temp_dir, basic_schema):
|
||||
zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=basic_schema,
|
||||
option=None,
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("read_only,enable_mmap", invalid_collection_options)
|
||||
def test_invalid_option(
|
||||
self, collection_temp_dir, basic_schema, read_only, enable_mmap
|
||||
):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=basic_schema,
|
||||
option=CollectionOption(read_only=read_only, enable_mmap=enable_mmap),
|
||||
)
|
||||
|
||||
assert CREATE_READ_ONLY_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field_name1,field_name2,vector_name1,vector_name2",
|
||||
duplicate_names_test,
|
||||
)
|
||||
def test_duplicate_field_names(
|
||||
self,
|
||||
collection_temp_dir,
|
||||
collection_option,
|
||||
field_name1,
|
||||
field_name2,
|
||||
vector_name1,
|
||||
vector_name2,
|
||||
):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="test_collection",
|
||||
fields=[
|
||||
FieldSchema(
|
||||
field_name1,
|
||||
DataType.INT64,
|
||||
nullable=False,
|
||||
index_param=InvertIndexParam(enable_range_optimization=True),
|
||||
),
|
||||
FieldSchema(
|
||||
field_name2,
|
||||
DataType.INT64,
|
||||
nullable=False,
|
||||
index_param=InvertIndexParam(enable_range_optimization=True),
|
||||
),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
vector_name1,
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
),
|
||||
VectorSchema(
|
||||
vector_name2,
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
|
||||
@pytest.mark.parametrize("long_name", long_names)
|
||||
def test_invalid_long_field_names(
|
||||
self, collection_option, collection_temp_dir, long_name
|
||||
):
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name=long_name,
|
||||
fields=[
|
||||
FieldSchema(
|
||||
long_name + "_field",
|
||||
DataType.INT64,
|
||||
nullable=False,
|
||||
index_param=InvertIndexParam(enable_range_optimization=True),
|
||||
),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
long_name + "_vector",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
|
||||
def test_invalid_empty_fields_and_vectors(
|
||||
self, collection_temp_dir, collection_option
|
||||
):
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="test_collection",
|
||||
fields=[], # Empty fields
|
||||
vectors=[], # Empty vectors
|
||||
)
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
coll = zvec.create_and_open(
|
||||
path=collection_temp_dir,
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert SCHEMA_VALIDATE_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
|
||||
@pytest.mark.parametrize("valid_path", valid_path_list)
|
||||
def test_valid_path(self, basic_schema, collection_option, valid_path):
|
||||
if os.path.exists(valid_path):
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(valid_path)
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=valid_path, schema=basic_schema, option=collection_option
|
||||
)
|
||||
|
||||
check_collection_info(coll, basic_schema, collection_option, valid_path)
|
||||
|
||||
coll.destroy()
|
||||
|
||||
@pytest.mark.parametrize("invalid_path", invalid_path_list)
|
||||
def test_invalid_path(self, basic_schema, collection_option, invalid_path):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
coll = zvec.create_and_open(
|
||||
path=invalid_path, schema=basic_schema, option=collection_option
|
||||
)
|
||||
|
||||
assert INVALID_PATH_ERROR_MSG in str(exc_info.value), str(exc_info.value)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,328 @@
|
||||
# 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 logging
|
||||
import pytest
|
||||
import numpy as np
|
||||
import zvec
|
||||
|
||||
from zvec import (
|
||||
CollectionOption,
|
||||
InvertIndexParam,
|
||||
HnswIndexParam,
|
||||
DataType,
|
||||
Collection,
|
||||
Doc,
|
||||
FieldSchema,
|
||||
Query,
|
||||
VectorSchema,
|
||||
)
|
||||
|
||||
|
||||
class TestCollectionExceptionHandling:
|
||||
@pytest.fixture(scope="function")
|
||||
def test_collection(self, tmp_path_factory):
|
||||
"""Fixture to create a test collection"""
|
||||
collection_schema = 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),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
),
|
||||
VectorSchema(
|
||||
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
collection_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
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"
|
||||
|
||||
yield coll
|
||||
|
||||
# Clean up
|
||||
if hasattr(coll, "destroy") and coll is not None:
|
||||
try:
|
||||
coll.destroy()
|
||||
except Exception as e:
|
||||
print(f"Warning: failed to destroy collection: {e}")
|
||||
|
||||
def test_create_and_open_missing_path(self, tmp_path_factory):
|
||||
collection_schema = 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(),
|
||||
),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
collection_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
coll = zvec.create_and_open(
|
||||
schema=collection_schema, option=collection_option
|
||||
)
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing path parameter"
|
||||
)
|
||||
|
||||
def test_create_and_open_missing_schema(self, tmp_path_factory):
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "test_collection"
|
||||
|
||||
collection_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
coll = zvec.create_and_open(
|
||||
path=str(collection_path), option=collection_option
|
||||
)
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing schema parameter"
|
||||
)
|
||||
|
||||
def test_open_missing_path(self):
|
||||
collection_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
coll = zvec.open(option=collection_option)
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing path parameter"
|
||||
)
|
||||
|
||||
def test_insert_missing_docs(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
result = test_collection.insert()
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing docs parameter"
|
||||
)
|
||||
|
||||
def test_update_missing_docs(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
result = test_collection.update()
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing docs parameter"
|
||||
)
|
||||
|
||||
def test_upsert_missing_docs(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
result = test_collection.upsert()
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing docs parameter"
|
||||
)
|
||||
|
||||
def test_delete_missing_ids(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
result = test_collection.delete()
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing ids parameter"
|
||||
)
|
||||
|
||||
def test_fetch_missing_ids(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
result = test_collection.fetch()
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing ids parameter"
|
||||
)
|
||||
|
||||
def test_query_missing_query_field_name(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
result = test_collection.query([Query()])
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing Query field_name parameter"
|
||||
)
|
||||
|
||||
def test_add_column_missing_field_schema(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
test_collection.add_column()
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing field_schema parameter"
|
||||
)
|
||||
|
||||
def test_alter_column_missing_old_name(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
test_collection.alter_column(new_name="new_name")
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing old_name parameter"
|
||||
)
|
||||
|
||||
def test_alter_column_missing_new_name(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
test_collection.alter_column(old_name="old_name")
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing new_name parameter"
|
||||
)
|
||||
|
||||
def test_drop_column_missing_field_name(self, test_collection: Collection):
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
test_collection.drop_column()
|
||||
assert exc_info.value is not None, (
|
||||
"Expected exception for missing field_name parameter"
|
||||
)
|
||||
|
||||
def test_invalid_parameter_types(self, test_collection: Collection):
|
||||
# This test depends on specific implementation details
|
||||
# Generally, we would expect TypeErrors or similar exceptions
|
||||
pass
|
||||
|
||||
def test_missing_required_parameters(self, test_collection: Collection):
|
||||
# This test depends on specific implementation details
|
||||
# Generally, we would expect TypeErrors or similar exceptions
|
||||
pass
|
||||
|
||||
def test_empty_collection_operations(self, tmp_path_factory):
|
||||
collection_schema = zvec.CollectionSchema(
|
||||
name="empty_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(),
|
||||
),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
collection_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "empty_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"
|
||||
|
||||
# Test fetch on empty collection
|
||||
result = coll.fetch(["1"])
|
||||
assert len(result) >= 0 # May be empty or have special handling
|
||||
|
||||
# Test query on empty collection
|
||||
result = coll.query()
|
||||
assert len(result) == 0
|
||||
|
||||
# Test update on empty collection
|
||||
doc = Doc(
|
||||
id="1",
|
||||
fields={"id": 1, "name": "test"},
|
||||
vectors={"dense": np.random.random(128).tolist()},
|
||||
)
|
||||
|
||||
result = coll.update(doc)
|
||||
# Should handle gracefully, possibly with NOT_FOUND status
|
||||
|
||||
# Clean up
|
||||
if hasattr(coll, "destroy") and coll is not None:
|
||||
try:
|
||||
coll.destroy()
|
||||
except Exception as e:
|
||||
print(f"Warning: failed to destroy collection: {e}")
|
||||
|
||||
def test_resource_management(self, test_collection: Collection):
|
||||
doc = Doc(
|
||||
id="1",
|
||||
fields={"id": 1, "name": "test", "weight": 80.5},
|
||||
vectors={
|
||||
"dense": np.random.random(128).tolist(),
|
||||
"sparse": {1: 1.0, 2: 2.0},
|
||||
},
|
||||
)
|
||||
|
||||
# Insert
|
||||
result = test_collection.insert(doc)
|
||||
assert result.ok()
|
||||
|
||||
# Fetch
|
||||
result = test_collection.fetch(["1"])
|
||||
assert len(result) == 1
|
||||
|
||||
# Query
|
||||
result = test_collection.query()
|
||||
assert len(result) >= 0
|
||||
|
||||
# Update
|
||||
result = test_collection.update(doc)
|
||||
assert result.ok()
|
||||
|
||||
# Delete
|
||||
result = test_collection.delete("1")
|
||||
assert result.ok()
|
||||
|
||||
def test_exception_resource_cleanup(self, test_collection: Collection):
|
||||
# This test would need to simulate exception conditions
|
||||
# which is difficult without specific failure injection points
|
||||
pass
|
||||
@@ -0,0 +1,967 @@
|
||||
# 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 threading
|
||||
import numpy as np
|
||||
|
||||
from fixture_helper import *
|
||||
|
||||
COLLECTION_OPTION_TEST_CASES_VALID = [
|
||||
# (read_only, enable_mmap, description)
|
||||
(False, True, "Read-write with mmap enabled"),
|
||||
(False, False, "Read-write with mmap disabled"),
|
||||
(True, True, "Read-only with mmap enabled"),
|
||||
(True, False, "Read-only with mmap disabled"),
|
||||
]
|
||||
|
||||
# Test data for invalid paths
|
||||
INVALID_PATH_LIST = [
|
||||
"/nonexistent/directory/test_collection",
|
||||
"invalid:path",
|
||||
"", # Empty path
|
||||
]
|
||||
|
||||
|
||||
@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=False, index_param=InvertIndexParam()
|
||||
),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
),
|
||||
VectorSchema(
|
||||
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def single_doc():
|
||||
id = 0
|
||||
return Doc(
|
||||
id=f"{id}",
|
||||
fields={"id": id, "name": "test"},
|
||||
vectors={
|
||||
"dense": [id + 0.1] * 128,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def test_collection(
|
||||
tmp_path_factory, collection_schema, collection_option
|
||||
) -> Generator[Any, Any, Collection]:
|
||||
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}")
|
||||
|
||||
|
||||
class TestCollectionOpen:
|
||||
def test_open_basic_functionality(
|
||||
self, tmp_path_factory, collection_schema, collection_option
|
||||
):
|
||||
import sys
|
||||
import time
|
||||
import os
|
||||
|
||||
# Create unique temp directory
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "test_collection"
|
||||
|
||||
# Ensure the path exists
|
||||
collection_path_str = str(collection_path)
|
||||
print(f"DEBUG: Collection path: {collection_path_str}")
|
||||
print(f"DEBUG: Temp directory exists: {temp_dir.exists()}")
|
||||
|
||||
# Create and open collection first
|
||||
created_coll = zvec.create_and_open(
|
||||
path=collection_path_str, schema=collection_schema, option=collection_option
|
||||
)
|
||||
|
||||
assert created_coll is not None, (
|
||||
f"Failed to create collection, returned None instead of valid Collection object. Path: {collection_path_str}"
|
||||
)
|
||||
assert created_coll.path == collection_path_str, (
|
||||
f"Collection path mismatch. Expected: {collection_path_str}, Actual: {created_coll.path}"
|
||||
)
|
||||
assert created_coll.schema.name == "test_collection", (
|
||||
f"Collection schema name mismatch. Expected: test_collection, Actual: {created_coll.schema.name}"
|
||||
)
|
||||
|
||||
# Insert multiple documents to verify persistence
|
||||
docs = []
|
||||
for i in range(3):
|
||||
doc = Doc(
|
||||
id=f"{i}",
|
||||
fields={"id": i, "name": f"test_{i}", "weight": float(i * 10)},
|
||||
vectors={
|
||||
"dense": [float(j + i) for j in range(128)],
|
||||
"sparse": {j: float(j + i) for j in range(5)},
|
||||
},
|
||||
)
|
||||
docs.append(doc)
|
||||
|
||||
result = created_coll.insert(docs)
|
||||
assert len(result) == 3, f"Expected 3 insertion results, but got {len(result)}"
|
||||
for i, res in enumerate(result):
|
||||
assert res.ok(), (
|
||||
f"Insertion result {i} is not OK. Status code: {res.code()}, Message: {res.message()}"
|
||||
)
|
||||
|
||||
# Verify documents were inserted using fetch interface
|
||||
fetched_docs_after_insert = created_coll.fetch(["0", "1", "2"])
|
||||
assert len(fetched_docs_after_insert) == 3, (
|
||||
f"Expected 3 fetched documents after insertion, but got {len(fetched_docs_after_insert)}"
|
||||
)
|
||||
assert "0" in fetched_docs_after_insert, (
|
||||
"Document with ID '0' not found in fetched results after insertion"
|
||||
)
|
||||
assert "1" in fetched_docs_after_insert, (
|
||||
"Document with ID '1' not found in fetched results after insertion"
|
||||
)
|
||||
assert "2" in fetched_docs_after_insert, (
|
||||
"Document with ID '2' not found in fetched results after insertion"
|
||||
)
|
||||
|
||||
# Verify fetched document content after insertion
|
||||
for i in range(3):
|
||||
doc = fetched_docs_after_insert[f"{i}"]
|
||||
assert doc is not None, (
|
||||
f"Fetched document with ID '{i}' is None after insertion"
|
||||
)
|
||||
assert doc.id == f"{i}", (
|
||||
f"Document ID mismatch for document '{i}' after insertion. Expected: {i}, Actual: {doc.id}"
|
||||
)
|
||||
assert doc.field("id") == i, (
|
||||
f"Document id field mismatch for document '{i}' after insertion. Expected: {i}, Actual: {doc.field('id')}"
|
||||
)
|
||||
assert doc.field("name") == f"test_{i}", (
|
||||
f"Document name field mismatch for document '{i}' after insertion. Expected: test_{i}, Actual: {doc.field('name')}"
|
||||
)
|
||||
assert doc.field("weight") == float(i * 10), (
|
||||
f"Document weight field mismatch for document '{i}' after insertion. Expected: {float(i * 10)}, Actual: {doc.field('weight')}"
|
||||
)
|
||||
|
||||
# Verify vector access after insertion
|
||||
assert doc.vector("dense") is not None, (
|
||||
f"Document {i} should have dense vector after insertion"
|
||||
)
|
||||
assert doc.vector("sparse") is not None, (
|
||||
f"Document {i} should have sparse vector after insertion"
|
||||
)
|
||||
|
||||
# Verify vector types after insertion
|
||||
assert isinstance(doc.vector("dense"), list), (
|
||||
f"Document {i} dense vector should be dict after insertion, got {type(doc.vector('dense'))}"
|
||||
)
|
||||
assert isinstance(doc.vector("sparse"), dict), (
|
||||
f"Document {i} sparse vector should be dict after insertion, got {type(doc.vector('sparse'))}"
|
||||
)
|
||||
|
||||
# Verify documents were inserted using stats
|
||||
stats = created_coll.stats
|
||||
assert stats is not None, "Collection stats should not be None"
|
||||
assert stats.doc_count == 3, (
|
||||
f"Document count mismatch after insertion. Expected: 3, Actual: {stats.doc_count}"
|
||||
)
|
||||
|
||||
# Store the collection path before cleanup
|
||||
collection_path = created_coll.path
|
||||
|
||||
# Clean up the created collection reference
|
||||
del created_coll
|
||||
|
||||
# Wait and verify the path still exists
|
||||
print(f"DEBUG: Collection path after destroy: {collection_path}")
|
||||
print(f"DEBUG: Path exists after destroy: {os.path.exists(collection_path)}")
|
||||
|
||||
# Now open the existing collection
|
||||
try:
|
||||
print(f"DEBUG: Path exists before open: {os.path.exists(collection_path)}")
|
||||
|
||||
# List contents of parent directory for debugging
|
||||
parent_dir = os.path.dirname(collection_path)
|
||||
if os.path.exists(parent_dir):
|
||||
print(f"DEBUG: Parent directory contents: {os.listdir(parent_dir)}")
|
||||
|
||||
opened_coll = zvec.open(path=collection_path, option=collection_option)
|
||||
|
||||
assert opened_coll is not None, (
|
||||
f"Failed to open existing collection at path: {collection_path}. Returned None instead of valid Collection object"
|
||||
)
|
||||
assert opened_coll.path == collection_path, (
|
||||
f"Opened collection path mismatch. Expected: {collection_path}, Actual: {opened_coll.path}"
|
||||
)
|
||||
assert opened_coll.schema.name == "test_collection", (
|
||||
f"Opened collection schema name mismatch. Expected: test_collection, Actual: {opened_coll.schema.name}"
|
||||
)
|
||||
|
||||
# Check reference count of opened collection
|
||||
opened_ref_count = sys.getrefcount(opened_coll)
|
||||
print(f"DEBUG: Reference count of opened collection: {opened_ref_count}")
|
||||
|
||||
# Verify data persistence
|
||||
# Verify data persistence using fetch interface
|
||||
fetched_docs = opened_coll.fetch(["0", "1", "2"])
|
||||
assert len(fetched_docs) == 3, (
|
||||
f"Expected 3 fetched documents after reopening, but got {len(fetched_docs)}"
|
||||
)
|
||||
assert "0" in fetched_docs, (
|
||||
"Document with ID '0' not found in fetched results after reopening"
|
||||
)
|
||||
assert "1" in fetched_docs, (
|
||||
"Document with ID '1' not found in fetched results after reopening"
|
||||
)
|
||||
assert "2" in fetched_docs, (
|
||||
"Document with ID '2' not found in fetched results after reopening"
|
||||
)
|
||||
|
||||
# Verify fetched document content after reopening collection
|
||||
for i in range(3):
|
||||
doc = fetched_docs[f"{i}"]
|
||||
assert doc is not None, (
|
||||
f"Fetched document with ID '{i}' is None after reopening collection"
|
||||
)
|
||||
assert doc.id == f"{i}", (
|
||||
f"Document ID mismatch for document '{i}' after reopening. Expected: {i}, Actual: {doc.id}"
|
||||
)
|
||||
assert doc.field("id") == i, (
|
||||
f"Document id field mismatch for document '{i}' after reopening. Expected: {i}, Actual: {doc.field('id')}"
|
||||
)
|
||||
assert doc.field("name") == f"test_{i}", (
|
||||
f"Document name field mismatch for document '{i}' after reopening. Expected: test_{i}, Actual: {doc.field('name')}"
|
||||
)
|
||||
assert doc.field("weight") == float(i * 10), (
|
||||
f"Document weight field mismatch for document '{i}' after reopening. Expected: {float(i * 10)}, Actual: {doc.field('weight')}"
|
||||
)
|
||||
|
||||
# Verify vector access after reopening
|
||||
assert doc.vector("dense") is not None, (
|
||||
f"Document {i} should have dense vector after reopening"
|
||||
)
|
||||
assert doc.vector("sparse") is not None, (
|
||||
f"Document {i} should have sparse vector after reopening"
|
||||
)
|
||||
|
||||
# Verify vector types after reopening
|
||||
assert isinstance(doc.vector("dense"), list), (
|
||||
f"Document {i} dense vector should be dict after reopening, got {type(doc.vector('dense'))}"
|
||||
)
|
||||
assert isinstance(doc.vector("sparse"), dict), (
|
||||
f"Document {i} sparse vector should be dict after reopening, got {type(doc.vector('sparse'))}"
|
||||
)
|
||||
|
||||
# Verify score attribute exists
|
||||
assert hasattr(doc, "score"), (
|
||||
f"Document {i} should have a score attribute after reopening"
|
||||
)
|
||||
assert isinstance(doc.score, (int, float)), (
|
||||
f"Document {i} score should be numeric after reopening, got {type(doc.score)}"
|
||||
)
|
||||
# For fetch operations, score is typically 0.0
|
||||
assert doc.score == 0.0, (
|
||||
f"Document {i} score should be 0.0 for fetch operation after reopening, but got {doc.score}"
|
||||
)
|
||||
|
||||
# Test query functionality
|
||||
query_result = opened_coll.query(include_vector=True)
|
||||
assert len(query_result) == 3, (
|
||||
f"Expected 3 query results, but got {len(query_result)}"
|
||||
)
|
||||
|
||||
# Verify query results have proper structure and content with detailed validation
|
||||
returned_doc_ids = set()
|
||||
for doc in query_result:
|
||||
# Verify basic document structure
|
||||
assert doc.id is not None, f"Query result document should have an ID"
|
||||
assert doc.id in ["0", "1", "2"], (
|
||||
f"Query result document ID should be one of ['0', '1', '2'], but got {doc.id}"
|
||||
)
|
||||
returned_doc_ids.add(doc.id)
|
||||
|
||||
# Verify field access
|
||||
assert doc.field("id") is not None, (
|
||||
f"Document {doc.id} should have id field"
|
||||
)
|
||||
assert doc.field("name") is not None, (
|
||||
f"Document {doc.id} should have name field"
|
||||
)
|
||||
assert doc.field("weight") is not None, (
|
||||
f"Document {doc.id} should have weight field"
|
||||
)
|
||||
|
||||
# Verify field values
|
||||
expected_id = int(doc.id)
|
||||
assert doc.field("id") == expected_id, (
|
||||
f"Document {doc.id} id field mismatch. Expected: {expected_id}, Actual: {doc.field('id')}"
|
||||
)
|
||||
assert doc.field("name") == f"test_{expected_id}", (
|
||||
f"Document {doc.id} name field mismatch. Expected: test_{expected_id}, Actual: {doc.field('name')}"
|
||||
)
|
||||
assert doc.field("weight") == float(expected_id * 10), (
|
||||
f"Document {doc.id} weight field mismatch. Expected: {float(expected_id * 10)}, Actual: {doc.field('weight')}"
|
||||
)
|
||||
|
||||
# Verify vector access
|
||||
assert doc.vector("dense") is not None, (
|
||||
f"Document {doc.id} should have dense vector"
|
||||
)
|
||||
assert doc.vector("sparse") is not None, (
|
||||
f"Document {doc.id} should have sparse vector"
|
||||
)
|
||||
|
||||
# Verify vector types
|
||||
assert isinstance(doc.vector("dense"), list), (
|
||||
f"Document {doc.id} dense vector should be list, got {type(doc.vector('dense'))}"
|
||||
)
|
||||
assert isinstance(doc.vector("sparse"), dict), (
|
||||
f"Document {doc.id} sparse vector should be dict, got {type(doc.vector('sparse'))}"
|
||||
)
|
||||
|
||||
# Verify score attribute exists
|
||||
assert hasattr(doc, "score"), (
|
||||
f"Document {doc.id} should have a score attribute"
|
||||
)
|
||||
assert isinstance(doc.score, (int, float)), (
|
||||
f"Document {doc.id} score should be numeric, got {type(doc.score)}"
|
||||
)
|
||||
|
||||
# Verify all expected documents are returned
|
||||
expected_doc_ids = {"0", "1", "2"}
|
||||
assert returned_doc_ids == expected_doc_ids, (
|
||||
f"Query should return all expected documents. Expected: {expected_doc_ids}, Actual: {returned_doc_ids}"
|
||||
)
|
||||
|
||||
# === Enhanced validation based on test_collection_dql_operations.py ===
|
||||
|
||||
# Verify vector field names accessibility for all documents
|
||||
for doc in query_result:
|
||||
vector_names = doc.vector_names()
|
||||
expected_vector_names = {"dense", "sparse"}
|
||||
assert set(vector_names) == expected_vector_names, (
|
||||
f"Document {doc.id} vector names mismatch. Expected: {expected_vector_names}, Actual: {set(vector_names)}"
|
||||
)
|
||||
|
||||
# Verify all vector fields can be accessed
|
||||
for vector_name in expected_vector_names:
|
||||
vector_data = doc.vector(vector_name)
|
||||
assert vector_data is not None, (
|
||||
f"Document {doc.id} should have accessible vector '{vector_name}'"
|
||||
)
|
||||
if vector_name == "dense":
|
||||
assert isinstance(vector_data, list), (
|
||||
f"Document {doc.id} vector '{vector_name}' should be list, got {type(vector_data)}"
|
||||
)
|
||||
else:
|
||||
assert isinstance(vector_data, dict), (
|
||||
f"Document {doc.id} vector '{vector_name}' should be dict, got {type(vector_data)}"
|
||||
)
|
||||
|
||||
# Test query with filter
|
||||
filtered_result = opened_coll.query(filter="id >= 1", include_vector=True)
|
||||
assert len(filtered_result) == 2, (
|
||||
f"Expected 2 filtered query results (id >= 1), but got {len(filtered_result)}"
|
||||
)
|
||||
|
||||
# Verify filtered query results
|
||||
filtered_doc_ids = set()
|
||||
for doc in filtered_result:
|
||||
assert doc.id is not None, (
|
||||
f"Filtered query result document should have an ID"
|
||||
)
|
||||
assert doc.id in ["1", "2"], (
|
||||
f"Filtered query result document ID should be one of ['1', '2'], but got {doc.id}"
|
||||
)
|
||||
filtered_doc_ids.add(doc.id)
|
||||
|
||||
# Verify filter condition is satisfied
|
||||
doc_id = int(doc.id)
|
||||
assert doc_id >= 1, (
|
||||
f"Document {doc.id} should satisfy filter condition id >= 1"
|
||||
)
|
||||
|
||||
# Verify document structure
|
||||
assert doc.field("id") is not None, (
|
||||
f"Document {doc.id} should have id field"
|
||||
)
|
||||
assert doc.field("name") is not None, (
|
||||
f"Document {doc.id} should have name field"
|
||||
)
|
||||
assert doc.field("weight") is not None, (
|
||||
f"Document {doc.id} should have weight field"
|
||||
)
|
||||
|
||||
# Verify field values
|
||||
assert doc.field("id") == doc_id, (
|
||||
f"Document {doc.id} id field mismatch. Expected: {doc_id}, Actual: {doc.field('id')}"
|
||||
)
|
||||
assert doc.field("name") == f"test_{doc_id}", (
|
||||
f"Document {doc.id} name field mismatch. Expected: test_{doc_id}, Actual: {doc.field('name')}"
|
||||
)
|
||||
assert doc.field("weight") == float(doc_id * 10), (
|
||||
f"Document {doc.id} weight field mismatch. Expected: {float(doc_id * 10)}, Actual: {doc.field('weight')}"
|
||||
)
|
||||
|
||||
# Verify vector access
|
||||
assert doc.vector("dense") is not None, (
|
||||
f"Document {doc.id} should have dense vector"
|
||||
)
|
||||
assert doc.vector("sparse") is not None, (
|
||||
f"Document {doc.id} should have sparse vector"
|
||||
)
|
||||
|
||||
# Verify score attribute exists
|
||||
assert hasattr(doc, "score"), (
|
||||
f"Document {doc.id} should have a score attribute"
|
||||
)
|
||||
assert isinstance(doc.score, (int, float)), (
|
||||
f"Document {doc.id} score should be numeric, got {type(doc.score)}"
|
||||
)
|
||||
|
||||
# Verify filtered documents
|
||||
expected_filtered_ids = {"1", "2"}
|
||||
assert filtered_doc_ids == expected_filtered_ids, (
|
||||
f"Filtered query should return expected documents. Expected: {expected_filtered_ids}, Actual: {filtered_doc_ids}"
|
||||
)
|
||||
|
||||
# Test vector query functionality for dense vectors
|
||||
query_vector_dense = [0.1] * 128
|
||||
vector_query_result = opened_coll.query(
|
||||
Query(field_name="dense", vector=query_vector_dense)
|
||||
)
|
||||
assert len(vector_query_result) > 0, (
|
||||
f"Expected at least 1 vector query result, but got {len(vector_query_result)}"
|
||||
)
|
||||
|
||||
# Verify vector query results structure
|
||||
for doc in vector_query_result[:3]: # Check first 3 results
|
||||
assert doc.id is not None, (
|
||||
f"Vector query result document should have an ID"
|
||||
)
|
||||
assert doc.id in ["0", "1", "2"], (
|
||||
f"Vector query result document ID should be one of ['0', '1', '2'], but got {doc.id}"
|
||||
)
|
||||
|
||||
# Verify document structure
|
||||
assert doc.field("id") is not None, (
|
||||
f"Document {doc.id} should have id field"
|
||||
)
|
||||
assert doc.field("name") is not None, (
|
||||
f"Document {doc.id} should have name field"
|
||||
)
|
||||
assert doc.field("weight") is not None, (
|
||||
f"Document {doc.id} should have weight field"
|
||||
)
|
||||
|
||||
# Verify vector access
|
||||
assert doc.vector("dense") is not None, (
|
||||
f"Document {doc.id} should have dense vector"
|
||||
)
|
||||
assert doc.vector("sparse") is not None, (
|
||||
f"Document {doc.id} should have sparse vector"
|
||||
)
|
||||
|
||||
# Verify score attribute exists and is numeric
|
||||
assert hasattr(doc, "score"), (
|
||||
f"Document {doc.id} should have a score attribute"
|
||||
)
|
||||
assert isinstance(doc.score, (int, float)), (
|
||||
f"Document {doc.id} score should be numeric, got {type(doc.score)}"
|
||||
)
|
||||
|
||||
# For dense vector queries, score should typically be non-negative (depending on metric)
|
||||
# Note: This may vary based on the metric type used
|
||||
assert doc.score >= 0 or doc.score < 0, (
|
||||
f"Document {doc.id} score should be a valid number"
|
||||
)
|
||||
|
||||
# Test vector query functionality for sparse vectors
|
||||
query_vector_sparse = {1: 1.0, 2: 2.0, 3: 3.0}
|
||||
sparse_vector_query_result = opened_coll.query(
|
||||
Query(field_name="sparse", vector=query_vector_sparse)
|
||||
)
|
||||
assert len(sparse_vector_query_result) > 0, (
|
||||
f"Expected at least 1 sparse vector query result, but got {len(sparse_vector_query_result)}"
|
||||
)
|
||||
|
||||
# Verify sparse vector query results structure
|
||||
for doc in sparse_vector_query_result[:3]: # Check first 3 results
|
||||
assert doc.id is not None, (
|
||||
f"Sparse vector query result document should have an ID"
|
||||
)
|
||||
assert doc.id in ["0", "1", "2"], (
|
||||
f"Sparse vector query result document ID should be one of ['0', '1', '2'], but got {doc.id}"
|
||||
)
|
||||
|
||||
# Verify document structure
|
||||
assert doc.field("id") is not None, (
|
||||
f"Document {doc.id} should have id field"
|
||||
)
|
||||
assert doc.field("name") is not None, (
|
||||
f"Document {doc.id} should have name field"
|
||||
)
|
||||
assert doc.field("weight") is not None, (
|
||||
f"Document {doc.id} should have weight field"
|
||||
)
|
||||
|
||||
# Verify vector access
|
||||
assert doc.vector("dense") is not None, (
|
||||
f"Document {doc.id} should have dense vector"
|
||||
)
|
||||
assert doc.vector("sparse") is not None, (
|
||||
f"Document {doc.id} should have sparse vector"
|
||||
)
|
||||
|
||||
# Verify score attribute exists and is numeric
|
||||
assert hasattr(doc, "score"), (
|
||||
f"Document {doc.id} should have a score attribute"
|
||||
)
|
||||
assert isinstance(doc.score, (int, float)), (
|
||||
f"Document {doc.id} score should be numeric, got {type(doc.score)}"
|
||||
)
|
||||
|
||||
# Clean up
|
||||
if hasattr(opened_coll, "destroy") and opened_coll is not None:
|
||||
opened_coll.destroy()
|
||||
print("DEBUG: Opened collection destroyed successfully")
|
||||
|
||||
except Exception as e:
|
||||
logging.error("Exception occurred: [{}]".format(e))
|
||||
raise e
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"read_only,enable_mmap,description", COLLECTION_OPTION_TEST_CASES_VALID
|
||||
)
|
||||
@pytest.mark.parametrize("createAndopen_enable_mmap", [True, False])
|
||||
def test_open_with_different_collection_options_valid(
|
||||
self,
|
||||
tmp_path_factory,
|
||||
createAndopen_enable_mmap,
|
||||
read_only,
|
||||
enable_mmap,
|
||||
description,
|
||||
collection_schema,
|
||||
):
|
||||
# Create collection with initial option
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "test_collection"
|
||||
|
||||
initial_option = CollectionOption(
|
||||
read_only=False, enable_mmap=createAndopen_enable_mmap
|
||||
)
|
||||
|
||||
# Create and open collection first
|
||||
created_coll = zvec.create_and_open(
|
||||
path=str(collection_path), schema=collection_schema, option=initial_option
|
||||
)
|
||||
|
||||
assert created_coll is not None, "Failed to create collection"
|
||||
|
||||
# Clean up the created collection reference
|
||||
del created_coll
|
||||
|
||||
# Now open with different options
|
||||
collection_option = CollectionOption(
|
||||
read_only=read_only, enable_mmap=enable_mmap
|
||||
)
|
||||
|
||||
try:
|
||||
opened_coll = zvec.open(path=str(collection_path), option=collection_option)
|
||||
|
||||
assert opened_coll is not None, (
|
||||
f"Failed to open collection with option: {description}. Returned None instead of valid Collection object. Path: {collection_path}"
|
||||
)
|
||||
assert opened_coll.path == str(collection_path), (
|
||||
f"Opened collection path mismatch. Expected: {collection_path}, Actual: {opened_coll.path}"
|
||||
)
|
||||
assert opened_coll.schema.name == collection_schema.name, (
|
||||
f"Opened collection schema name mismatch. Expected: {collection_schema.name}, Actual: {opened_coll.schema.name}"
|
||||
)
|
||||
assert opened_coll.option.read_only == read_only, (
|
||||
f"Opened collection read_only option mismatch. Expected: {read_only}, Actual: {opened_coll.option.read_only}"
|
||||
)
|
||||
assert opened_coll.option.enable_mmap == createAndopen_enable_mmap, (
|
||||
f"Opened collection mmap option mismatch. Expected: {createAndopen_enable_mmap}, Actual: {opened_coll.option.enable_mmap}"
|
||||
)
|
||||
|
||||
# Clean up
|
||||
if (
|
||||
hasattr(opened_coll, "destroy")
|
||||
and opened_coll is not None
|
||||
and read_only == False
|
||||
):
|
||||
opened_coll.destroy()
|
||||
|
||||
except Exception as e:
|
||||
logging.error("Exception occurred: [{}]".format(e))
|
||||
pytest.fail(f"Failed to open collection with different options: {e}")
|
||||
|
||||
def test_open_with_none_option(self, tmp_path_factory, collection_schema):
|
||||
# Create collection
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "test_collection"
|
||||
|
||||
initial_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
# Create and open collection first
|
||||
created_coll = zvec.create_and_open(
|
||||
path=str(collection_path), schema=collection_schema, option=initial_option
|
||||
)
|
||||
|
||||
assert created_coll is not None, (
|
||||
f"Failed to create collection. Returned None instead of valid Collection object. Path: {collection_path}"
|
||||
)
|
||||
|
||||
# Clean up the created collection reference
|
||||
del created_coll
|
||||
|
||||
# Now open with None option
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
zvec.open(path=str(collection_path), option=None)
|
||||
|
||||
assert "incompatible function arguments" in str(exc_info.value), (
|
||||
f"Expected 'incompatible function arguments' error, but got: {exc_info.value}"
|
||||
)
|
||||
|
||||
def test_reopen_collection(self, tmp_path_factory):
|
||||
# Prepare schema
|
||||
collection_schema = 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(
|
||||
"description",
|
||||
DataType.STRING,
|
||||
nullable=True,
|
||||
index_param=InvertIndexParam(),
|
||||
),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
collection_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
# Create collection
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "test_collection"
|
||||
|
||||
# Create and open collection
|
||||
coll1 = zvec.create_and_open(
|
||||
path=str(collection_path),
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert coll1 is not None, "Failed to create and open collection"
|
||||
|
||||
# Insert some data
|
||||
doc = Doc(
|
||||
id="1",
|
||||
fields={"id": 1, "name": "test", "description": "这是一个中文描述。"},
|
||||
vectors={"dense": np.random.random(128).tolist()},
|
||||
)
|
||||
|
||||
result = coll1.insert(doc)
|
||||
assert result.ok()
|
||||
|
||||
# Close the first collection (delete reference)
|
||||
del coll1
|
||||
|
||||
# Reopen the collection
|
||||
coll2 = zvec.open(path=str(collection_path), option=collection_option)
|
||||
|
||||
assert coll2 is not None, "Failed to reopen collection"
|
||||
assert coll2.path == str(collection_path)
|
||||
assert coll2.schema.name == collection_schema.name
|
||||
|
||||
# Verify data is still there
|
||||
fetched_docs = coll2.fetch(["1"])
|
||||
assert "1" in fetched_docs
|
||||
fetched_doc = fetched_docs["1"]
|
||||
assert fetched_doc.id == "1"
|
||||
assert fetched_doc.field("name") == "test"
|
||||
assert fetched_doc.field("description") == "这是一个中文描述。"
|
||||
|
||||
# Clean up
|
||||
if hasattr(coll2, "destroy") and coll2 is not None:
|
||||
try:
|
||||
coll2.destroy()
|
||||
except Exception as e:
|
||||
print(f"Warning: failed to destroy collection: {e}")
|
||||
|
||||
def test_open_concurrent_same_path(self, tmp_path_factory):
|
||||
# First create a collection
|
||||
collection_schema = 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(),
|
||||
),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
collection_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
# Create collection path
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "test_collection"
|
||||
|
||||
# First create the collection
|
||||
created_coll = zvec.create_and_open(
|
||||
path=str(collection_path),
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert created_coll is not None, "Failed to create collection"
|
||||
|
||||
# Close the collection so we can test concurrent opening
|
||||
if hasattr(created_coll, "close") and created_coll is not None:
|
||||
created_coll.close()
|
||||
|
||||
# Shared variables to collect results from threads
|
||||
results = []
|
||||
errors = []
|
||||
|
||||
# Lock for thread-safe operations
|
||||
lock = threading.Lock()
|
||||
# Clean up the created collection reference
|
||||
del created_coll
|
||||
|
||||
# Function to be executed by each thread
|
||||
def open_collection_thread(thread_id):
|
||||
try:
|
||||
coll = zvec.open(path=str(collection_path), option=collection_option)
|
||||
with lock:
|
||||
results.append((thread_id, coll))
|
||||
# Close the collection if opened successfully
|
||||
if hasattr(coll, "close") and coll is not None:
|
||||
coll.close()
|
||||
except Exception as e:
|
||||
with lock:
|
||||
errors.append((thread_id, str(e)))
|
||||
|
||||
# Create 5 threads to call open concurrently
|
||||
threads = []
|
||||
for i in range(5):
|
||||
thread = threading.Thread(target=open_collection_thread, args=(i,))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Wait for all threads to complete
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# Verify concurrency safety: only one should succeed, others should fail
|
||||
assert len(results) == 1, (
|
||||
f"Expected exactly one successful open, but got {len(results)}"
|
||||
)
|
||||
assert len(errors) == 4, (
|
||||
f"Expected exactly four failures, but got {len(errors)}"
|
||||
)
|
||||
|
||||
# Additional verification: check that the successful open has a valid collection
|
||||
successful_thread_id, successful_collection = results[0]
|
||||
assert successful_collection is not None, (
|
||||
"Successful open should return a valid collection"
|
||||
)
|
||||
assert successful_collection.path == str(collection_path), (
|
||||
"Collection path mismatch"
|
||||
)
|
||||
|
||||
# Clean up the successfully opened collection
|
||||
if (
|
||||
hasattr(successful_collection, "destroy")
|
||||
and successful_collection is not None
|
||||
):
|
||||
try:
|
||||
successful_collection.destroy()
|
||||
except Exception as e:
|
||||
print(f"Warning: failed to destroy collection: {e}")
|
||||
|
||||
def test_open_with_corrupted_files(self, tmp_path_factory):
|
||||
# First create a collection
|
||||
collection_schema = 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(),
|
||||
),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"dense",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=HnswIndexParam(),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
collection_option = CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
# Create collection path
|
||||
temp_dir = tmp_path_factory.mktemp("zvec")
|
||||
collection_path = temp_dir / "test_collection"
|
||||
|
||||
# First create the collection
|
||||
created_coll = zvec.create_and_open(
|
||||
path=str(collection_path),
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert created_coll is not None, "Failed to create collection"
|
||||
|
||||
# Close the collection so we can manipulate its files
|
||||
if hasattr(created_coll, "close") and created_coll is not None:
|
||||
created_coll.close()
|
||||
|
||||
# Test case 1: Delete some files in the collection directory (simulate partial corruption)
|
||||
import os
|
||||
import shutil
|
||||
import random
|
||||
|
||||
# Get the collection directory path
|
||||
collection_dir = str(collection_path)
|
||||
|
||||
# List all files in the collection directory
|
||||
files_in_dir = []
|
||||
for root, dirs, files in os.walk(collection_dir):
|
||||
for file in files:
|
||||
files_in_dir.append(os.path.join(root, file))
|
||||
|
||||
# Randomly delete approximately half of the files to simulate partial corruption
|
||||
if files_in_dir:
|
||||
# Shuffle the list to randomly select files
|
||||
random.shuffle(files_in_dir)
|
||||
files_to_delete = files_in_dir[: len(files_in_dir) // 2]
|
||||
for file_path in files_to_delete:
|
||||
try:
|
||||
os.remove(file_path)
|
||||
except Exception as e:
|
||||
pass # Ignore errors during deletion
|
||||
|
||||
# Try to open the collection with missing files - should raise an exception
|
||||
with pytest.raises(Exception):
|
||||
zvec.open(path=str(collection_path), option=collection_option)
|
||||
|
||||
# Test case 2: Delete all files in the collection directory (simulate complete corruption)
|
||||
# Recreate the collection
|
||||
recreated_coll = zvec.create_and_open(
|
||||
path=str(collection_path) + "_all",
|
||||
schema=collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert recreated_coll is not None, "Failed to recreate collection"
|
||||
|
||||
# Close the collection so we can manipulate its files
|
||||
if hasattr(recreated_coll, "close") and recreated_coll is not None:
|
||||
recreated_coll.close()
|
||||
|
||||
# Delete all files in the collection directory
|
||||
try:
|
||||
shutil.rmtree(collection_dir)
|
||||
os.makedirs(collection_dir) # Recreate empty directory
|
||||
except Exception as e:
|
||||
pass # Ignore errors during deletion
|
||||
|
||||
# Try to open the collection with missing files - should raise an exception
|
||||
with pytest.raises(Exception):
|
||||
zvec.open(path=str(collection_path), option=collection_option)
|
||||
@@ -0,0 +1,740 @@
|
||||
# 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
|
||||
@@ -0,0 +1,307 @@
|
||||
# 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 logging
|
||||
import pytest
|
||||
import tempfile
|
||||
import os
|
||||
import sys
|
||||
import subprocess
|
||||
|
||||
import zvec
|
||||
import zvec
|
||||
from zvec import LogType, LogLevel
|
||||
|
||||
# Error messages
|
||||
INITIALIZATION_ERROR_MSG = "initialization failed"
|
||||
RUNTIME_ERROR_MSG = "RuntimeError"
|
||||
VALUE_ERROR_MSG = "ValueError"
|
||||
TYPE_ERROR_MSG = "TypeError"
|
||||
|
||||
|
||||
# ==================== helper ====================
|
||||
def run_in_subprocess(func):
|
||||
def wrapper(*args, **kwargs):
|
||||
if os.getenv("RUNNING_IN_SUBPROCESS"):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
env = os.environ.copy()
|
||||
env["RUNNING_IN_SUBPROCESS"] = "1"
|
||||
env["PYTEST_CURRENT_TEST"] = func.__name__
|
||||
|
||||
import inspect
|
||||
|
||||
filepath = inspect.getfile(func)
|
||||
qualname = func.__qualname__.replace(".", "::")
|
||||
test_id = f"{filepath}::{qualname}"
|
||||
|
||||
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
env["PYTHONPATH"] = project_root + ":" + env.get("PYTHONPATH", "")
|
||||
|
||||
cmd = [sys.executable, "-m", "pytest", "-v", "-s", test_id]
|
||||
|
||||
result = subprocess.run(cmd, env=env, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
pytest.fail(
|
||||
f"Subprocess test {func.__name__} failed with code {result.returncode}\n"
|
||||
f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
|
||||
)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
# ==================== Fixtures ====================
|
||||
@pytest.fixture(scope="function")
|
||||
def temp_log_dir(tmp_path_factory):
|
||||
return tmp_path_factory.mktemp("logs")
|
||||
|
||||
|
||||
# ==================== Tests ====================
|
||||
class TestDbConfigInitialization:
|
||||
@run_in_subprocess
|
||||
def test_init_default(self):
|
||||
# default config
|
||||
# log_type: Optional[LogType] = LogType.CONSOLE,
|
||||
# log_level: Optional[LogLevel] = LogLevel.WARN,
|
||||
# log_dir: Optional[str] = "./logs",
|
||||
# log_basename: Optional[str] = "zvec.log",
|
||||
# log_file_size: Optional[int] = 2048,
|
||||
# log_overdue_days: Optional[int] = 7,
|
||||
zvec.init()
|
||||
|
||||
@run_in_subprocess
|
||||
def test_init_file_logger(self):
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
|
||||
zvec.init(
|
||||
log_level=LogLevel.DEBUG,
|
||||
log_type=LogType.FILE,
|
||||
)
|
||||
# assert logdir exist
|
||||
log_dir = Path("./logs")
|
||||
assert log_dir.exists()
|
||||
|
||||
# validate write log
|
||||
col = zvec.create_and_open(
|
||||
"/tmp/test/1",
|
||||
zvec.CollectionSchema(
|
||||
name="test",
|
||||
vectors=zvec.VectorSchema(
|
||||
dimension=4,
|
||||
data_type=zvec.DataType.VECTOR_FP32,
|
||||
name="image",
|
||||
),
|
||||
),
|
||||
)
|
||||
col.insert(docs=[zvec.Doc(id="1", vectors={"image": [1.0, 2.0, 3.0, 4.0]})])
|
||||
assert any(log_dir.glob("zvec.log.*"))
|
||||
|
||||
# clear
|
||||
col.destroy()
|
||||
shutil.rmtree(log_dir, ignore_errors=True)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_init_with_mixed_config(self):
|
||||
zvec.init(
|
||||
memory_limit_mb=128,
|
||||
log_type=LogType.FILE,
|
||||
query_threads=1,
|
||||
log_level=LogLevel.WARN,
|
||||
)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_repeated_initialization(self):
|
||||
# Calling init() repeatedly is allowed:
|
||||
# it succeeds but becomes a no-op after the first successful init()
|
||||
zvec.init()
|
||||
|
||||
|
||||
class TestDbConfigMemoryLimitValidation:
|
||||
@run_in_subprocess
|
||||
def test_memory_limit_min_valid(self):
|
||||
# MIN_MEMORY_LIMIT_BYTES is 100M
|
||||
with pytest.raises(RuntimeError):
|
||||
zvec.init(memory_limit_mb=99)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_memory_limit_invalid_value(self):
|
||||
# memory_limit_mb must >= 0 and must be int and if None, set default value
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(memory_limit_mb=0)
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(memory_limit_mb=-1)
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(memory_limit_mb="512")
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(memory_limit_mb=512.5)
|
||||
|
||||
|
||||
class TestDbConfigThreadValidation:
|
||||
@run_in_subprocess
|
||||
def test_query_threads(self):
|
||||
zvec.init(query_threads=1)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_query_threads_invalid(self):
|
||||
# query_threads must >= 0 and must be int and if None, set default value
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(query_threads=0)
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(query_threads=-1)
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(query_threads="value")
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(query_threads=512.5)
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(query_threads="512")
|
||||
|
||||
@run_in_subprocess
|
||||
def test_optimize_threads(self):
|
||||
zvec.init(optimize_threads=1)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_optimize_threads_invalid(self):
|
||||
# optimize_threads must >= 0 and must be int and if None, set default value
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(optimize_threads=0)
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(optimize_threads=-1)
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(optimize_threads="value")
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(optimize_threads=512.5)
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(optimize_threads="512")
|
||||
|
||||
|
||||
class TestDbConfigRatioValidation:
|
||||
@run_in_subprocess
|
||||
def test_init_invert_to_forward_scan_ratio(self):
|
||||
# must be in [0,1]
|
||||
zvec.init(invert_to_forward_scan_ratio=0.8)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_init_invert_to_forward_scan_ratio_invalid(self):
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(invert_to_forward_scan_ratio=1.1)
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(invert_to_forward_scan_ratio=-0.1)
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(invert_to_forward_scan_ratio="0.8")
|
||||
|
||||
@run_in_subprocess
|
||||
def test_init_brute_force_by_keys_ratio(self):
|
||||
zvec.init(brute_force_by_keys_ratio=0.8)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_init_brute_force_by_keys_ratio_invalid(self):
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(brute_force_by_keys_ratio=1.1)
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(brute_force_by_keys_ratio=-0.1)
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(brute_force_by_keys_ratio="0.8")
|
||||
|
||||
|
||||
class TestDbConfigLogValidation:
|
||||
@run_in_subprocess
|
||||
def test_log_type_valid(self):
|
||||
zvec.init(log_type=LogType.CONSOLE)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_log_type_invalid(self):
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_type="FILE")
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_type="")
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_type="invalid")
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_type=123)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_log_level_valid(self):
|
||||
zvec.init(log_level=LogLevel.ERROR)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_log_level_invalid(self):
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_level="WARN")
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_level="")
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_level="invalid")
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_level=123)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_init_file_logger(self):
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
|
||||
temp_dir = tempfile.mkdtemp(prefix="log_test_")
|
||||
abs_temp_dir = os.path.abspath(temp_dir)
|
||||
|
||||
zvec.init(
|
||||
log_level=LogLevel.DEBUG,
|
||||
log_type=LogType.FILE,
|
||||
log_dir=abs_temp_dir,
|
||||
log_basename="test",
|
||||
)
|
||||
|
||||
# assert logdir exist
|
||||
log_dir = Path(abs_temp_dir)
|
||||
assert log_dir.exists()
|
||||
|
||||
# validate write log
|
||||
col = zvec.create_and_open(
|
||||
"/tmp/test/1",
|
||||
zvec.CollectionSchema(
|
||||
name="test",
|
||||
vectors=zvec.VectorSchema(
|
||||
dimension=4,
|
||||
data_type=zvec.DataType.VECTOR_FP32,
|
||||
name="image",
|
||||
),
|
||||
),
|
||||
)
|
||||
col.insert(docs=[zvec.Doc(id="1", vectors={"image": [1.0, 2.0, 3.0, 4.0]})])
|
||||
assert any(log_dir.glob("test.*"))
|
||||
|
||||
# clear
|
||||
col.destroy()
|
||||
shutil.rmtree(log_dir, ignore_errors=True)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_log_file_size_invalid(self):
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_type=LogType.FILE, log_file_size="df")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(log_type=LogType.FILE, log_file_size=0)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(log_type=LogType.FILE, log_file_size=-1)
|
||||
|
||||
@run_in_subprocess
|
||||
def test_log_overdue_days_invalid(self):
|
||||
with pytest.raises(TypeError):
|
||||
zvec.init(log_type=LogType.FILE, log_overdue_days="df")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(log_type=LogType.FILE, log_overdue_days=0)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
zvec.init(log_type=LogType.FILE, log_overdue_days=-1)
|
||||
Reference in New Issue
Block a user