chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:42 +08:00
commit be3ef883e1
1214 changed files with 431743 additions and 0 deletions
+392
View File
@@ -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
+464
View File
@@ -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
+652
View File
@@ -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
+209
View File
@@ -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"
+126
View File
@@ -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
+967
View File
@@ -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
+307
View File
@@ -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)
File diff suppressed because it is too large Load Diff
+600
View File
@@ -0,0 +1,600 @@
# 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.
"""End-to-end collection tests for the DiskAnn index.
Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn plugin.
Two platform-level prerequisites are enforced at module import time:
1. DiskAnn is currently built only for Linux x86_64 — other platforms are
skipped wholesale.
2. The DiskAnn backend lives in a *runtime-loaded* plugin
(``libzvec_diskann_plugin.so``). It must be loaded with ``RTLD_GLOBAL |
RTLD_NOW`` BEFORE ``import zvec`` so that the plugin's ``IndexFactory``
singleton is unified with the one inside ``_zvec.so``. After ``import
zvec`` we must also call ``zvec.load_diskann_plugin()`` exactly once.
If either prerequisite fails the whole module is skipped so the rest of the
test-suite is not affected.
"""
from __future__ import annotations
import math
import os
import platform
import sys
import pytest
# --------------------------------------------------------------------------- #
# Platform gating (must happen BEFORE we touch zvec).
# --------------------------------------------------------------------------- #
pytestmark = pytest.mark.skipif(
not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")),
reason="DiskAnn plugin is only supported on Linux x86_64",
)
# Promote all symbols in subsequently-loaded DSOs to the global namespace and
# resolve relocations eagerly. This is REQUIRED so the DiskAnn plugin can see
# the ``IndexFactory`` singleton that lives in ``_zvec.so`` and vice versa.
# See: DiskAnn RTLD_GLOBAL + RTLD_NOW Requirement.
if sys.platform == "linux":
sys.setdlopenflags(sys.getdlopenflags() | os.RTLD_GLOBAL | os.RTLD_NOW)
import zvec # noqa: E402
from zvec import ( # noqa: E402
Collection,
CollectionOption,
DataType,
DiskAnnIndexParam,
DiskAnnQueryParam,
Doc,
FieldSchema,
MetricType,
Query,
VectorSchema,
)
from zvec.typing import QuantizeType # noqa: E402
@pytest.fixture(scope="session")
def diskann_collection_schema():
"""Create a collection schema with a DiskAnn index."""
return zvec.CollectionSchema(
name="test_diskann_collection",
fields=[
FieldSchema("id", DataType.INT64, nullable=False),
FieldSchema("name", DataType.STRING, nullable=False),
],
vectors=[
VectorSchema(
"embedding",
DataType.VECTOR_FP32,
dimension=128,
index_param=DiskAnnIndexParam(
metric_type=MetricType.L2,
max_degree=64,
list_size=100,
pq_chunk_num=0,
quantize_type=QuantizeType.UNDEFINED,
),
),
],
)
@pytest.fixture(scope="session")
def collection_option():
"""Create collection options."""
return CollectionOption(read_only=False, enable_mmap=True)
@pytest.fixture
def single_doc():
"""Create a single document for testing."""
return Doc(
id="0",
fields={"id": 0, "name": "test_doc_0"},
vectors={"embedding": [0.1 + i * 0.01 for i in range(128)]},
)
@pytest.fixture
def multiple_docs():
"""Create multiple documents for testing."""
return [
Doc(
id=f"{i}",
fields={"id": i, "name": f"test_doc_{i}"},
vectors={"embedding": [i * 0.1 + j * 0.01 for j in range(128)]},
)
for i in range(1, 101)
]
@pytest.fixture(scope="function")
def diskann_collection(
tmp_path_factory, diskann_collection_schema, collection_option
) -> Collection:
"""
Function-scoped fixture: creates and opens a collection with DiskAnn index.
"""
temp_dir = tmp_path_factory.mktemp("zvec_diskann")
collection_path = temp_dir / "test_diskann_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=diskann_collection_schema,
option=collection_option,
)
assert coll is not None, "Failed to create and open DiskAnn collection"
assert coll.path == str(collection_path)
assert coll.schema.name == diskann_collection_schema.name
try:
yield coll
finally:
if hasattr(coll, "destroy") and coll is not None:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
@pytest.fixture
def collection_with_single_doc(
diskann_collection: Collection, single_doc: Doc
) -> Collection:
"""Setup: insert single doc into collection."""
assert diskann_collection.stats.doc_count == 0
result = diskann_collection.insert(single_doc)
assert bool(result)
assert result.ok()
assert diskann_collection.stats.doc_count == 1
yield diskann_collection
# Teardown: delete single doc
diskann_collection.delete(single_doc.id)
assert diskann_collection.stats.doc_count == 0
@pytest.fixture
def collection_with_multiple_docs(
diskann_collection: Collection, multiple_docs: list[Doc]
) -> Collection:
"""Setup: insert multiple docs into collection."""
assert diskann_collection.stats.doc_count == 0
result = diskann_collection.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
assert diskann_collection.stats.doc_count == len(multiple_docs)
yield diskann_collection
# Teardown: delete multiple docs
diskann_collection.delete([doc.id for doc in multiple_docs])
# ==================== Tests ====================
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionCreation:
"""Test DiskAnn collection creation and schema validation."""
def test_collection_creation(
self, diskann_collection: Collection, diskann_collection_schema
):
"""Test that collection is created with correct schema."""
assert diskann_collection is not None
assert diskann_collection.schema.name == diskann_collection_schema.name
assert len(diskann_collection.schema.fields) == len(
diskann_collection_schema.fields
)
assert len(diskann_collection.schema.vectors) == len(
diskann_collection_schema.vectors
)
def test_vector_schema_validation(self, diskann_collection: Collection):
"""Test that vector schema has correct DiskAnn configuration."""
vector_schema = diskann_collection.schema.vector("embedding")
assert vector_schema is not None
assert vector_schema.name == "embedding"
assert vector_schema.data_type == DataType.VECTOR_FP32
assert vector_schema.dimension == 128
index_param = vector_schema.index_param
assert index_param is not None
assert index_param.metric_type == MetricType.L2
assert index_param.max_degree == 64
assert index_param.list_size == 100
assert index_param.pq_chunk_num == 0
def test_collection_stats(self, diskann_collection: Collection):
"""Test initial collection statistics."""
stats = diskann_collection.stats
assert stats is not None
assert stats.doc_count == 0
assert len(stats.index_completeness) == 1
assert stats.index_completeness["embedding"] == 1
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionInsert:
"""Test document insertion into DiskAnn collection."""
def test_insert_single_doc(self, diskann_collection: Collection, single_doc: Doc):
"""Test inserting a single document."""
result = diskann_collection.insert(single_doc)
assert bool(result)
assert result.ok()
stats = diskann_collection.stats
assert stats is not None
assert stats.doc_count == 1
def test_insert_multiple_docs(
self, diskann_collection: Collection, multiple_docs: list[Doc]
):
"""Test inserting multiple documents."""
result = diskann_collection.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
stats = diskann_collection.stats
assert stats is not None
assert stats.doc_count == len(multiple_docs)
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionFetch:
"""Test document fetching from DiskAnn collection."""
def test_fetch_single_doc(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test fetching a single document by ID."""
result = collection_with_single_doc.fetch(ids=[single_doc.id])
assert bool(result)
assert single_doc.id in result.keys()
doc = result[single_doc.id]
assert doc is not None
assert doc.id == single_doc.id
assert doc.field("id") == single_doc.field("id")
assert doc.field("name") == single_doc.field("name")
def test_fetch_multiple_docs(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test fetching multiple documents by IDs."""
ids = [doc.id for doc in multiple_docs[:10]]
result = collection_with_multiple_docs.fetch(ids=ids)
assert bool(result)
assert len(result) == len(ids)
for doc_id in ids:
assert doc_id in result
doc = result[doc_id]
assert doc is not None
assert doc.id == doc_id
def test_fetch_nonexistent_doc(self, collection_with_single_doc: Collection):
"""Test fetching a non-existent document."""
result = collection_with_single_doc.fetch(ids=["nonexistent_id"])
assert len(result) == 0
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionQuery:
"""Test vector search queries on DiskAnn collection."""
def test_query_by_vector(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying by vector with DiskAnn index."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result = collection_with_multiple_docs.query(queries=query, topk=10)
assert len(result) > 0
assert len(result) <= 10
# First result should be the query document itself (or very close)
first_doc = result[0]
assert first_doc is not None
assert first_doc.id is not None
def test_query_by_id(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying by document ID with DiskAnn index."""
query = Query(
field_name="embedding",
id=multiple_docs[0].id,
param=DiskAnnQueryParam(list_size=100),
)
result = collection_with_multiple_docs.query(queries=query, topk=10)
assert len(result) > 0
assert len(result) <= 10
def test_query_with_different_list_size(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with different list_size parameter values."""
query_vector = multiple_docs[0].vector("embedding")
# Test with list_size=50
query_small = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=50),
)
result_small = collection_with_multiple_docs.query(queries=query_small, topk=10)
assert len(result_small) > 0
# Test with list_size=200
query_large = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=200),
)
result_large = collection_with_multiple_docs.query(queries=query_large, topk=10)
assert len(result_large) > 0
def test_query_with_topk(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with different topk values."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
# Test topk=5
result_5 = collection_with_multiple_docs.query(queries=query, topk=5)
assert len(result_5) <= 5
# Test topk=20
result_20 = collection_with_multiple_docs.query(queries=query, topk=20)
assert len(result_20) <= 20
def test_query_with_filter(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with filter conditions."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
# Query with id filter
result = collection_with_multiple_docs.query(
queries=query, topk=10, filter="id < 50"
)
assert len(result) > 0
for doc in result:
assert doc.field("id") < 50
def test_query_with_output_fields(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with specific output fields."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result = collection_with_multiple_docs.query(
queries=query, topk=10, output_fields=["id", "name"]
)
assert len(result) > 0
first_doc = result[0]
assert "id" in first_doc.field_names()
assert "name" in first_doc.field_names()
def test_query_with_include_vector(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with vector data included in results."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result = collection_with_multiple_docs.query(
queries=query, topk=10, include_vector=True
)
assert len(result) > 0
first_doc = result[0]
assert first_doc.vector("embedding") is not None
assert len(first_doc.vector("embedding")) == 128
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionUpdate:
"""Test document update in DiskAnn collection."""
def test_update_doc_fields(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test updating document fields."""
updated_doc = Doc(
id=single_doc.id,
fields={"id": single_doc.field("id"), "name": "updated_name"},
)
result = collection_with_single_doc.update(updated_doc)
assert bool(result)
assert result.ok()
# Verify update
fetched = collection_with_single_doc.fetch(ids=[single_doc.id])
assert single_doc.id in fetched
doc = fetched[single_doc.id]
assert doc.field("name") == "updated_name"
def test_update_doc_vector(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test updating document vector."""
new_vector = [0.5 + i * 0.01 for i in range(128)]
updated_doc = Doc(
id=single_doc.id,
vectors={"embedding": new_vector},
)
result = collection_with_single_doc.update(updated_doc)
assert bool(result)
assert result.ok()
# Verify update
fetched = collection_with_single_doc.fetch(
ids=[single_doc.id],
)
assert single_doc.id in fetched
doc = fetched[single_doc.id]
assert doc.vector("embedding") is not None
embedding = doc.vector("embedding")
assert len(embedding) == 128
# Verify vector values are approximately equal (float comparison)
for i in range(128):
assert math.isclose(embedding[i], new_vector[i], rel_tol=1e-5)
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionDelete:
"""Test document deletion from DiskAnn collection."""
def test_delete_single_doc(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test deleting a single document."""
result = collection_with_single_doc.delete(single_doc.id)
assert bool(result)
assert result.ok()
stats = collection_with_single_doc.stats
assert stats.doc_count == 0
def test_delete_multiple_docs(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test deleting multiple documents."""
ids_to_delete = [doc.id for doc in multiple_docs[:10]]
result = collection_with_multiple_docs.delete(ids_to_delete)
assert len(result) == len(ids_to_delete)
for item in result:
assert item.ok()
stats = collection_with_multiple_docs.stats
assert stats.doc_count == len(multiple_docs) - len(ids_to_delete)
@pytest.mark.usefixtures("diskann_collection")
class TestDiskAnnCollectionOptimizeAndReopen:
"""Test collection optimize and reopen functionality."""
def test_optimize_close_reopen_and_query(
self,
tmp_path_factory,
diskann_collection_schema,
collection_option,
multiple_docs: list[Doc],
):
"""Test inserting 100 docs, optimize, close, reopen and query."""
# Create collection and insert 100 documents
temp_dir = tmp_path_factory.mktemp("zvec_diskann_optimize")
collection_path = temp_dir / "test_optimize_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=diskann_collection_schema,
option=collection_option,
)
assert coll is not None
assert coll.stats.doc_count == 0
# Insert 100 documents
result = coll.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
assert coll.stats.doc_count == len(multiple_docs)
# Call optimize
from zvec import OptimizeOption
coll.optimize(option=OptimizeOption())
# Verify data is still accessible after optimize
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result_before_close = coll.query(queries=query, topk=10)
assert len(result_before_close) > 0
# Close collection (destroy will close it)
collection_path_str = str(collection_path)
del coll
# Reopen collection
reopened_coll = zvec.open(path=collection_path_str, option=collection_option)
assert reopened_coll is not None
assert reopened_coll.stats.doc_count == len(multiple_docs)
# Execute query on reopened collection
query_after_reopen = Query(
field_name="embedding",
vector=query_vector,
param=DiskAnnQueryParam(list_size=100),
)
result_after_reopen = reopened_coll.query(queries=query_after_reopen, topk=10)
assert len(result_after_reopen) > 0
assert len(result_after_reopen) <= 10
# Verify query results are valid
first_doc = result_after_reopen[0]
assert first_doc is not None
assert first_doc.id is not None
assert first_doc.field("id") is not None
assert first_doc.field("name") is not None
# Cleanup
reopened_coll.destroy()
+188
View File
@@ -0,0 +1,188 @@
# 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.
"""End-to-end tests for FTS-only collections (no vector field).
The schema validation rule "must have at least one vector field" has been
lifted; these tests pin the new behavior so insert / query / delete /
optimize all work on a vector-less collection.
"""
from __future__ import annotations
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
FtsIndexParam,
OptimizeOption,
)
from zvec.model.param.query import Fts, Query
# ==================== Fixtures ====================
@pytest.fixture(scope="function")
def fts_collection(tmp_path_factory) -> Collection:
"""FTS-only collection: a STRING field for forward + an FTS-indexed STRING."""
temp_dir = tmp_path_factory.mktemp("zvec_fts_only")
collection_path = temp_dir / "fts_collection"
schema = zvec.CollectionSchema(
name="fts_only",
fields=[
FieldSchema("title", DataType.STRING, nullable=False),
FieldSchema(
"content",
DataType.STRING,
nullable=False,
index_param=FtsIndexParam(
tokenizer_name="standard",
filters=["lowercase"],
),
),
],
# vectors omitted on purpose — schema validation must accept this.
)
coll = zvec.create_and_open(
path=str(collection_path),
schema=schema,
option=CollectionOption(read_only=False, enable_mmap=True),
)
assert coll is not None
try:
yield coll
finally:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def _make_docs() -> list[Doc]:
"""5-doc corpus where 4 contain 'hello' and doc 4 is the only outlier."""
return [
Doc(id="pk_0", fields={"title": "intro", "content": "hello world"}),
Doc(id="pk_1", fields={"title": "guide", "content": "hello foo bar"}),
Doc(id="pk_2", fields={"title": "tips", "content": "hello baz"}),
Doc(id="pk_3", fields={"title": "more", "content": "hello hello"}),
Doc(id="pk_4", fields={"title": "other", "content": "nothing relevant"}),
]
def _fts_query(coll: Collection, term: str) -> list[Doc]:
"""Run a single-term FTS match query against the `content` field."""
return coll.query(
queries=Query(field_name="content", fts=Fts(match_string=term)),
topk=10,
)
# ==================== Tests ====================
class TestFtsOnlyCollectionSchema:
def test_create_and_open_without_vectors(self, fts_collection: Collection):
"""Schema with zero vector fields must be accepted by validate()."""
assert fts_collection.schema.name == "fts_only"
assert {f.name for f in fts_collection.schema.fields} == {"title", "content"}
# Empty vectors is the whole point of the test.
assert list(fts_collection.schema.vectors) == []
assert fts_collection.stats.doc_count == 0
def test_create_schema_omitting_vectors_kwarg(self):
"""Constructing CollectionSchema without `vectors=` argument is valid."""
schema = zvec.CollectionSchema(
name="bare_fts",
fields=[
FieldSchema(
"content",
DataType.STRING,
nullable=False,
index_param=FtsIndexParam(),
),
],
)
assert list(schema.vectors) == []
assert {f.name for f in schema.fields} == {"content"}
class TestFtsOnlyCollectionLifecycle:
def test_insert_and_fts_query(self, fts_collection: Collection):
"""FTS-only collection supports insert + FTS query end-to-end."""
results = fts_collection.insert(_make_docs())
assert all(r.ok() for r in results)
assert fts_collection.stats.doc_count == 5
hits = _fts_query(fts_collection, "hello")
assert len(hits) == 4
assert {doc.id for doc in hits} == {"pk_0", "pk_1", "pk_2", "pk_3"}
# Term that nothing in the surviving corpus contains.
assert _fts_query(fts_collection, "missing_term_xyz") == []
def test_delete_then_query(self, fts_collection: Collection):
"""Tombstone filter must drop deleted docs from FTS results."""
fts_collection.insert(_make_docs())
statuses = fts_collection.delete(["pk_0", "pk_4"])
assert all(s.ok() for s in statuses)
assert fts_collection.stats.doc_count == 3
hits = _fts_query(fts_collection, "hello")
assert len(hits) == 3
assert {doc.id for doc in hits} == {"pk_1", "pk_2", "pk_3"}
# pk_4's unique term is filtered out post-delete.
assert _fts_query(fts_collection, "nothing") == []
def test_optimize_rebuilds_fts(self, fts_collection: Collection):
"""Optimize with >30% deletes triggers ReduceFts; recall unchanged."""
fts_collection.insert(_make_docs())
# 40% delete ratio — above COMPACT_DELETE_RATIO_THRESHOLD=0.3, so
# build_compact_task picks the rebuild path and ReduceFts runs.
fts_collection.delete(["pk_0", "pk_4"])
before = {doc.id for doc in _fts_query(fts_collection, "hello")}
assert before == {"pk_1", "pk_2", "pk_3"}
fts_collection.optimize(option=OptimizeOption())
assert fts_collection.stats.doc_count == 3
after = {doc.id for doc in _fts_query(fts_collection, "hello")}
assert after == before
assert _fts_query(fts_collection, "nothing") == []
class TestFtsOnlyCollectionQueryValidation:
def test_vector_query_rejected(self, fts_collection: Collection):
"""Vector query on a no-vector collection must raise."""
with pytest.raises(ValueError, match="No vector field found"):
fts_collection.query(
queries=Query(field_name="content", vector=[0.1, 0.2, 0.3]),
topk=5,
)
def test_id_query_rejected(self, fts_collection: Collection):
"""ID-based query on a no-vector collection must raise."""
fts_collection.insert(_make_docs()[:1])
with pytest.raises(ValueError, match="No vector field found"):
fts_collection.query(
queries=Query(field_name="content", id="pk_0"),
topk=5,
)
@@ -0,0 +1,391 @@
# 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.
"""Tests for FTS + vector hybrid retrieval via multi-query with reranker."""
from __future__ import annotations
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
FtsIndexParam,
HnswIndexParam,
VectorSchema,
)
from zvec.extension.multi_vector_reranker import RrfReRanker, WeightedReRanker
from zvec.model.param.query import Fts, Query
DIM = 16
# ==================== Fixtures ====================
@pytest.fixture(scope="function")
def hybrid_collection(tmp_path_factory) -> Collection:
"""Collection with one vector field + one FTS field."""
temp_dir = tmp_path_factory.mktemp("zvec_hybrid")
collection_path = temp_dir / "hybrid_collection"
schema = zvec.CollectionSchema(
name="hybrid_test",
fields=[
FieldSchema("title", DataType.STRING, nullable=False),
FieldSchema(
"content",
DataType.STRING,
nullable=False,
index_param=FtsIndexParam(
tokenizer_name="standard",
filters=["lowercase"],
),
),
],
vectors=[
VectorSchema(
"embedding",
DataType.VECTOR_FP32,
dimension=DIM,
index_param=HnswIndexParam(),
),
],
)
coll = zvec.create_and_open(
path=str(collection_path),
schema=schema,
option=CollectionOption(read_only=False, enable_mmap=True),
)
assert coll is not None
try:
yield coll
finally:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def _make_docs() -> list[Doc]:
"""Corpus with both text content and vectors.
Docs 0-2: AI/ML topic, vectors clustered in one region.
Docs 3-4: retrieval topic, vectors clustered in another region.
Doc 5: unrelated topic.
"""
# AI cluster vectors
ai_vec = [1.0] * 8 + [0.0] * 8
# Retrieval cluster vectors
ret_vec = [0.0] * 8 + [1.0] * 8
# Unrelated vector
other_vec = [0.5] * 16
return [
Doc(
id="pk_0",
fields={
"title": "ML Intro",
"content": "machine learning is a branch of artificial intelligence",
},
vectors={"embedding": ai_vec},
),
Doc(
id="pk_1",
fields={
"title": "Deep Learning",
"content": "deep learning uses neural networks for pattern recognition",
},
vectors={"embedding": [0.9] * 8 + [0.1] * 8},
),
Doc(
id="pk_2",
fields={
"title": "NLP",
"content": "natural language processing handles text with artificial intelligence",
},
vectors={"embedding": [0.8] * 8 + [0.2] * 8},
),
Doc(
id="pk_3",
fields={
"title": "Search Engine",
"content": "search engine uses inverted index for text retrieval",
},
vectors={"embedding": ret_vec},
),
Doc(
id="pk_4",
fields={
"title": "Vector DB",
"content": "vector database enables similarity retrieval and search",
},
vectors={"embedding": [0.1] * 8 + [0.9] * 8},
),
Doc(
id="pk_5",
fields={
"title": "Cooking",
"content": "baking bread requires flour water yeast and salt",
},
vectors={"embedding": other_vec},
),
]
@pytest.fixture(scope="function")
def hybrid_collection_with_docs(hybrid_collection: Collection) -> Collection:
"""Hybrid collection pre-populated with test documents."""
results = hybrid_collection.insert(_make_docs())
assert all(r.ok() for r in results)
return hybrid_collection
# ==================== Tests ====================
class TestFtsVectorHybridQuery:
"""Test FTS + vector hybrid retrieval using multi-query with RRF reranker."""
def test_hybrid_fts_and_vector_basic(self, hybrid_collection_with_docs: Collection):
"""FTS + vector multi-query with RRF reranker returns results."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval")),
Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 5
# Results should have scores
for doc in result:
assert doc.score > 0
def test_hybrid_fts_and_vector_ranking(
self, hybrid_collection_with_docs: Collection
):
"""Docs relevant in both FTS and vector should rank higher."""
reranker = RrfReRanker(rank_constant=60)
# FTS: "retrieval search" matches pk_3, pk_4
# Vector: ret_vec cluster matches pk_3, pk_4
# Both signals agree: pk_3 and pk_4 should rank top
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval search")),
Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8),
],
topk=5,
reranker=reranker,
)
top_ids = {doc.id for doc in result[:3]}
assert "pk_3" in top_ids or "pk_4" in top_ids
def test_hybrid_scores_descending(self, hybrid_collection_with_docs: Collection):
"""Hybrid query results must be sorted by score descending."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="intelligence")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=6,
reranker=reranker,
)
assert len(result) >= 2
scores = [doc.score for doc in result]
assert scores == sorted(scores, reverse=True)
def test_hybrid_with_filter(self, hybrid_collection_with_docs: Collection):
"""Hybrid query respects SQL filter."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="learning")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=10,
reranker=reranker,
filter="title like '%Learning%'",
)
for doc in result:
assert "Learning" in doc.fields["title"]
def test_hybrid_fts_no_match_still_returns_vector_results(
self, hybrid_collection_with_docs: Collection
):
"""When FTS matches nothing, vector results still appear."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(
field_name="content",
fts=Fts(match_string="nonexistent_term_xyz"),
),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=5,
reranker=reranker,
)
# Vector query alone should still produce results
assert len(result) > 0
def test_hybrid_query_string_syntax(self, hybrid_collection_with_docs: Collection):
"""Hybrid query works with FTS query_string (advanced syntax)."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(
field_name="content",
fts=Fts(query_string="artificial AND intelligence"),
),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
# pk_0 and pk_2 contain "artificial intelligence"
hit_ids = {doc.id for doc in result}
assert "pk_0" in hit_ids or "pk_2" in hit_ids
class TestFtsVectorHybridValidation:
"""Test validation rules for FTS + vector hybrid queries."""
def test_hybrid_requires_reranker(self, hybrid_collection_with_docs: Collection):
"""Multi-query with FTS + vector without reranker should raise."""
with pytest.raises(ValueError, match="[Rr]eranker"):
hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="learning")),
Query(field_name="embedding", vector=[1.0] * DIM),
],
topk=5,
)
def test_duplicate_field_name_allowed(
self, hybrid_collection_with_docs: Collection
):
"""Multi-query with duplicate field names is allowed and returns results."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="learning")),
Query(field_name="content", fts=Fts(match_string="intelligence")),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 5
def test_multiple_vectors_allowed(self, hybrid_collection_with_docs: Collection):
"""Two vector queries on the same field are allowed with a reranker."""
reranker = RrfReRanker(rank_constant=60)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="embedding", vector=[1.0] * DIM),
Query(field_name="embedding", vector=[0.5] * DIM),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 5
class TestFtsVectorHybridWeightedReranker:
"""Test FTS + vector hybrid retrieval using WeightedReranker."""
def test_weighted_reranker_fts_and_vector(
self, hybrid_collection_with_docs: Collection
):
"""WeightedReranker correctly normalizes FTS scores alongside vector scores."""
weights = [0.5, 0.5]
reranker = WeightedReRanker(weights=weights)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval search")),
Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 5
for doc in result:
assert doc.score > 0
def test_weighted_reranker_scores_descending(
self, hybrid_collection_with_docs: Collection
):
"""WeightedReranker hybrid results are sorted by score descending."""
weights = [0.4, 0.6]
reranker = WeightedReRanker(weights=weights)
result = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="intelligence")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=6,
reranker=reranker,
)
assert len(result) >= 2
scores = [doc.score for doc in result]
assert scores == sorted(scores, reverse=True)
def test_weighted_reranker_fts_weight_influence(
self, hybrid_collection_with_docs: Collection
):
"""Higher FTS weight should boost FTS-relevant docs in ranking."""
# High FTS weight: FTS signal dominates
weights_fts_heavy = [0.9, 0.1]
reranker_fts = WeightedReRanker(weights=weights_fts_heavy)
result_fts = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=5,
reranker=reranker_fts,
)
# High vector weight: vector signal dominates
weights_vec_heavy = [0.1, 0.9]
reranker_vec = WeightedReRanker(weights=weights_vec_heavy)
result_vec = hybrid_collection_with_docs.query(
queries=[
Query(field_name="content", fts=Fts(match_string="retrieval")),
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
],
topk=5,
reranker=reranker_vec,
)
# Both should return results
assert len(result_fts) > 0
assert len(result_vec) > 0
# With FTS-heavy weight, FTS-relevant docs (pk_3, pk_4) should rank higher
fts_top = [doc.id for doc in result_fts[:2]]
vec_top = [doc.id for doc in result_vec[:2]]
# The rankings should differ due to weight difference
assert fts_top != vec_top or len(result_fts) == len(result_vec) == 1
+574
View File
@@ -0,0 +1,574 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import platform
import sys
import pytest
import math
import zvec
pytestmark = pytest.mark.skipif(
not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")),
reason="HNSW RaBitQ only supported on Linux x86_64",
)
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
HnswRabitqIndexParam,
HnswRabitqQueryParam,
MetricType,
VectorSchema,
Query,
)
# ==================== Fixtures ====================
@pytest.fixture(scope="session")
def hnsw_rabitq_collection_schema():
"""Create a collection schema with HNSW RaBitQ index."""
return zvec.CollectionSchema(
name="test_hnsw_rabitq_collection",
fields=[
FieldSchema("id", DataType.INT64, nullable=False),
FieldSchema("name", DataType.STRING, nullable=False),
],
vectors=[
VectorSchema(
"embedding",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswRabitqIndexParam(
metric_type=MetricType.L2,
m=16,
ef_construction=200,
total_bits=7,
num_clusters=64,
),
),
],
)
@pytest.fixture(scope="session")
def collection_option():
"""Create collection options."""
return CollectionOption(read_only=False, enable_mmap=True)
@pytest.fixture
def single_doc():
"""Create a single document for testing."""
return Doc(
id="0",
fields={"id": 0, "name": "test_doc_0"},
vectors={"embedding": [0.1 + i * 0.01 for i in range(128)]},
)
@pytest.fixture
def multiple_docs():
"""Create multiple documents for testing."""
return [
Doc(
id=f"{i}",
fields={"id": i, "name": f"test_doc_{i}"},
vectors={"embedding": [i * 0.1 + j * 0.01 for j in range(128)]},
)
for i in range(1, 101)
]
@pytest.fixture(scope="function")
def hnsw_rabitq_collection(
tmp_path_factory, hnsw_rabitq_collection_schema, collection_option
) -> Collection:
"""
Function-scoped fixture: creates and opens a collection with HNSW RaBitQ index.
"""
temp_dir = tmp_path_factory.mktemp("zvec_hnsw_rabitq")
collection_path = temp_dir / "test_hnsw_rabitq_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=hnsw_rabitq_collection_schema,
option=collection_option,
)
assert coll is not None, "Failed to create and open HNSW RaBitQ collection"
assert coll.path == str(collection_path)
assert coll.schema.name == hnsw_rabitq_collection_schema.name
try:
yield coll
finally:
if hasattr(coll, "destroy") and coll is not None:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
@pytest.fixture
def collection_with_single_doc(
hnsw_rabitq_collection: Collection, single_doc: Doc
) -> Collection:
"""Setup: insert single doc into collection."""
assert hnsw_rabitq_collection.stats.doc_count == 0
result = hnsw_rabitq_collection.insert(single_doc)
assert bool(result)
assert result.ok()
assert hnsw_rabitq_collection.stats.doc_count == 1
yield hnsw_rabitq_collection
# Teardown: delete single doc
hnsw_rabitq_collection.delete(single_doc.id)
assert hnsw_rabitq_collection.stats.doc_count == 0
@pytest.fixture
def collection_with_multiple_docs(
hnsw_rabitq_collection: Collection, multiple_docs: list[Doc]
) -> Collection:
"""Setup: insert multiple docs into collection."""
assert hnsw_rabitq_collection.stats.doc_count == 0
result = hnsw_rabitq_collection.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
assert hnsw_rabitq_collection.stats.doc_count == len(multiple_docs)
yield hnsw_rabitq_collection
# Teardown: delete multiple docs
hnsw_rabitq_collection.delete([doc.id for doc in multiple_docs])
# ==================== Tests ====================
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionCreation:
"""Test HNSW RaBitQ collection creation and schema validation."""
def test_collection_creation(
self, hnsw_rabitq_collection: Collection, hnsw_rabitq_collection_schema
):
"""Test that collection is created with correct schema."""
assert hnsw_rabitq_collection is not None
assert hnsw_rabitq_collection.schema.name == hnsw_rabitq_collection_schema.name
assert len(hnsw_rabitq_collection.schema.fields) == len(
hnsw_rabitq_collection_schema.fields
)
assert len(hnsw_rabitq_collection.schema.vectors) == len(
hnsw_rabitq_collection_schema.vectors
)
def test_vector_schema_validation(self, hnsw_rabitq_collection: Collection):
"""Test that vector schema has correct HNSW RaBitQ configuration."""
vector_schema = hnsw_rabitq_collection.schema.vector("embedding")
assert vector_schema is not None
assert vector_schema.name == "embedding"
assert vector_schema.data_type == DataType.VECTOR_FP32
assert vector_schema.dimension == 128
index_param = vector_schema.index_param
assert index_param is not None
assert index_param.metric_type == MetricType.L2
assert index_param.m == 16
assert index_param.ef_construction == 200
assert index_param.total_bits == 7
assert index_param.num_clusters == 64
def test_collection_stats(self, hnsw_rabitq_collection: Collection):
"""Test initial collection statistics."""
stats = hnsw_rabitq_collection.stats
assert stats is not None
assert stats.doc_count == 0
assert len(stats.index_completeness) == 1
assert stats.index_completeness["embedding"] == 1
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionInsert:
"""Test document insertion into HNSW RaBitQ collection."""
def test_insert_single_doc(
self, hnsw_rabitq_collection: Collection, single_doc: Doc
):
"""Test inserting a single document."""
result = hnsw_rabitq_collection.insert(single_doc)
assert bool(result)
assert result.ok()
stats = hnsw_rabitq_collection.stats
assert stats is not None
assert stats.doc_count == 1
def test_insert_multiple_docs(
self, hnsw_rabitq_collection: Collection, multiple_docs: list[Doc]
):
"""Test inserting multiple documents."""
result = hnsw_rabitq_collection.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
stats = hnsw_rabitq_collection.stats
assert stats is not None
assert stats.doc_count == len(multiple_docs)
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionFetch:
"""Test document fetching from HNSW RaBitQ collection."""
def test_fetch_single_doc(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test fetching a single document by ID."""
result = collection_with_single_doc.fetch(ids=[single_doc.id])
assert bool(result)
assert single_doc.id in result.keys()
doc = result[single_doc.id]
assert doc is not None
assert doc.id == single_doc.id
assert doc.field("id") == single_doc.field("id")
assert doc.field("name") == single_doc.field("name")
def test_fetch_multiple_docs(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test fetching multiple documents by IDs."""
ids = [doc.id for doc in multiple_docs[:10]]
result = collection_with_multiple_docs.fetch(ids=ids)
assert bool(result)
assert len(result) == len(ids)
for doc_id in ids:
assert doc_id in result
doc = result[doc_id]
assert doc is not None
assert doc.id == doc_id
def test_fetch_nonexistent_doc(self, collection_with_single_doc: Collection):
"""Test fetching a non-existent document."""
result = collection_with_single_doc.fetch(ids=["nonexistent_id"])
assert len(result) == 0
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionQuery:
"""Test vector search queries on HNSW RaBitQ collection."""
def test_query_by_vector(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying by vector with HNSW RaBitQ index."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result = collection_with_multiple_docs.query(queries=query, topk=10)
assert len(result) > 0
assert len(result) <= 10
# First result should be the query document itself (or very close)
first_doc = result[0]
assert first_doc is not None
assert first_doc.id is not None
def test_query_by_id(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying by document ID with HNSW RaBitQ index."""
query = Query(
field_name="embedding",
id=multiple_docs[0].id,
param=HnswRabitqQueryParam(ef=300),
)
result = collection_with_multiple_docs.query(queries=query, topk=10)
assert len(result) > 0
assert len(result) <= 10
def test_query_with_different_ef_values(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with different ef parameter values."""
query_vector = multiple_docs[0].vector("embedding")
# Test with ef=100
query_100 = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=100),
)
result_100 = collection_with_multiple_docs.query(queries=query_100, topk=10)
assert len(result_100) > 0
# Test with ef=500
query_500 = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=500),
)
result_500 = collection_with_multiple_docs.query(queries=query_500, topk=10)
assert len(result_500) > 0
def test_query_with_topk(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with different topk values."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
# Test topk=5
result_5 = collection_with_multiple_docs.query(queries=query, topk=5)
assert len(result_5) <= 5
# Test topk=20
result_20 = collection_with_multiple_docs.query(queries=query, topk=20)
assert len(result_20) <= 20
def test_query_with_filter(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with filter conditions."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
# Query with id filter
result = collection_with_multiple_docs.query(
queries=query, topk=10, filter="id < 50"
)
assert len(result) > 0
for doc in result:
assert doc.field("id") < 50
def test_query_with_output_fields(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with specific output fields."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result = collection_with_multiple_docs.query(
queries=query, topk=10, output_fields=["id", "name"]
)
assert len(result) > 0
first_doc = result[0]
assert "id" in first_doc.field_names()
assert "name" in first_doc.field_names()
def test_query_with_include_vector(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test querying with vector data included in results."""
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result = collection_with_multiple_docs.query(
queries=query, topk=10, include_vector=True
)
assert len(result) > 0
first_doc = result[0]
assert first_doc.vector("embedding") is not None
assert len(first_doc.vector("embedding")) == 128
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionUpdate:
"""Test document update in HNSW RaBitQ collection."""
def test_update_doc_fields(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test updating document fields."""
updated_doc = Doc(
id=single_doc.id,
fields={"id": single_doc.field("id"), "name": "updated_name"},
)
result = collection_with_single_doc.update(updated_doc)
assert bool(result)
assert result.ok()
# Verify update
fetched = collection_with_single_doc.fetch(ids=[single_doc.id])
assert single_doc.id in fetched
doc = fetched[single_doc.id]
assert doc.field("name") == "updated_name"
def test_update_doc_vector(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test updating document vector."""
new_vector = [0.5 + i * 0.01 for i in range(128)]
updated_doc = Doc(
id=single_doc.id,
vectors={"embedding": new_vector},
)
result = collection_with_single_doc.update(updated_doc)
assert bool(result)
assert result.ok()
# Verify update
fetched = collection_with_single_doc.fetch(
ids=[single_doc.id],
)
assert single_doc.id in fetched
doc = fetched[single_doc.id]
assert doc.vector("embedding") is not None
embedding = doc.vector("embedding")
assert len(embedding) == 128
# Verify vector values are approximately equal (float comparison)
for i in range(128):
assert math.isclose(embedding[i], new_vector[i], rel_tol=1e-5)
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionDelete:
"""Test document deletion from HNSW RaBitQ collection."""
def test_delete_single_doc(
self, collection_with_single_doc: Collection, single_doc: Doc
):
"""Test deleting a single document."""
result = collection_with_single_doc.delete(single_doc.id)
assert bool(result)
assert result.ok()
stats = collection_with_single_doc.stats
assert stats.doc_count == 0
def test_delete_multiple_docs(
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
):
"""Test deleting multiple documents."""
ids_to_delete = [doc.id for doc in multiple_docs[:10]]
result = collection_with_multiple_docs.delete(ids_to_delete)
assert len(result) == len(ids_to_delete)
for item in result:
assert item.ok()
stats = collection_with_multiple_docs.stats
assert stats.doc_count == len(multiple_docs) - len(ids_to_delete)
@pytest.mark.usefixtures("hnsw_rabitq_collection")
class TestHnswRabitqCollectionOptimizeAndReopen:
"""Test collection optimize and reopen functionality."""
def test_optimize_close_reopen_and_query(
self,
tmp_path_factory,
hnsw_rabitq_collection_schema,
collection_option,
multiple_docs: list[Doc],
):
"""Test inserting 100 docs, optimize, close, reopen and query."""
# Create collection and insert 100 documents
temp_dir = tmp_path_factory.mktemp("zvec_hnsw_rabitq_optimize")
collection_path = temp_dir / "test_optimize_collection"
coll = zvec.create_and_open(
path=str(collection_path),
schema=hnsw_rabitq_collection_schema,
option=collection_option,
)
assert coll is not None
assert coll.stats.doc_count == 0
# Insert 100 documents
result = coll.insert(multiple_docs)
assert len(result) == len(multiple_docs)
for item in result:
assert item.ok()
assert coll.stats.doc_count == len(multiple_docs)
# Call optimize
from zvec import OptimizeOption
coll.optimize(option=OptimizeOption())
# Verify data is still accessible after optimize
query_vector = multiple_docs[0].vector("embedding")
query = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result_before_close = coll.query(query, topk=10)
assert len(result_before_close) > 0
# Close collection (destroy will close it)
collection_path_str = str(collection_path)
del coll
# Reopen collection
reopened_coll = zvec.open(path=collection_path_str, option=collection_option)
assert reopened_coll is not None
assert reopened_coll.stats.doc_count == len(multiple_docs)
# Execute query on reopened collection
query_after_reopen = Query(
field_name="embedding",
vector=query_vector,
param=HnswRabitqQueryParam(ef=300),
)
result_after_reopen = reopened_coll.query(query_after_reopen, topk=10)
assert len(result_after_reopen) > 0
assert len(result_after_reopen) <= 10
# Verify query results are valid
first_doc = result_after_reopen[0]
assert first_doc is not None
assert first_doc.id is not None
assert first_doc.field("id") is not None
assert first_doc.field("name") is not None
# Cleanup
reopened_coll.destroy()
+584
View File
@@ -0,0 +1,584 @@
from __future__ import annotations
import math
import pytest
from zvec._zvec import _Doc
from zvec.model.convert import convert_to_py_doc, convert_to_cpp_doc
from zvec import Doc, CollectionSchema, DataType, FieldSchema, VectorSchema
# ----------------------------
# Convert Cpp Doc Test Case
# ----------------------------
class TestConvertCppDoc:
def test_default(self):
doc = Doc(id="1")
schema = CollectionSchema(
name="test_collection",
fields=FieldSchema("name", DataType.STRING),
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
def test_with_field_notin_schema(self):
doc = Doc(id="1", fields={"name": "Tom"})
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("id", DataType.UINT64),
FieldSchema("salary", DataType.UINT32),
FieldSchema("age", DataType.INT32),
FieldSchema("create_at", DataType.INT64),
FieldSchema("author", DataType.STRING),
FieldSchema("weight", DataType.FLOAT),
],
)
with pytest.raises(ValueError):
convert_to_cpp_doc(doc, collection_schema=schema)
def test_with_scalar_fields(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("id", DataType.UINT64),
FieldSchema("salary", DataType.UINT32),
FieldSchema("age", DataType.INT32),
FieldSchema("create_at", DataType.INT64),
FieldSchema("author", DataType.STRING),
FieldSchema("weight", DataType.FLOAT),
FieldSchema("bmi", DataType.DOUBLE),
FieldSchema("is_male", DataType.BOOL),
],
)
doc = Doc(
id="1",
fields={
"id": 1,
"salary": 1000,
"age": 18,
"create_at": 1640995200,
"bmi": 80.0 / 200.0,
"author": "Tom",
"weight": 80.0,
"is_male": True,
},
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
assert cpp_doc.get_any("id", DataType.UINT64) == 1
assert cpp_doc.get_any("salary", DataType.UINT32) == 1000
assert cpp_doc.get_any("age", DataType.INT32) == 18
assert cpp_doc.get_any("create_at", DataType.INT64) == 1640995200
assert cpp_doc.get_any("author", DataType.STRING) == "Tom"
assert math.isclose(
cpp_doc.get_any("weight", DataType.FLOAT), 80.0, rel_tol=1e-6
)
assert math.isclose(
cpp_doc.get_any("bmi", DataType.DOUBLE), 80.0 / 200.0, rel_tol=1e-6
)
assert cpp_doc.get_any("is_male", DataType.BOOL) == True
def test_with_array_fields(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("tags", DataType.ARRAY_STRING),
FieldSchema("ids", DataType.ARRAY_UINT64),
FieldSchema("marks", DataType.ARRAY_UINT32),
FieldSchema("x", DataType.ARRAY_INT32),
FieldSchema("y", DataType.ARRAY_INT64),
FieldSchema("scores", DataType.ARRAY_FLOAT),
FieldSchema("ratios", DataType.ARRAY_DOUBLE),
FieldSchema("results", DataType.ARRAY_BOOL),
],
)
doc = Doc(
id="1",
fields={
"tags": ["tag1", "tag2", "tag3"],
"ids": [111111111111, 222222222222, 333333333333],
"marks": [100, 200, 300],
"x": [1, 2, 3],
"y": [100, 200, 300],
"scores": [1.1, 2.2, 3.3],
"ratios": [0.1, 0.2, 0.3],
"results": [True, False, True],
},
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
assert cpp_doc.get_any("tags", DataType.ARRAY_STRING) == doc.field("tags")
assert cpp_doc.get_any("ids", DataType.ARRAY_UINT64) == doc.field("ids")
assert cpp_doc.get_any("marks", DataType.ARRAY_UINT32) == doc.field("marks")
assert cpp_doc.get_any("x", DataType.ARRAY_INT32) == doc.field("x")
assert cpp_doc.get_any("y", DataType.ARRAY_INT64) == doc.field("y")
scores = cpp_doc.get_any("scores", DataType.ARRAY_FLOAT)
for i in range(len(doc.field("scores"))):
assert math.isclose(scores[i], doc.field("scores")[i], rel_tol=1e-1)
ratios = cpp_doc.get_any("ratios", DataType.ARRAY_DOUBLE)
for i in range(len(doc.field("ratios"))):
assert math.isclose(ratios[i], doc.field("ratios")[i], rel_tol=1e-1)
results = cpp_doc.get_any("results", DataType.ARRAY_BOOL)
for i in range(len(doc.field("results"))):
assert results[i] == doc.field("results")[i]
def test_with_dense_vector_fields(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="embedding",
data_type=DataType.VECTOR_FP16,
dimension=4,
),
VectorSchema(
name="image",
data_type=DataType.VECTOR_FP32,
dimension=8,
),
VectorSchema(
name="text",
data_type=DataType.VECTOR_INT8,
dimension=32,
),
],
)
doc = Doc(
id="1",
vectors={
"embedding": [1.1] * 4,
"image": [2.2] * 8,
"text": [4] * 32,
},
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
embedding_vector = cpp_doc.get_any("embedding", DataType.VECTOR_FP16)
assert len(embedding_vector) == 4
for i in range(4):
assert math.isclose(
embedding_vector[i], doc.vector("embedding")[i], rel_tol=1e-1
)
image_vector = cpp_doc.get_any("image", DataType.VECTOR_FP32)
assert len(image_vector) == 8
for i in range(8):
assert math.isclose(image_vector[i], doc.vector("image")[i], rel_tol=1e-1)
text_vector = cpp_doc.get_any("text", DataType.VECTOR_INT8)
assert len(text_vector) == 32
for i in range(32):
assert text_vector[i] == doc.vectors["text"][i]
def test_with_sparse_vector_fields(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="author",
data_type=DataType.SPARSE_VECTOR_FP32,
),
VectorSchema(
name="content",
data_type=DataType.SPARSE_VECTOR_FP16,
),
],
)
doc = Doc(
id="1",
vectors={
"author": {1: 1.1, 2: 2.2, 3: 3.3},
"content": {4: 4.4, 5: 5.5, 6: 6.6},
},
)
cpp_doc = convert_to_cpp_doc(doc, collection_schema=schema)
assert cpp_doc is not None
assert cpp_doc.pk() == doc.id
author_vector = cpp_doc.get_any("author", DataType.SPARSE_VECTOR_FP32)
assert isinstance(author_vector, dict)
for key, value in doc.vector("author").items():
assert math.isclose(author_vector[key], value, rel_tol=1e-1)
content_vector = cpp_doc.get_any("content", DataType.SPARSE_VECTOR_FP16)
assert isinstance(content_vector, dict)
for key, value in doc.vector("content").items():
assert math.isclose(content_vector[key], value, rel_tol=1e-1)
def test_with_scalar_fields_error_datatype(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("id", DataType.UINT64),
FieldSchema("salary", DataType.UINT32),
FieldSchema("age", DataType.INT32),
FieldSchema("create_at", DataType.INT64),
FieldSchema("author", DataType.STRING),
FieldSchema("weight", DataType.FLOAT),
FieldSchema("bmi", DataType.DOUBLE),
FieldSchema("is_male", DataType.BOOL),
],
)
doc = Doc(
id="1",
fields={
"id": "1",
},
)
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"salary": "1000"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"age": "18"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"create_at": "2021-01-01"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"author": 1})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"weight": "80.5"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"bmi": "25.0"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"is_male": "true"})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
def test_with_array_fields_error_datatype(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("tags", DataType.ARRAY_STRING),
FieldSchema("ids", DataType.ARRAY_UINT64),
FieldSchema("marks", DataType.ARRAY_UINT32),
FieldSchema("x", DataType.ARRAY_INT32),
FieldSchema("y", DataType.ARRAY_INT64),
FieldSchema("scores", DataType.ARRAY_FLOAT),
FieldSchema("ratios", DataType.ARRAY_DOUBLE),
FieldSchema("results", DataType.ARRAY_BOOL),
],
)
doc = Doc(id="1", fields={"tags": [1, 2, 3]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"ids": ["1", "2", "3"]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"marks": [1.1, 2.2, 3.3]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"x": [1.1, 2.2, 3.3]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"y": [1.1, 2.2, 3.3]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"scores": ["1", "2", "3"]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"ratios": ["1", "2", "3"]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", fields={"results": ["1", "2", "3"]})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
def test_with_vector_fields_error_datatype(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="embedding",
data_type=DataType.VECTOR_FP16,
dimension=4,
),
VectorSchema(
name="image",
data_type=DataType.VECTOR_FP32,
dimension=8,
),
VectorSchema(
name="text",
data_type=DataType.VECTOR_INT8,
dimension=32,
),
],
)
doc = Doc(id="1", vectors={"image": ["1.1"] * 4})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", vectors={"text": ["1"] * 4})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(id="1", vectors={"embedding": ["1"] * 4})
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
def test_with_sparse_vector_error_datatype(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="author",
data_type=DataType.SPARSE_VECTOR_FP32,
),
VectorSchema(
name="content",
data_type=DataType.SPARSE_VECTOR_FP16,
),
],
)
doc = Doc(
id="1",
vectors={
"author": {"1": 1.1, "2": 2.2, "3": 3.3},
},
)
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(
id="1",
vectors={
"content": {"1": 1.1, "2": 2.2, "3": 3.3},
},
)
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
doc = Doc(
id="1",
vectors={
"author": {1: "1", 2: "2", 3: "3"},
},
)
with pytest.raises(TypeError):
convert_to_cpp_doc(doc, collection_schema=schema)
# ----------------------------
# Convert Py Doc Test Case
# ----------------------------
class TestConvertPyDoc:
def test_default(self):
doc = _Doc()
doc.set_pk("1")
doc.set_score(1.0)
schema = CollectionSchema(
name="test_collection",
fields=FieldSchema("name", DataType.STRING),
)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.id == "1"
assert py_doc.score == 1.0
def test_with_scalar_fields(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("id", DataType.UINT64),
FieldSchema("salary", DataType.UINT32),
FieldSchema("age", DataType.INT32),
FieldSchema("create_at", DataType.INT64),
FieldSchema("author", DataType.STRING),
FieldSchema("weight", DataType.FLOAT),
FieldSchema("bmi", DataType.DOUBLE),
FieldSchema("is_male", DataType.BOOL),
],
)
doc = _Doc()
doc.set_pk("1")
doc.set_any("id", schema.field("id")._get_object(), 1111111111111111)
doc.set_any("salary", schema.field("salary")._get_object(), 1000)
doc.set_any("age", schema.field("age")._get_object(), 18)
doc.set_any("create_at", schema.field("create_at")._get_object(), 1640995200)
doc.set_any("author", schema.field("author")._get_object(), "Tom")
doc.set_any("weight", schema.field("weight")._get_object(), 80.0)
doc.set_any("bmi", schema.field("bmi")._get_object(), 80.0 / 200.0)
doc.set_any("is_male", schema.field("is_male")._get_object(), True)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.id == "1"
assert py_doc.field("id") == 1111111111111111
assert py_doc.field("salary") == 1000
assert py_doc.field("age") == 18
assert py_doc.field("create_at") == 1640995200
assert py_doc.field("author") == "Tom"
assert py_doc.field("weight") == 80.0
assert py_doc.field("bmi") == 80.0 / 200.0
assert py_doc.field("is_male") == True
def test_with_array_fields(self):
schema = CollectionSchema(
name="test_collection",
fields=[
FieldSchema("tags", DataType.ARRAY_STRING),
FieldSchema("ids", DataType.ARRAY_UINT64),
FieldSchema("marks", DataType.ARRAY_UINT32),
FieldSchema("x", DataType.ARRAY_INT32),
FieldSchema("y", DataType.ARRAY_INT64),
FieldSchema("scores", DataType.ARRAY_FLOAT),
FieldSchema("ratios", DataType.ARRAY_DOUBLE),
FieldSchema("results", DataType.ARRAY_BOOL),
],
)
doc = _Doc()
doc.set_pk("1")
doc.set_any(
"tags", schema.field("tags")._get_object(), ["tag1", "tag2", "tag3"]
)
doc.set_any(
"ids",
schema.field("ids")._get_object(),
[111111111111, 222222222222, 3333333333333],
)
doc.set_any("marks", schema.field("marks")._get_object(), [1000, 2000, 3000])
doc.set_any("x", schema.field("x")._get_object(), [1, 2, 3])
doc.set_any("y", schema.field("y")._get_object(), [100, 200, 300])
doc.set_any("scores", schema.field("scores")._get_object(), [0.1, 0.2, 0.3])
doc.set_any("ratios", schema.field("ratios")._get_object(), [0.1, 0.2, 0.3])
doc.set_any(
"results", schema.field("results")._get_object(), [True, False, True]
)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.field("tags") == ["tag1", "tag2", "tag3"]
assert py_doc.field("ids") == [111111111111, 222222222222, 3333333333333]
assert py_doc.field("marks") == [1000, 2000, 3000]
assert py_doc.field("x") == [1, 2, 3]
assert py_doc.field("y") == [100, 200, 300]
scores = doc.get_any("scores", DataType.ARRAY_FLOAT)
for i in range(len(scores)):
assert math.isclose(scores[i], py_doc.field("scores")[i], rel_tol=1e-1)
ratios = doc.get_any("ratios", DataType.ARRAY_DOUBLE)
for i in range(len(ratios)):
assert math.isclose(ratios[i], py_doc.field("ratios")[i], rel_tol=1e-1)
results = doc.get_any("results", DataType.ARRAY_BOOL)
for i in range(len(results)):
assert results[i] == py_doc.field("results")[i]
def test_with_dense_vector_fields(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="embedding",
data_type=DataType.VECTOR_FP16,
dimension=4,
),
VectorSchema(
name="image",
data_type=DataType.VECTOR_FP32,
dimension=8,
),
VectorSchema(
name="text",
data_type=DataType.VECTOR_INT8,
dimension=32,
),
],
)
doc = _Doc()
doc.set_pk("1")
doc.set_any("embedding", schema.vector("embedding")._get_object(), [1.1] * 4)
doc.set_any("image", schema.vector("image")._get_object(), [2.2] * 8)
doc.set_any("text", schema.vector("text")._get_object(), [4] * 32)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.id == "1"
embedding_vector = py_doc.vector("embedding")
assert len(embedding_vector) == 4
for i in range(4):
assert math.isclose(
py_doc.vector("embedding")[i], embedding_vector[i], rel_tol=1e-1
)
image_vector = py_doc.vector("image")
assert len(image_vector) == 8
for i in range(8):
assert math.isclose(
py_doc.vector("image")[i], image_vector[i], rel_tol=1e-1
)
text_vector = py_doc.vector("text")
assert len(text_vector) == 32
for i in range(32):
assert py_doc.vector("text")[i] == text_vector[i]
def test_with_sparse_vector_fields(self):
schema = CollectionSchema(
name="test_collection",
vectors=[
VectorSchema(
name="author",
data_type=DataType.SPARSE_VECTOR_FP32,
),
VectorSchema(
name="content",
data_type=DataType.SPARSE_VECTOR_FP16,
),
],
)
doc = _Doc()
doc.set_pk("1")
doc.set_any(
"author", schema.vector("author")._get_object(), {1: 1.1, 2: 2.2, 3: 3.3}
)
doc.set_any(
"content", schema.vector("content")._get_object(), {4: 4.4, 5: 5.5, 6: 6.6}
)
py_doc = convert_to_py_doc(doc, schema)
assert py_doc.id == "1"
author_vector = py_doc.vector("author")
assert isinstance(author_vector, dict)
for key, value in doc.get_any("author", DataType.SPARSE_VECTOR_FP32).items():
assert math.isclose(author_vector[key], value, rel_tol=1e-1)
content_vector = py_doc.vector("content")
assert isinstance(content_vector, dict)
for key, value in doc.get_any("content", DataType.SPARSE_VECTOR_FP16).items():
assert math.isclose(content_vector[key], value, rel_tol=1e-1)
+269
View File
@@ -0,0 +1,269 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import math
import pytest
from zvec._zvec import _Doc
from zvec import FieldSchema, VectorSchema, Doc, DataType
# ----------------------------
# PyDoc Test Case
# ----------------------------
class TestPyDoc:
def test_default(self):
Doc(id="1")
def test_with_single_vector(self):
doc = Doc(id="1", vectors={"dense": [1, 2, 3]})
assert doc is not None
assert doc.id == "1"
assert doc.vector("dense") == [1, 2, 3]
def test_with_hybrid_vectors(self):
doc = Doc(
id="1", vectors={"dense": [1, 2, 3], "sparse": {1: 1.0, 2: 2.0, 3: 3.0}}
)
assert doc is not None
assert doc.id == "1"
assert doc.vector("dense") == [1, 2, 3]
assert doc.vector("sparse") == {1: 1.0, 2: 2.0, 3: 3.0}
def test_with_multi_vectors(self):
doc = Doc(
id="1",
vectors={
"image": [1, 2, 3],
"description": [4, 5, 6],
"keys": {1: 1.0, 2: 2.0, 3: 3.0},
},
fields={"author": "Tom", "age": 19, "is_male": True, "weight": 60.5},
)
assert doc is not None
assert doc.id == "1"
assert doc.vector("image") == [1, 2, 3]
assert doc.vector("description") == [4, 5, 6]
assert doc.vector("keys") == {1: 1.0, 2: 2.0, 3: 3.0}
assert doc.field("author") == "Tom"
assert doc.field("age") == 19
assert doc.field("is_male") == True
assert doc.field("weight") == 60.5
def test_with_numpy_array(self):
import numpy as np
doc = Doc._from_tuple(
(
"1",
0.0,
None,
{
"image": np.array([1, 2, 3]),
"description": np.random.random(512),
"keys": {1: 1.0, 2: 2.0, 3: 3.0},
},
)
)
assert doc is not None
assert doc.id == "1"
assert doc.vector("image") == [1, 2, 3]
assert doc.vector("keys") == {1: 1.0, 2: 2.0, 3: 3.0}
# ----------------------------
# CppDoc Test Case
# ----------------------------
class TestCppDoc:
def test_default(self):
doc = _Doc()
assert doc is not None
def test_doc_set_pk(self):
doc = _Doc()
doc.set_pk("1")
assert doc.pk() == "1"
def test_doc_set_score(self):
doc = _Doc()
doc.set_score(0.9)
assert math.isclose(doc.score(), 0.9, rel_tol=1e-6)
def test_doc_get_null_field(self):
doc = _Doc()
schema = FieldSchema("author", DataType.STRING, nullable=True)
doc.set_any("author", schema._get_object(), None)
assert doc.has_field("author")
assert doc.get_any("author", schema.data_type) is None
def test_doc_get_set_has_null_field(self):
doc = _Doc()
schema = FieldSchema("author", DataType.STRING, nullable=False)
with pytest.raises(ValueError):
doc.set_any("author", schema._get_object(), None)
def test_doc_get_set_has_string_field(self):
doc = _Doc()
schema = FieldSchema("author", DataType.STRING)
doc.set_any("author", schema._get_object(), "Tom")
assert doc.has_field("author")
assert doc.get_any("author", DataType.STRING) == "Tom"
def test_doc_get_set_has_bool_field(self):
doc = _Doc()
schema = FieldSchema("is_male", DataType.BOOL)
doc.set_any("is_male", schema._get_object(), True)
assert doc.has_field("is_male")
assert doc.get_any("is_male", DataType.BOOL) == True
def test_doc_get_set_has_int32_field(self):
doc = _Doc()
schema = FieldSchema("age", DataType.INT32)
doc.set_any("age", schema._get_object(), 19)
assert doc.has_field("age")
assert doc.get_any("age", DataType.INT32) == 19
def test_doc_get_set_has_int64_field(self):
doc = _Doc()
schema = FieldSchema("id", DataType.INT64)
doc.set_any("id", schema._get_object(), 1111111111111111111)
assert doc.has_field("id")
assert doc.get_any("id", DataType.INT64) == 1111111111111111111
def test_doc_get_set_has_float_field(self):
doc = _Doc()
schema = FieldSchema("weight", DataType.FLOAT)
doc.set_any("weight", schema._get_object(), 60.5)
assert doc.has_field("weight")
assert math.isclose(doc.get_any("weight", DataType.FLOAT), 60.5, rel_tol=1e-6)
def test_doc_get_set_has_double_field(self):
doc = _Doc()
schema = FieldSchema("height", DataType.DOUBLE)
doc.set_any("height", schema._get_object(), 1.77777777777)
assert doc.has_field("height")
assert math.isclose(
doc.get_any("height", DataType.DOUBLE), 1.7777777777, rel_tol=1e-9
)
def test_doc_get_set_has_uint32_field(self):
doc = _Doc()
schema = FieldSchema("id", DataType.UINT32)
doc.set_any("id", schema._get_object(), 4294967295)
assert doc.has_field("id")
assert doc.get_any("id", DataType.UINT32) == 4294967295
def test_doc_get_set_has_uint64_field(self):
doc = _Doc()
schema = FieldSchema("id", DataType.UINT64)
doc.set_any("id", schema._get_object(), 18446744073709551615)
assert doc.has_field("id")
assert doc.get_any("id", DataType.UINT64) == 18446744073709551615
def test_doc_get_set_has_array_string_field(self):
doc = _Doc()
schema = FieldSchema("tags", DataType.ARRAY_STRING)
doc.set_any("tags", schema._get_object(), ["tag1", "tag2", "tag3"])
assert doc.has_field("tags")
assert doc.get_any("tags", DataType.ARRAY_STRING) == ["tag1", "tag2", "tag3"]
def test_doc_get_set_has_array_int32_field(self):
doc = _Doc()
schema = FieldSchema("ids", DataType.ARRAY_INT32)
doc.set_any("ids", schema._get_object(), [1, 2, 3])
assert doc.has_field("ids")
assert doc.get_any("ids", DataType.ARRAY_INT32) == [1, 2, 3]
def test_doc_get_set_has_array_int64_field(self):
doc = _Doc()
schema = FieldSchema("ids", DataType.ARRAY_INT64)
doc.set_any("ids", schema._get_object(), [1, 2, 3])
assert doc.has_field("ids")
assert doc.get_any("ids", DataType.ARRAY_INT64) == [1, 2, 3]
def test_doc_get_set_has_array_float_field(self):
doc = _Doc()
schema = FieldSchema("weights", DataType.ARRAY_FLOAT)
doc.set_any("weights", schema._get_object(), [1.0, 2.0, 3.0])
assert doc.has_field("weights")
assert doc.get_any("weights", DataType.ARRAY_FLOAT) == [1.0, 2.0, 3.0]
def test_doc_get_set_has_array_double_field(self):
doc = _Doc()
schema = FieldSchema("heights", DataType.ARRAY_DOUBLE)
doc.set_any("heights", schema._get_object(), [1.0, 2.0, 3.0])
assert doc.has_field("heights")
assert doc.get_any("heights", DataType.ARRAY_DOUBLE) == [1.0, 2.0, 3.0]
def test_doc_get_set_has_array_bool_field(self):
doc = _Doc()
schema = FieldSchema("bools", DataType.ARRAY_BOOL)
doc.set_any("bools", schema._get_object(), [True, False, True])
assert doc.has_field("bools")
assert doc.get_any("bools", DataType.ARRAY_BOOL) == [True, False, True]
def test_doc_get_set_has_vector_fp16(self):
doc = _Doc()
schema = VectorSchema("image", DataType.VECTOR_FP16)
doc.set_any("image", schema._get_object(), [1.0, 2.0, 3.0])
assert doc.has_field("image")
image_vector = doc.get_any("image", DataType.VECTOR_FP16)
assert image_vector is not None
for i in range(len(image_vector)):
assert math.isclose(image_vector[i], [1.0, 2.0, 3.0][i], rel_tol=1e-6)
def test_doc_get_set_has_vector_fp32(self):
doc = _Doc()
schema = VectorSchema("image", DataType.VECTOR_FP32)
doc.set_any("image", schema._get_object(), [1.111111, 2.222222, 3.333333])
assert doc.has_field("image")
vector = doc.get_any("image", DataType.VECTOR_FP32)
assert vector is not None
for i in range(len(vector)):
assert math.isclose(
vector[i], [1.111111, 2.222222, 3.333333][i], rel_tol=1e-6
)
def test_doc_get_set_has_vector_int8(self):
doc = _Doc()
schema = VectorSchema("image", DataType.VECTOR_INT8)
doc.set_any("image", schema._get_object(), [1, 2, 3])
assert doc.has_field("image")
assert doc.get_any("image", DataType.VECTOR_INT8) == [1, 2, 3]
def test_doc_get_set_has_sparse_vector_fp32(self):
doc = _Doc()
sparse = {1: 1.111111, 2: 2.222222, 3: 3.333333}
schema = VectorSchema("key", DataType.SPARSE_VECTOR_FP32)
doc.set_any("key", schema._get_object(), sparse)
assert doc.has_field("key")
vector = doc.get_any("key", DataType.SPARSE_VECTOR_FP32)
assert vector is not None
assert isinstance(vector, dict)
for key, value in sparse.items():
assert math.isclose(vector[key], value, rel_tol=1e-6)
def test_doc_get_set_has_sparse_vector_fp16(self):
doc = _Doc()
sparse = {1: 1.1, 2: 2.2, 3: 3.3}
schema = VectorSchema("key", DataType.SPARSE_VECTOR_FP16)
doc.set_any("key", schema._get_object(), sparse)
assert doc.has_field("key")
vector = doc.get_any("key", DataType.SPARSE_VECTOR_FP16)
assert vector is not None
assert isinstance(vector, dict)
for key, value in sparse.items():
assert math.isclose(vector[key], value, rel_tol=1e-1)
File diff suppressed because it is too large Load Diff
+158
View File
@@ -0,0 +1,158 @@
# 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.
"""Tests for FTS (Full-Text Search) query support in the Python SDK."""
import pickle
import pytest
from zvec.model.param.query import Fts, Query
class TestFtsQueryValidation:
"""Test FTS parameter validation in Query dataclass."""
def test_fts_query_string_only(self):
"""Query with only query_string in Fts should be valid."""
q = Query(
field_name="content", fts=Fts(query_string='+hello -world "exact phrase"')
)
q._validate()
assert q.fts.query_string == '+hello -world "exact phrase"'
assert q.fts.match_string is None
assert q.has_fts() is True
def test_fts_match_string_only(self):
"""Query with only match_string in Fts should be valid."""
q = Query(field_name="content", fts=Fts(match_string="machine learning"))
q._validate()
assert q.fts.match_string == "machine learning"
assert q.fts.query_string is None
assert q.has_fts() is True
def test_fts_query_string_and_match_string_mutually_exclusive(self):
"""Cannot provide both query_string and match_string in Fts."""
q = Query(
field_name="content",
fts=Fts(query_string="+hello", match_string="hello world"),
)
with pytest.raises(ValueError, match="mutually exclusive"):
q._validate()
def test_no_fts(self):
"""Query without FTS fields should have has_fts() == False."""
q = Query(field_name="embedding", vector=[0.1, 0.2, 0.3])
assert q.has_fts() is False
def test_vector_and_fts_mutually_exclusive(self):
"""Cannot combine vector search with FTS in a single Query."""
q = Query(
field_name="embedding",
vector=[0.1, 0.2, 0.3],
fts=Fts(match_string="deep learning"),
)
with pytest.raises(ValueError, match="Cannot combine fts with vector search"):
q._validate()
def test_fts_without_vector_or_id(self):
"""Query with only FTS (no vector, no id) should be valid."""
q = Query(field_name="content", fts=Fts(query_string="hello"))
q._validate()
assert q.has_vector() is False
assert q.has_id() is False
assert q.has_fts() is True
class TestFtsQueryBinding:
"""Test FTS binding layer (_Fts)."""
def test_import_fts_query(self):
"""_Fts should be importable from _zvec.param."""
from zvec._zvec.param import _Fts
fts = _Fts()
assert fts.query_string == ""
assert fts.match_string == ""
def test_fts_query_set_fields(self):
"""Setting fields on _Fts should work."""
from zvec._zvec.param import _Fts
fts = _Fts()
fts.query_string = "+hello -world"
assert fts.query_string == "+hello -world"
fts2 = _Fts()
fts2.match_string = "machine learning"
assert fts2.match_string == "machine learning"
def test_fts_query_pickle(self):
"""_Fts should support pickling."""
from zvec._zvec.param import _Fts
fts = _Fts()
fts.query_string = "+vector search"
fts.match_string = ""
data = pickle.dumps(fts)
restored = pickle.loads(data)
assert restored.query_string == "+vector search"
assert restored.match_string == ""
def test_search_query_fts_field(self):
"""_SearchQuery should have fts field."""
from zvec._zvec.param import _Fts, _SearchQuery
vq = _SearchQuery()
# fts should be None by default (optional)
assert vq.fts is None
# set fts
fts = _Fts()
fts.query_string = "hello"
vq.fts = fts
assert vq.fts is not None
assert vq.fts.query_string == "hello"
def test_search_query_pickle_with_fts(self):
"""_SearchQuery with fts should survive pickling."""
from zvec._zvec.param import _Fts, _SearchQuery
vq = _SearchQuery()
vq.topk = 10
vq.field_name = "embedding"
fts = _Fts()
fts.match_string = "test query"
vq.fts = fts
data = pickle.dumps(vq)
restored = pickle.loads(data)
assert restored.topk == 10
assert restored.field_name == "embedding"
assert restored.fts is not None
assert restored.fts.match_string == "test query"
def test_search_query_pickle_without_fts(self):
"""_SearchQuery without fts should survive pickling."""
from zvec._zvec.param import _SearchQuery
vq = _SearchQuery()
vq.topk = 5
vq.field_name = "vec"
data = pickle.dumps(vq)
restored = pickle.loads(data)
assert restored.topk == 5
assert restored.field_name == "vec"
assert restored.fts is None
+273
View File
@@ -0,0 +1,273 @@
# 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.
"""Tests to verify that the GIL is released during native C++ query calls,
enabling true thread-level concurrency for multi-threaded Python applications."""
from __future__ import annotations
import os
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
HnswIndexParam,
Query,
VectorSchema,
)
@pytest.fixture(scope="module")
def gil_test_collection(tmp_path_factory) -> Collection:
"""Create a collection with enough data to make queries take measurable time."""
schema = zvec.CollectionSchema(
name="gil_test",
fields=[
FieldSchema("id", DataType.INT64, nullable=False),
],
vectors=[
VectorSchema(
"vec",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
],
)
option = CollectionOption(read_only=False, enable_mmap=True)
temp_dir = tmp_path_factory.mktemp("zvec_gil_test")
collection_path = temp_dir / "gil_test_collection"
coll = zvec.create_and_open(path=str(collection_path), schema=schema, option=option)
# Insert enough docs to make queries non-trivial
docs = [
Doc(
id=str(i),
fields={"id": i},
vectors={"vec": [float(i % 100) + 0.1 * j for j in range(128)]},
)
for i in range(500)
]
result = coll.insert(docs)
for r in result:
assert r.ok()
yield coll
try:
coll.destroy()
except Exception:
pass
class TestGILRelease:
"""Verify that C++ query calls release the GIL, allowing true thread concurrency."""
def test_gil_released_during_query(self, gil_test_collection: Collection):
"""Prove the GIL is explicitly released during C++ Query calls.
Strategy:
- Calibrate per-query latency on the current platform (slow archs like
RISC-V can be 10x slower than x86), then dynamically pick a query count
whose total runtime fits comfortably inside switch_interval.
- Set switch_interval well above the projected total query time so that
CPython's involuntary GIL switching will NOT trigger during the run.
- A background thread (using time.sleep(0) to avoid deadlock) counts how
many times it got to run.
- Since total query time < switch_interval, the bg thread can ONLY run if
the C++ code explicitly releases the GIL.
- Reset counter just before queries; check counter > 0 after queries.
"""
query_vec = [1.0] * 128
def run_query():
gil_test_collection.query(
Query(field_name="vec", vector=query_vec),
topk=100,
)
# --- Calibrate: estimate per-query latency on this platform ---
# Warm up to avoid first-call overhead skewing the measurement.
for _ in range(3):
run_query()
calib_iters = 10
calib_start = time.monotonic()
for _ in range(calib_iters):
run_query()
per_query = max((time.monotonic() - calib_start) / calib_iters, 1e-6)
# Target total query window ~200ms, capped to a sane range so the test
# remains meaningful on both fast and slow archs.
target_total = 0.2
num_iters = max(1, min(500, int(target_total / per_query)))
projected_total = per_query * num_iters
# Pick switch_interval with a large safety margin (>=10x, >=2s) to absorb
# GC pauses, CPU throttling, and noisy-neighbor effects on CI / shared VMs.
switch_interval = max(2.0, projected_total * 10.0)
old_interval = sys.getswitchinterval()
sys.setswitchinterval(switch_interval)
try:
counter = {"value": 0}
stop_event = threading.Event()
def background_counter():
while not stop_event.is_set():
counter["value"] += 1
time.sleep(0) # Yield GIL to prevent deadlock
bg_thread = threading.Thread(target=background_counter, daemon=True)
bg_thread.start()
# Let bg thread start (sleep releases GIL)
time.sleep(0.05)
# --- Critical section: reset counter, run queries, capture counter ---
counter["value"] = 0
start = time.monotonic()
for _ in range(num_iters):
run_query()
elapsed = time.monotonic() - start
count_during_queries = counter["value"]
# --- End critical section ---
stop_event.set()
time.sleep(0.01)
bg_thread.join(timeout=5)
print(
f"\nPer-query: {per_query * 1000:.2f}ms, iters: {num_iters}, "
f"elapsed: {elapsed:.4f}s, switch_interval: {switch_interval:.2f}s"
)
print(f"Counter during queries: {count_during_queries}")
# Verify queries completed within the switch_interval window.
# If they did NOT, the run was contaminated by external jitter (GC,
# throttling, noisy neighbor) rather than a real GIL-release defect,
# so skip instead of failing to avoid flaky CI noise.
if elapsed >= switch_interval:
pytest.skip(
f"Queries took {elapsed:.3f}s >= switch_interval "
f"({switch_interval:.3f}s); calibration was outpaced by "
"runtime jitter, result is inconclusive."
)
# If elapsed < switch_interval, the ONLY way bg thread could run is
# via explicit GIL release.
assert count_during_queries > 0, (
"Background thread could not run during C++ execution despite "
"query time < switch_interval. GIL was NOT released."
)
finally:
sys.setswitchinterval(old_interval)
def test_parallel_queries_correctness(self, gil_test_collection: Collection):
"""Verify parallel queries return correct results and print timing info.
NOTE: The definitive proof of GIL release is test_gil_released_during_query
(counter + setswitchinterval). This test focuses on parallel correctness and
logs timing for manual inspection, since CI timing is too noisy for assertions.
"""
num_queries = 1000
query_vec = [1.0] * 128
def do_query():
return gil_test_collection.query(
Query(field_name="vec", vector=query_vec),
topk=100,
)
# Serial execution (baseline)
start_serial = time.monotonic()
for _ in range(num_queries):
do_query()
serial_time = time.monotonic() - start_serial
# Parallel execution
num_workers = os.cpu_count() or 2
start_parallel = time.monotonic()
with ThreadPoolExecutor(max_workers=num_workers) as executor:
futures = [executor.submit(do_query) for _ in range(num_queries)]
for future in as_completed(futures):
result = future.result()
assert len(result) > 0
parallel_time = time.monotonic() - start_parallel
print(f"\nSerial time: {serial_time:.4f}s, Parallel time: {parallel_time:.4f}s")
print(
f"Speedup ratio: {serial_time / parallel_time:.2f}x (workers={num_workers})"
)
def test_thread_safety_concurrent_queries(self, gil_test_collection: Collection):
"""Verify no crashes or data corruption under concurrent query load."""
num_threads = 8
queries_per_thread = 10
errors = []
def worker(thread_id):
try:
for i in range(queries_per_thread):
vec = [float(thread_id + i) + 0.1 * j for j in range(128)]
result = gil_test_collection.query(
Query(field_name="vec", vector=vec),
topk=10,
)
assert len(result) > 0
except Exception as e:
errors.append((thread_id, e))
threads = [
threading.Thread(target=worker, args=(tid,)) for tid in range(num_threads)
]
for t in threads:
t.start()
for t in threads:
t.join(timeout=60)
assert len(errors) == 0, f"Errors in threads: {errors}"
def test_concurrent_fetch_release_gil(self, gil_test_collection: Collection):
"""Verify Fetch operations also release the GIL correctly."""
num_threads = 4
errors = []
def worker(thread_id):
try:
ids = [str(i) for i in range(thread_id * 10, thread_id * 10 + 10)]
result = gil_test_collection.fetch(ids)
assert len(result) > 0
except Exception as e:
errors.append((thread_id, e))
threads = [
threading.Thread(target=worker, args=(tid,)) for tid in range(num_threads)
]
for t in threads:
t.start()
for t in threads:
t.join(timeout=30)
assert len(errors) == 0, f"Errors in threads: {errors}"
+415
View File
@@ -0,0 +1,415 @@
# 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.
"""
Tests for the ``use_contiguous_memory`` optimization on ``HnswIndexParam``.
The HNSW streamer supports two allocation strategies for graph nodes:
* ``use_contiguous_memory=False`` (default): each node allocates its own
linked buffer. Lower peak memory usage, worse cache locality.
* ``use_contiguous_memory=True``: a single contiguous arena holds every
node. Higher peak memory usage, better cache locality and search
throughput.
These tests exercise the Python surface end-to-end and make sure that
when a collection is created / reopened with ``use_contiguous_memory=True``
the underlying HNSW streamer entity is constructed correctly and serves
search traffic.
"""
from __future__ import annotations
import pickle
import sys
import numpy as np
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
CollectionSchema,
Doc,
FieldSchema,
HnswIndexParam,
HnswQueryParam,
InvertIndexParam,
Query,
VectorSchema,
)
from zvec.typing import DataType, IndexType, MetricType, QuantizeType
DIMENSION = 32
NUM_DOCS = 128
TOPK = 5
# ---------------------------------------------------------------------------
def _debug_hnsw_storage_mode(coll: Collection, column: str = "dense") -> str:
"""Return the internal HNSW entity storage mode for ``column``.
Exposes the debug-only introspection hook on the pybind11 ``_Collection``.
Only meaningful after ``optimize()`` has built a persisted HNSW index; on
a pure writing segment it will raise ``KeyError``.
"""
underlying = coll._obj # type: ignore[attr-defined]
return underlying._debug_hnsw_storage_mode(column)
def _build_schema(name: str, *, use_contiguous_memory: bool) -> CollectionSchema:
"""Create a simple schema with a single FP32 HNSW vector column."""
return CollectionSchema(
name=name,
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=DIMENSION,
index_param=HnswIndexParam(
metric_type=MetricType.IP,
m=16,
ef_construction=100,
use_contiguous_memory=use_contiguous_memory,
),
),
],
)
def _generate_docs(rng: np.random.Generator, num: int = NUM_DOCS) -> list[Doc]:
"""Produce deterministic documents for insertion."""
docs: list[Doc] = []
for i in range(num):
vec = rng.standard_normal(DIMENSION).astype(np.float32)
docs.append(
Doc(
id=str(i),
fields={"id": i},
vectors={"dense": vec.tolist()},
)
)
return docs
def _assert_query_matches(coll: Collection, query_vec: list[float]) -> list[str]:
"""Run a top-k vector query and return the returned ids in order."""
vector_query = Query(
field_name="dense",
vector=query_vec,
param=HnswQueryParam(ef=128),
)
hits = coll.query(vector_query, topk=TOPK)
# Expect a single result group for the single vector query.
assert hits is not None, "query returned None"
assert len(hits) >= 1, f"expected at least one hit, got {hits!r}"
return [doc.id for doc in hits]
# ---------------------------------------------------------------------------
# 1) Pure Python surface: construction / property / to_dict / repr / pickle
# ---------------------------------------------------------------------------
class TestHnswIndexParamContiguousMemorySurface:
"""Verify the Python binding exposes ``use_contiguous_memory`` correctly."""
def test_default_is_false(self):
param = HnswIndexParam()
assert param.use_contiguous_memory is False
def test_custom_true(self):
param = HnswIndexParam(use_contiguous_memory=True)
assert param.use_contiguous_memory is True
assert param.type == IndexType.HNSW
# other fields keep their default values
assert param.m == 50
assert param.ef_construction == 500
def test_to_dict_includes_use_contiguous_memory(self):
param = HnswIndexParam(
metric_type=MetricType.L2,
m=16,
ef_construction=100,
quantize_type=QuantizeType.FP16,
use_contiguous_memory=True,
)
data = param.to_dict()
assert data["use_contiguous_memory"] is True
# Make sure existing fields are still present.
assert data["metric_type"] == "L2"
assert data["m"] == 16
assert data["ef_construction"] == 100
assert data["quantize_type"] == "FP16"
def test_repr_contains_flag(self):
on = repr(HnswIndexParam(use_contiguous_memory=True))
off = repr(HnswIndexParam(use_contiguous_memory=False))
assert "use_contiguous_memory" in on
assert "use_contiguous_memory" in off
assert "true" in on
assert "false" in off
def test_readonly_property(self):
param = HnswIndexParam(use_contiguous_memory=True)
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
param.use_contiguous_memory = False # type: ignore[misc]
def test_pickle_roundtrip(self):
original = HnswIndexParam(
metric_type=MetricType.COSINE,
m=24,
ef_construction=150,
quantize_type=QuantizeType.INT8,
use_contiguous_memory=True,
)
restored = pickle.loads(pickle.dumps(original))
assert restored.use_contiguous_memory is True
assert restored.metric_type == MetricType.COSINE
assert restored.m == 24
assert restored.ef_construction == 150
assert restored.quantize_type == QuantizeType.INT8
# ---------------------------------------------------------------------------
# 2) End-to-end: create collection, insert, query with contiguous memory on
# ---------------------------------------------------------------------------
@pytest.fixture
def rng() -> np.random.Generator:
return np.random.default_rng(seed=42)
# NOTE: the ``enable_mmap=False`` (BufferPool) variant is intentionally
# omitted from this fixture. Building a persisted HNSW index via
# ``optimize()`` / ``create_vector_index`` / ``drop_vector_index``
# currently requires mmap-backed storage, because the BufferPool backend
# has not implemented the ``create_new`` semantics yet and the guard in
# ``SegmentImpl::merge_vector_indexer`` rejects that combination. Once
# BufferPool gains write support, re-add ``False`` to ``params`` (and
# drop the guard in segment.cc) so these end-to-end tests cover both
# storage modes again.
@pytest.fixture(params=[True], ids=["mmap_on"])
def collection_option(request) -> CollectionOption:
return CollectionOption(read_only=False, enable_mmap=request.param)
# Building a new persisted HNSW index currently requires mmap-backed storage
# because the BufferPool backend has not implemented `create_new` semantics
# yet. Collections opened with ``enable_mmap=False`` therefore cannot run
# optimize()/create_vector_index/drop_vector_index. Tests use this fixture
# to know which behaviour to assert, and once BufferPool gains write support
# the guard in segment.cc (and these branches) can be removed together.
@pytest.fixture
def build_index_supported(collection_option: CollectionOption) -> bool:
return bool(collection_option.enable_mmap)
# Error message fragments emitted by the NotSupported guard in
# SegmentImpl::merge_vector_indexer / drop_vector_index. If the C++ message
# changes, update these together.
_BUILD_NOT_SUPPORTED_FRAGMENTS = ("not yet supported", "enable_mmap=false")
class TestHnswContiguousMemoryEndToEnd:
"""End-to-end: schema -> create_and_open -> insert -> query works."""
def test_create_with_contiguous_memory_and_query(
self,
tmp_path_factory,
collection_option,
rng,
):
"""With the flag on, the schema round-trips and search works end-to-end.
After ``optimize()`` the writing segment is compacted into a persisted
segment backed by the configured HNSW entity. We assert both the
user-observable behaviour (schema + search) and, via the debug hook,
that the entity type actually honours ``use_contiguous_memory``.
"""
schema = _build_schema("hnsw_contig_create", use_contiguous_memory=True)
path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_create"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
# Schema round-trips with the flag set.
vec_schema = coll.schema.vectors[0]
assert vec_schema.index_param.use_contiguous_memory is True
docs = _generate_docs(rng)
insert_result = coll.insert(docs=docs)
for r in insert_result:
assert r.ok(), f"insert failed: code={r.code()}"
assert coll.stats.doc_count == NUM_DOCS
# Build persisted HNSW index; this is where the contiguous entity
# is actually instantiated.
coll.optimize()
assert _debug_hnsw_storage_mode(coll) == "contiguous", (
"use_contiguous_memory=True should produce a contiguous entity"
)
# Pick an existing vector as the query; top-1 must be itself.
query_vec = docs[0].vector("dense")
ids = _assert_query_matches(coll, query_vec)
assert ids[0] == "0", f"expected self-recall, got top-1 id={ids[0]}"
finally:
coll.destroy()
def test_create_without_contiguous_memory_uses_mmap_entity(
self,
tmp_path_factory,
collection_option,
rng,
):
"""Baseline: when the flag is omitted the default (mmap) entity is used."""
schema = _build_schema("hnsw_contig_default", use_contiguous_memory=False)
path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_default"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
vec_schema = coll.schema.vectors[0]
assert vec_schema.index_param.use_contiguous_memory is False
docs = _generate_docs(rng)
for r in coll.insert(docs=docs):
assert r.ok()
assert coll.stats.doc_count == NUM_DOCS
coll.optimize()
# With the flag off and mmap on, the persisted entity must be the
# default mmap layout — specifically, not the contiguous arena.
assert _debug_hnsw_storage_mode(coll) == "mmap", (
"use_contiguous_memory=False + enable_mmap=True should "
"produce the mmap entity"
)
# Search still functions with the default entity backing.
query_vec = docs[0].vector("dense")
ids = _assert_query_matches(coll, query_vec)
assert ids[0] == "0"
finally:
coll.destroy()
def test_close_and_reopen_with_contiguous_memory(
self,
tmp_path_factory,
collection_option,
rng,
):
"""Reopening a collection must preserve the ``use_contiguous_memory`` flag.
The core property: the flag survives the schema persist/reload
round-trip so the HNSW streamer entity — constructed lazily on first
persisted-segment build — honours the user's choice. We run
``optimize()`` after reopen and confirm the contiguous entity was
materialized.
"""
schema = _build_schema("hnsw_contig_reopen", use_contiguous_memory=True)
path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_reopen"
path_str = str(path)
created = zvec.create_and_open(
path=path_str, schema=schema, option=collection_option
)
docs = _generate_docs(rng)
for r in created.insert(docs=docs):
assert r.ok()
assert created.stats.doc_count == NUM_DOCS
# Persist pending writes so that reopen reconstructs state from disk.
created.flush()
del created # close the handle
reopened = zvec.open(path=path_str, option=collection_option)
try:
assert reopened is not None
assert reopened.stats.doc_count == NUM_DOCS
# Schema persisted the flag across the reopen boundary.
vec_schema = reopened.schema.vectors[0]
assert vec_schema.index_param.use_contiguous_memory is True
reopened.optimize()
assert _debug_hnsw_storage_mode(reopened) == "contiguous"
# Entity actually works: exact self-recall + fetch parity.
query_vec = docs[7].vector("dense")
ids = _assert_query_matches(reopened, query_vec)
assert ids[0] == "7"
fetched = reopened.fetch([d.id for d in docs[:10]])
assert len(fetched) == 10
finally:
reopened.destroy()
def test_result_parity_with_and_without_contiguous_memory(
self,
tmp_path_factory,
rng,
):
"""
Two collections built from the same documents must return the same
top-k neighbors regardless of whether contiguous memory is enabled:
the flag is a memory-layout optimization and must not alter recall
for identical graph construction parameters on the same data.
"""
docs = _generate_docs(rng)
query_vec = docs[3].vector("dense")
def _build_and_query(tag: str, flag: bool) -> list[str]:
schema = _build_schema(f"hnsw_parity_{tag}", use_contiguous_memory=flag)
option = CollectionOption(read_only=False, enable_mmap=True)
path = tmp_path_factory.mktemp("zvec") / f"hnsw_parity_{tag}"
coll = zvec.create_and_open(path=str(path), schema=schema, option=option)
try:
for r in coll.insert(docs=docs):
assert r.ok()
coll.optimize()
expected_mode = "contiguous" if flag else "mmap"
assert _debug_hnsw_storage_mode(coll) == expected_mode, (
f"{tag}: unexpected entity type"
)
return _assert_query_matches(coll, query_vec)
finally:
coll.destroy()
ids_off = _build_and_query("off", flag=False)
ids_on = _build_and_query("on", flag=True)
# The graph is built with the same (m, ef_construction, data, order),
# so top-k results must match exactly.
assert ids_on == ids_off, (
f"top-{TOPK} results diverged between use_contiguous_memory modes: "
f"on={ids_on}, off={ids_off}"
)
# Sanity: self-recall is still perfect.
assert ids_on[0] == "3"
+201
View File
@@ -0,0 +1,201 @@
# 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.
"""End-to-end: jieba FTS works without any user configuration.
`import zvec` is supposed to register the wheel-bundled jieba dict
directory via `set_default_jieba_dict_dir`. With that in place a user can
declare an FTS field with `tokenizer_name="jieba"`, leave `extra_params`
empty, and Chinese full-text search just works.
Falls back to GTEST_SKIP-equivalent when running against a build that did
not bundle the dict (e.g., source-tree dev install without the install
step). In that case CI will rely on the C++ unit tests instead.
"""
from __future__ import annotations
import os
import sys
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
DataType,
Doc,
FieldSchema,
FtsIndexParam,
)
from zvec.model.param.query import Fts, Query
def _bundled_dict_dir() -> str:
"""Path zvec.__init__ would have registered; empty when not bundled."""
return zvec.get_default_jieba_dict_dir()
def _bundled_dict_files_exist() -> bool:
"""Whether the registered default actually contains the dict files.
`importlib.resources` happily returns a path even when the data dir was
not installed (e.g. source-tree dev runs); only an installed wheel has
the files on disk.
"""
import os
base = _bundled_dict_dir()
if not base:
return False
return os.path.isfile(os.path.join(base, "jieba.dict.utf8")) and os.path.isfile(
os.path.join(base, "hmm_model.utf8")
)
@pytest.fixture(scope="module", autouse=True)
def _require_bundled_dict():
if not _bundled_dict_files_exist():
pytest.skip(
"Bundled jieba dict not found at zvec/data/jieba_dict/ — "
"this test requires an installed wheel (not a source-tree dev "
"build without the install step).",
)
@pytest.fixture(scope="function")
def jieba_collection(tmp_path_factory) -> Collection:
"""FTS-only collection using jieba tokenizer and no explicit dict path."""
# env-var shadows GlobalConfig in the priority chain.
if os.environ.get("ZVEC_JIEBA_DICT_DIR"):
pytest.skip("ZVEC_JIEBA_DICT_DIR shadows the bundled default")
temp_dir = tmp_path_factory.mktemp("zvec_jieba_default")
collection_path = temp_dir / "fts_jieba"
schema = zvec.CollectionSchema(
name="fts_jieba_default",
fields=[
FieldSchema("title", DataType.STRING, nullable=False),
FieldSchema(
"content",
DataType.STRING,
nullable=False,
# Deliberately omit extra_params — the bundled default must
# be picked up via GlobalConfig.jieba_dict_dir.
index_param=FtsIndexParam(
tokenizer_name="jieba",
filters=["lowercase"],
),
),
],
)
coll = zvec.create_and_open(
path=str(collection_path),
schema=schema,
option=CollectionOption(read_only=False, enable_mmap=True),
)
assert coll is not None
try:
yield coll
finally:
try:
coll.destroy()
except Exception as e:
print(f"Warning: failed to destroy collection: {e}")
def test_jieba_works_without_explicit_dict_path(jieba_collection: Collection):
"""User opens collection, inserts CJK doc, searches — no init() / no
extra_params / no env var / no manual setter call. Just `import zvec`."""
docs = [
Doc(id="pk_1", fields={"title": "t1", "content": "中华人民共和国成立"}),
Doc(id="pk_2", fields={"title": "t2", "content": "无关文档"}),
]
insert_results = jieba_collection.insert(docs)
assert all(r.ok() for r in insert_results)
hits = jieba_collection.query(
queries=Query(field_name="content", fts=Fts(match_string="中华")),
topk=10,
)
ids = {doc.id for doc in hits}
assert "pk_1" in ids
assert "pk_2" not in ids
def test_default_dict_dir_is_registered_on_import():
"""Sanity check: zvec.__init__ registered a non-empty default."""
assert _bundled_dict_dir() != ""
def test_user_can_override_default_at_runtime():
"""zvec.set_default_jieba_dict_dir can be called any time to override."""
saved = zvec.get_default_jieba_dict_dir()
try:
zvec.set_default_jieba_dict_dir("/tmp/zvec/jieba-override")
assert zvec.get_default_jieba_dict_dir() == "/tmp/zvec/jieba-override"
finally:
zvec.set_default_jieba_dict_dir(saved)
@pytest.mark.skipif(
sys.platform == "win32",
reason="os.environ writes may not propagate across CRT to zvec.pyd",
)
def test_env_var_overrides_global_config(monkeypatch, tmp_path_factory):
"""ZVEC_JIEBA_DICT_DIR beats GlobalConfig in jieba's resolution chain."""
bundled = _bundled_dict_dir()
monkeypatch.setenv("ZVEC_JIEBA_DICT_DIR", bundled)
saved_global = zvec.get_default_jieba_dict_dir()
try:
zvec.set_default_jieba_dict_dir("/zvec/intentionally/missing/global")
temp_dir = tmp_path_factory.mktemp("zvec_jieba_env")
schema = zvec.CollectionSchema(
name="fts_jieba_env",
fields=[
FieldSchema("title", DataType.STRING, nullable=False),
FieldSchema(
"content",
DataType.STRING,
nullable=False,
index_param=FtsIndexParam(
tokenizer_name="jieba",
filters=["lowercase"],
),
),
],
)
coll = zvec.create_and_open(
path=str(temp_dir / "fts_jieba_env"),
schema=schema,
option=CollectionOption(read_only=False, enable_mmap=True),
)
assert coll is not None
try:
results = coll.insert(
[
Doc(id="pk_1", fields={"title": "t", "content": "搜索引擎技术"}),
]
)
assert all(r.ok() for r in results)
hits = coll.query(
queries=Query(field_name="content", fts=Fts(match_string="搜索")),
topk=10,
)
assert {d.id for d in hits} == {"pk_1"}
finally:
coll.destroy()
finally:
zvec.set_default_jieba_dict_dir(saved_global)
+540
View File
@@ -0,0 +1,540 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import sys
import time
import numpy as np
import pytest
from zvec import (
AddColumnOption,
AlterColumnOption,
CollectionOption,
FlatIndexParam,
HnswIndexParam,
IndexOption,
InvertIndexParam,
IVFIndexParam,
OptimizeOption,
HnswQueryParam,
IVFQueryParam,
Query,
VectorQuery,
IndexType,
MetricType,
QuantizeType,
QuantizerParam,
DataType,
VectorSchema,
)
from zvec._zvec.param import _SearchQuery
# ----------------------------
# Invert Index Param Test Case
# ----------------------------
class TestInvertIndexParam:
def test_default(self):
param = InvertIndexParam()
assert param.enable_range_optimization is False
assert param.enable_extended_wildcard is False
assert param.type == IndexType.INVERT
def test_custom(self):
param = InvertIndexParam(
enable_range_optimization=True, enable_extended_wildcard=True
)
assert param.enable_range_optimization is True
assert param.enable_extended_wildcard is True
def test_readonly(self):
param = InvertIndexParam()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
param.enable_range_optimization = False
param.enable_extended_wildcard = False
# ----------------------------
# Hnsw Index Param Test Case
# ----------------------------
class TestHnswIndexParam:
def test_default(self):
param = HnswIndexParam()
assert param.metric_type == MetricType.IP
assert param.m == 50
assert param.ef_construction == 500
assert param.quantize_type == QuantizeType.UNDEFINED
assert param.type == IndexType.HNSW
def test_custom(self):
param = HnswIndexParam(
metric_type=MetricType.L2,
m=10,
ef_construction=1000,
quantize_type=QuantizeType.FP16,
)
assert param.metric_type == MetricType.L2
assert param.m == 10
assert param.ef_construction == 1000
assert param.quantize_type == QuantizeType.FP16
@pytest.mark.parametrize(
"attr", ["metric_type", "m", "ef_construction", "quantize_type"]
)
def test_readonly_attributes(self, attr):
param = HnswIndexParam()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# Flat Index Param Test Case
# ----------------------------
class TestFlatIndexParam:
def test_default(self):
param = FlatIndexParam()
assert param.type == IndexType.FLAT
assert param.quantize_type == QuantizeType.UNDEFINED
assert param.metric_type == MetricType.IP
def test_custom(self):
param = FlatIndexParam(
metric_type=MetricType.L2, quantize_type=QuantizeType.INT8
)
assert param.metric_type == MetricType.L2
assert param.quantize_type == QuantizeType.INT8
@pytest.mark.parametrize("attr", ["metric_type", "quantize_type"])
def test_readonly_attributes(self, attr):
param = FlatIndexParam()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# Ivf Index Param Test Case
# ----------------------------
class TestIVFIndexParam:
def test_default(self):
param = IVFIndexParam()
assert param.metric_type == MetricType.IP
assert param.n_list == 10
assert param.quantize_type == QuantizeType.UNDEFINED
assert param.type == IndexType.IVF
def test_custom(self):
param = IVFIndexParam(
metric_type=MetricType.L2, n_list=1000, quantize_type=QuantizeType.FP16
)
assert param.metric_type == MetricType.L2
assert param.n_list == 1000
assert param.quantize_type == QuantizeType.FP16
assert param.type == IndexType.IVF
@pytest.mark.parametrize("attr", ["metric_type", "n_list", "quantize_type"])
def test_readonly_attributes(self, attr):
param = IVFIndexParam()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# CollectionOption Test Case
# ----------------------------
class TestCollectionOption:
def test_default(self):
option = CollectionOption()
assert option is not None
assert option.read_only == False
assert option.enable_mmap == True
def test_custom(self):
option = CollectionOption(read_only=True, enable_mmap=False)
assert option.read_only == True
assert option.enable_mmap == False
option = CollectionOption(read_only=False, enable_mmap=True)
assert option.read_only == False
assert option.enable_mmap == True
@pytest.mark.parametrize("attr", ["read_only", "enable_mmap"])
def test_readonly_attributes(self, attr):
param = CollectionOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# IndexOption Test Case
# ----------------------------
class TestIndexOption:
def test_default(self):
option = IndexOption()
assert option is not None
assert option.concurrency == 0
def test_custom(self):
option = IndexOption(concurrency=10)
assert option.concurrency == 10
@pytest.mark.parametrize("attr", ["concurrency"])
def test_readonly_attributes(self, attr):
param = IndexOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# AddColumnOption Test Case
# ----------------------------
class TestAddColumnOption:
def test_default(self):
option = AddColumnOption()
assert option is not None
assert option.concurrency == 0
def test_custom(self):
option = AddColumnOption(concurrency=10)
assert option.concurrency == 10
@pytest.mark.parametrize("attr", ["concurrency"])
def test_readonly_attributes(self, attr):
param = AddColumnOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# AlterColumnOption Test Case
# ----------------------------
class TestAlterColumnOption:
def test_default(self):
option = AlterColumnOption()
assert option is not None
assert option.concurrency == 0
def test_custom(self):
option = AlterColumnOption(concurrency=10)
assert option.concurrency == 10
@pytest.mark.parametrize("attr", ["concurrency"])
def test_readonly_attributes(self, attr):
param = AlterColumnOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# OptimizeOption Test Case
# ----------------------------
class TestOptimizeOption:
def test_default(self):
option = OptimizeOption()
assert option is not None
assert option.concurrency == 0
def test_custom(self):
option = OptimizeOption(concurrency=10)
assert option.concurrency == 10
@pytest.mark.parametrize("attr", ["concurrency"])
def test_readonly_attributes(self, attr):
param = OptimizeOption()
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, attr, getattr(param, attr))
# ----------------------------
# HnswQueryParam Test Case
# ----------------------------
class TestHnswQueryParam:
def test_default(self):
param = HnswQueryParam()
assert param is not None
assert param.ef == 300
assert param.is_using_refiner == False
assert param.radius == 0
assert param.is_linear == False
assert param.prefetch_offset == 8
assert param.prefetch_lines == 0
def test_custom(self):
param = HnswQueryParam(
ef=10,
is_using_refiner=True,
radius=30,
is_linear=True,
extra_params={
"prefetch_offset": 16,
"prefetch_lines": 4,
},
)
assert param.ef == 10
assert param.is_using_refiner == True
assert param.radius == 30
assert param.is_linear == True
assert param.prefetch_offset == 16
assert param.prefetch_lines == 4
def test_readonly_attributes(self):
param = HnswQueryParam()
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
param.ef = 10
param.is_using_refiner = True
param.radius = 30
param.is_linear = True
# # ----------------------------
# # IVFQueryParam Test Case
# # ----------------------------
# class TestIVFQueryParam:
# def test_default(self):
# param = IVFQueryParam()
# assert param is not None
# assert param.nprobe == 10
# assert param.is_using_refiner == False
# assert param.radius == 0
# assert param.is_linear == False
# assert param.scale_factor == 10
#
# def test_custom(self):
# param = IVFQueryParam(
# nprobe=20,
# is_using_refiner=True,
# radius=30,
# is_linear=True,
# scale_factor=40
# )
# assert param.nprobe == 20
# assert param.is_using_refiner == True
# assert param.radius == 30
# assert param.is_linear == True
# assert param.scale_factor == 40
class TestQuery:
def test_init_with_valid_id(self):
vq = Query(field_name="embedding", id="doc123")
assert vq.field_name == "embedding"
assert vq.id == "doc123"
assert vq.vector is None
assert vq.param is None
def test_init_with_valid_vector(self):
vec = [0.1, 0.2, 0.3]
param = HnswQueryParam(ef=300)
vq = Query(field_name="embedding", vector=vec, param=param)
assert vq.field_name == "embedding"
assert vq.vector == vec
assert vq.param == param
def test_init_both_id_and_vector_raises_error(self):
with pytest.raises(ValueError):
Query(field_name="embedding", id="doc123", vector=[0.1])._validate()
def test_init_without_field_name_raises_error(self):
with pytest.raises(ValueError):
Query(field_name=None)._validate()
def test_has_id_returns_true_when_id_set(self):
vq = Query(field_name="embedding", id="doc123")
assert vq.has_id()
def test_has_id_returns_false_when_no_id(self):
vq = Query(field_name="embedding", vector=[0.1])
assert not vq.has_id()
def test_has_vector_returns_true_with_non_empty_vector(self):
vq = Query(field_name="embedding", vector=[0.1])
assert vq.has_vector()
def test_validate_fails_on_both_id_and_vector(self):
vq = Query(field_name="test", id="doc123", vector=[0.1])
with pytest.raises(ValueError):
vq._validate()
def test_validate_fails_on_both_id_and_numpy_vector(self):
vq = Query(field_name="test", id="doc123", vector=np.array([0.1]))
with pytest.raises(ValueError, match="Cannot provide both id and vector"):
vq._validate()
class TestVectorQueryDeprecated:
def test_deprecation_warning(self):
import warnings
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
vq = VectorQuery(field_name="embedding", id="doc123")
assert len(w) == 1
assert issubclass(w[0].category, DeprecationWarning)
assert "Query" in str(w[0].message)
def test_isinstance_compatibility(self):
import warnings
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
vq = VectorQuery(field_name="embedding", id="doc123")
assert isinstance(vq, Query)
# ----------------------------
# QuantizerParam Test Case
# ----------------------------
class TestQuantizerParam:
def test_default(self):
qp = QuantizerParam()
assert qp.enable_rotate is False
def test_enable_rotate_true(self):
qp = QuantizerParam(enable_rotate=True)
assert qp.enable_rotate is True
def test_enable_rotate_false(self):
qp = QuantizerParam(enable_rotate=False)
assert qp.enable_rotate is False
def test_equality(self):
qp1 = QuantizerParam(enable_rotate=True)
qp2 = QuantizerParam(enable_rotate=True)
qp3 = QuantizerParam(enable_rotate=False)
assert qp1 == qp2
assert qp1 != qp3
def test_to_dict(self):
qp = QuantizerParam(enable_rotate=True)
d = qp.to_dict()
assert isinstance(d, dict)
assert d.get("enable_rotate") is True
def test_repr(self):
qp = QuantizerParam(enable_rotate=True)
r = repr(qp)
assert "enable_rotate" in r or "QuantizerParam" in r
def test_pickle_roundtrip(self):
import pickle
qp = QuantizerParam(enable_rotate=True)
data = pickle.dumps(qp)
qp2 = pickle.loads(data)
assert qp2.enable_rotate is True
assert qp == qp2
# ----------------------------
# HnswIndexParam with QuantizerParam
# ----------------------------
class TestHnswIndexParamQuantizer:
def test_default_quantizer_param(self):
param = HnswIndexParam()
assert param.quantizer_param is not None
assert param.quantizer_param.enable_rotate is False
def test_with_quantizer_param(self):
qp = QuantizerParam(enable_rotate=True)
param = HnswIndexParam(
metric_type=MetricType.L2,
quantize_type=QuantizeType.INT8,
quantizer_param=qp,
)
assert param.quantizer_param.enable_rotate is True
assert param.quantize_type == QuantizeType.INT8
# ----------------------------
# FlatIndexParam with QuantizerParam
# ----------------------------
class TestFlatIndexParamQuantizer:
def test_with_quantizer_param(self):
qp = QuantizerParam(enable_rotate=True)
param = FlatIndexParam(
metric_type=MetricType.L2,
quantize_type=QuantizeType.INT8,
quantizer_param=qp,
)
assert param.quantizer_param.enable_rotate is True
assert param.quantize_type == QuantizeType.INT8
+329
View File
@@ -0,0 +1,329 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Dict, Union
from unittest.mock import MagicMock, patch
import numpy as np
import math
from zvec._zvec.param import _SearchQuery
import pytest
from zvec.executor.query_executor import (
QueryContext,
QueryExecutor,
)
from zvec import (
RrfReRanker,
WeightedReRanker,
HnswQueryParam,
CollectionSchema,
VectorSchema,
DataType,
MetricType,
Query,
VectorQuery,
)
from zvec.extension.multi_vector_reranker import CallbackReRanker
# ----------------------------
# Mock Collection Schema
# ----------------------------
class MockCollectionSchema(CollectionSchema):
def __init__(self, vectors=Union[VectorSchema, Dict[str, VectorSchema]]):
self._vectors = (
[vectors] if not isinstance(vectors, Dict) else list(vectors.values())
)
@property
def vectors(self):
return self._vectors
# ----------------------------
# VectorQuery Test Case
# ----------------------------
class TestQuery:
def test_init(self):
query = Query(field_name="test_field")
assert query.field_name == "test_field"
assert query.id is None
assert query.vector is None
assert query.param is None
param = HnswQueryParam()
query = Query(
field_name="test_field", id="test_id", vector=[1, 2, 3], param=param
)
assert query.field_name == "test_field"
assert query.id == "test_id"
assert query.vector == [1, 2, 3]
assert query.param == param
def test_has_id(self):
query = Query(field_name="test_field")
assert not query.has_id()
query = Query(field_name="test_field", id="test_id")
assert query.has_id()
def test_has_vector(self):
query = Query(field_name="test_field")
assert not query.has_vector()
query = Query(field_name="test_field", vector=[])
assert not query.has_vector()
query = Query(field_name="test_field", vector=[1, 2, 3])
assert query.has_vector()
def test_validate_dense_fp16_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP16)
vec = np.array([1.1, 2.1, 3.1], dtype=np.float16)
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
assert np.array_equal(vec, ret)
def test_validate_dense_fp32_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32)
vec = np.array([1.1, 2.1, 3.1], dtype=np.float32)
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
assert np.array_equal(vec, ret)
def test_validate_dense_fp64_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP64)
vec = np.array([1.1, 2.1, 3.1], dtype=np.float64)
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
assert np.array_equal(vec, ret)
def test_validate_dense_int8_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.VECTOR_INT8)
vec = np.array([1, 2, 3], dtype=np.int8)
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
assert np.array_equal(vec, ret)
def test_validate_sparse_fp32_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.SPARSE_VECTOR_FP32)
vec = {1: 1.1, 2: 2.2, 3: 3.3}
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
for k in vec.keys():
assert math.isclose(vec[k], ret[k], abs_tol=1e-6)
def test_validate_sparse_fp16_convert(self):
v = _SearchQuery()
schema = VectorSchema(name="test", data_type=DataType.SPARSE_VECTOR_FP16)
vec = {1: 1.1, 2: 2.2, 3: 3.3}
v.set_vector(schema._get_object(), vec)
ret = v.get_vector(schema._get_object())
for k in vec.keys():
assert math.isclose(np.float16(vec[k]), ret[k], abs_tol=1e-6)
class TestVectorQueryDeprecated:
def test_deprecation_warning(self):
import warnings
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
vq = VectorQuery(field_name="test_field")
assert len(w) == 1
assert issubclass(w[0].category, DeprecationWarning)
assert "Query" in str(w[0].message)
def test_isinstance_compatibility(self):
import warnings
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
vq = VectorQuery(field_name="test_field")
assert isinstance(vq, Query)
class TestQueryContext:
def test_init(self):
ctx = QueryContext(topk=10)
assert ctx.topk == 10
assert ctx.queries == []
assert ctx.filter is None
assert ctx.reranker is None
assert ctx.output_fields is None
assert ctx.include_vector is False
def test_properties(self):
queries = [Query(field_name="test")]
reranker = RrfReRanker()
output_fields = ["field1", "field2"]
ctx = QueryContext(
topk=5,
filter="test_filter",
include_vector=True,
queries=queries,
output_fields=output_fields,
reranker=reranker,
)
assert ctx.topk == 5
assert ctx.queries == queries
assert ctx.filter == "test_filter"
assert ctx.reranker == reranker
assert ctx.output_fields == output_fields
assert ctx.include_vector is True
def test_properties_with_weighted_reranker(self):
queries = [Query(field_name="test")]
reranker = WeightedReRanker(
weights=[1.0],
)
ctx = QueryContext(
topk=5,
queries=queries,
reranker=reranker,
)
assert ctx.reranker == reranker
assert ctx.reranker.weights == [1.0]
def test_properties_with_callback_reranker(self):
queries = [Query(field_name="test")]
cb = lambda query_results, topn: []
reranker = CallbackReRanker(callback=cb)
ctx = QueryContext(
topk=5,
queries=queries,
reranker=reranker,
)
assert ctx.reranker == reranker
class TestQueryExecutor:
def test_init(self):
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
assert isinstance(executor, QueryExecutor)
def test_do_build_without_queries(self):
# When no queries are given, build a single vector-less query.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5, filter="test_filter")
result = executor._build_queries(ctx, MagicMock())
assert len(result) == 1
assert result[0].topk == 5
assert result[0].filter == "test_filter"
def test_do_build_query_wo_vector(self):
# Vector-less core query should carry the context query params.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
ctx = QueryContext(topk=7, filter="f", include_vector=True)
core_vector = executor._build_base_search_query(ctx)
assert core_vector.topk == 7
assert core_vector.filter == "f"
assert core_vector.include_vector is True
def test_do_merge_rerank_results_single_without_reranker(self):
# A single result list without a reranker is returned as-is.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
docs_list = [["doc1", "doc2"]]
result = executor._merge_and_rerank(ctx, docs_list)
assert result == ["doc1", "doc2"]
def test_do_merge_rerank_results_empty(self):
# Empty results should raise an error.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
with pytest.raises(ValueError, match="Query results is empty"):
executor._merge_and_rerank(ctx, [])
def test_do_merge_rerank_results_with_reranker(self):
# Multiple result lists are merged through the reranker.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
reranker = MagicMock()
reranker.rerank.return_value = ["merged"]
ctx = QueryContext(
topk=5,
queries=[Query(field_name="test1"), Query(field_name="test2")],
reranker=reranker,
)
docs_list = [["d1"], ["d2"]]
result = executor._merge_and_rerank(ctx, docs_list)
assert result == ["merged"]
reranker.rerank.assert_called_once_with(docs_list, ctx.topk)
def test_execute_python_pipeline(self):
# Each query is executed serially and converted into a result list.
schema = MockCollectionSchema()
executor = QueryExecutor(schema)
collection = MagicMock()
collection.Query.side_effect = [["raw1"], ["raw2"]]
vectors = [MagicMock(), MagicMock()]
with patch(
"zvec.executor.query_executor.convert_to_py_doc",
side_effect=lambda doc, schema: doc,
):
results = executor._execute_python_pipeline(vectors, collection)
assert results == [["raw1"], ["raw2"]]
assert collection.Query.call_count == 2
def test_build_search_query_by_missing_id_raises_value_error(self):
vector_schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32)
schema = CollectionSchema(name="test_collection", vectors=[vector_schema])
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
collection = MagicMock()
collection.Fetch.return_value = {}
with pytest.raises(ValueError, match="Document with id 'missing' not found"):
executor._build_search_query(
ctx, Query(field_name="test", id="missing"), collection
)
def test_build_search_query_validates_query(self):
vector_schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32)
schema = CollectionSchema(name="test_collection", vectors=[vector_schema])
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
collection = MagicMock()
with pytest.raises(ValueError, match="Cannot provide both id and vector"):
executor._build_search_query(
ctx,
Query(field_name="test", id="doc1", vector=np.array([0.1])),
collection,
)
+948
View File
@@ -0,0 +1,948 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from unittest.mock import patch, MagicMock
import pytest
import os
from zvec import Doc, MetricType, VectorSchema, DataType, FlatIndexParam
from zvec.extension.multi_vector_reranker import (
CallbackReRanker,
RrfReRanker,
WeightedReRanker,
)
from zvec.extension.sentence_transformer_rerank_function import (
DefaultLocalReRanker,
)
from zvec.extension.qwen_rerank_function import QwenReRanker
# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API tests
RUN_INTEGRATION_TESTS = os.environ.get("ZVEC_RUN_INTEGRATION_TESTS", "0") == "1"
# ----------------------------
# RrfReRanker Test Case
# ----------------------------
class TestRrfReRanker:
def test_init(self):
reranker = RrfReRanker(rank_constant=100)
assert reranker.rank_constant == 100
def test_default_rank_constant(self):
reranker = RrfReRanker()
assert reranker.rank_constant == 60
def test_rerank(self):
reranker = RrfReRanker(rank_constant=60)
doc1 = Doc(id="1", score=0.8)
doc2 = Doc(id="2", score=0.7)
doc3 = Doc(id="3", score=0.9)
doc4 = Doc(id="4", score=0.6)
query_results = [[doc1, doc2, doc3], [doc3, doc1, doc4]]
results = reranker.rerank(query_results, topn=3)
assert len(results) <= 3
for doc in results:
assert hasattr(doc, "score")
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True)
# ----------------------------
# WeightedReRanker Test Case
# ----------------------------
class TestWeightedReRanker:
@staticmethod
def _make_fields(metrics):
return [
VectorSchema(
name=f"vector{i}",
data_type=DataType.VECTOR_FP32,
dimension=4,
index_param=FlatIndexParam(metric_type=metric),
)
for i, metric in enumerate(metrics)
]
def test_init(self):
reranker = WeightedReRanker([0.7, 0.3])
assert reranker.weights == [0.7, 0.3]
def test_rerank(self):
reranker = WeightedReRanker([0.7, 0.3])
doc1 = Doc(id="1", score=0.8)
doc2 = Doc(id="2", score=0.7)
doc3 = Doc(id="3", score=0.9)
query_results = [[doc1, doc2], [doc2, doc3]]
fields = self._make_fields([MetricType.L2, MetricType.L2])
results = reranker.rerank(query_results, topn=3, fields=fields)
assert len(results) <= 3
for doc in results:
assert hasattr(doc, "score")
# ----------------------------
# CallbackReRanker Test Case
# ----------------------------
class TestCallbackReRanker:
def test_rerank(self):
def my_callback(query_results, fields, topn):
all_docs = []
for docs in query_results:
all_docs.extend(docs)
all_docs.sort(key=lambda d: d.score, reverse=True)
return all_docs[:topn]
reranker = CallbackReRanker(my_callback)
doc1 = Doc(id="1", score=0.8)
doc2 = Doc(id="2", score=0.9)
doc3 = Doc(id="3", score=0.7)
doc4 = Doc(id="4", score=0.6)
query_results = [[doc1, doc2], [doc3, doc4]]
results = reranker.rerank(query_results, topn=3)
assert len(results) == 3
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True)
def test_callback_with_topn(self):
received_topn = []
def my_callback(query_results, fields, topn):
received_topn.append(topn)
return []
reranker = CallbackReRanker(my_callback)
reranker.rerank([[Doc(id="1", score=0.5)]], topn=7)
assert received_topn == [7]
# ----------------------------
# QwenReRanker Test Case
# ----------------------------
class TestQwenReRanker:
def test_init_without_query(self):
with pytest.raises(ValueError, match="Query is required for QwenReRanker"):
QwenReRanker(api_key="test_key")
def test_init_without_api_key(self):
with patch.dict(os.environ, {}, clear=True):
with pytest.raises(ValueError, match="DashScope API key is required"):
QwenReRanker(query="test")
@patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"})
def test_init_with_env_api_key(self):
reranker = QwenReRanker(query="test", rerank_field="content")
assert reranker.query == "test"
assert reranker._api_key == "test_key"
assert reranker.rerank_field == "content"
def test_init_with_explicit_api_key(self):
reranker = QwenReRanker(
query="test", api_key="explicit_key", rerank_field="content"
)
assert reranker.query == "test"
assert reranker._api_key == "explicit_key"
def test_model_property(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
assert reranker.model == "gte-rerank-v2"
reranker = QwenReRanker(
query="test",
model="custom-model",
api_key="test_key",
rerank_field="content",
)
assert reranker.model == "custom-model"
def test_query_property(self):
reranker = QwenReRanker(
query="test query", api_key="test_key", rerank_field="content"
)
assert reranker.query == "test query"
def test_rerank_field_property(self):
reranker = QwenReRanker(query="test", api_key="test_key", rerank_field="title")
assert reranker.rerank_field == "title"
def test_rerank_empty_results(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
results = reranker.rerank({})
assert results == []
def test_rerank_no_valid_documents(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
# Document without the rerank_field
query_results = {"vector1": [Doc(id="1")]}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
def test_rerank_skip_empty_content(self):
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": ""}),
Doc(id="2", fields={"content": " "}),
]
}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_success(self, mock_require_module):
# Mock dashscope module
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
# Mock API response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.output = {
"results": [
{"index": 0, "relevance_score": 0.95},
{"index": 1, "relevance_score": 0.85},
]
}
mock_dashscope.TextReRank.call.return_value = mock_response
reranker = QwenReRanker(
query="test query", api_key="test_key", rerank_field="content"
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Document 1"}),
Doc(id="2", fields={"content": "Document 2"}),
]
}
results = reranker.rerank(query_results, topn=2)
assert len(results) == 2
assert results[0].id == "1"
assert results[0].score == 0.95
assert results[1].id == "2"
assert results[1].score == 0.85
# Verify API call
mock_dashscope.TextReRank.call.assert_called_once_with(
model="gte-rerank-v2",
query="test query",
documents=["Document 1", "Document 2"],
top_n=2,
return_documents=False,
)
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_deduplicate_documents(self, mock_require_module):
# Mock dashscope module
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
# Mock API response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.output = {
"results": [
{"index": 0, "relevance_score": 0.9},
]
}
mock_dashscope.TextReRank.call.return_value = mock_response
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
# Same document in multiple vector results
doc1 = Doc(id="1", fields={"content": "Document 1"})
query_results = {"vector1": [doc1], "vector2": [doc1]}
results = reranker.rerank(query_results, topn=5)
# Should only call API with document once
call_args = mock_dashscope.TextReRank.call.call_args
assert len(call_args[1]["documents"]) == 1
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_api_error(self, mock_require_module):
# Mock dashscope module
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
# Mock API error response
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.message = "Invalid request"
mock_response.code = "InvalidParameter"
mock_dashscope.TextReRank.call.return_value = mock_response
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
with pytest.raises(ValueError, match="DashScope API error"):
reranker.rerank(query_results)
@patch("zvec.extension.qwen_function.require_module")
def test_rerank_runtime_error(self, mock_require_module):
# Mock dashscope module that raises exception
mock_dashscope = MagicMock()
mock_require_module.return_value = mock_dashscope
mock_dashscope.TextReRank.call.side_effect = Exception("Network error")
reranker = QwenReRanker(
query="test", api_key="test_key", rerank_field="content"
)
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
with pytest.raises(RuntimeError, match="Failed to call DashScope API"):
reranker.rerank(query_results)
@pytest.mark.skipif(
not RUN_INTEGRATION_TESTS,
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
)
def test_real_qwen_rerank(self):
"""Integration test with real DashScope TextReRank API.
To run this test, set environment variables:
export ZVEC_RUN_INTEGRATION_TESTS=1
export DASHSCOPE_API_KEY=your-api-key
"""
# Create reranker with real API
reranker = QwenReRanker(
query="What is machine learning?",
rerank_field="content",
model="gte-rerank-v2",
)
# Prepare test documents
query_results = {
"vector1": [
Doc(
id="1",
score=0.8,
fields={
"content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data."
},
),
Doc(
id="2",
score=0.7,
fields={
"content": "The weather is nice today with clear skies and sunshine."
},
),
Doc(
id="3",
score=0.75,
fields={
"content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers."
},
),
],
"vector2": [
Doc(
id="4",
score=0.6,
fields={
"content": "Python is a popular programming language for data science and machine learning applications."
},
),
Doc(
id="5",
score=0.65,
fields={
"content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder."
},
),
],
}
# Call real API
results = reranker.rerank(query_results, topn=3)
# Verify results
assert len(results) <= 3, "Should return at most topn documents"
assert len(results) > 0, "Should return at least one document"
# All results should have valid scores
for doc in results:
assert hasattr(doc, "score"), "Each document should have a score"
assert isinstance(doc.score, (int, float)), "Score should be numeric"
assert doc.score > 0, "Score should be positive"
# Verify scores are in descending order
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True), (
"Results should be sorted by score in descending order"
)
# Verify relevant documents are ranked higher
# Document 1 and 3 are about machine learning, should rank higher than weather/recipe docs
result_ids = [doc.id for doc in results]
# At least one of the ML-related documents should be in top results
ml_related_docs = {"1", "3", "4"}
assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (
"ML-related documents should rank higher"
)
# Print results for manual verification (useful during development)
print("\nReranking results:")
for i, doc in enumerate(results, 1):
print(f"{i}. ID={doc.id}, Score={doc.score:.4f}")
if doc.fields:
content = doc.field("content")
if content:
print(f" Content: {content[:80]}...")
# ----------------------------
# DefaultLocalReRanker Test Case
# ----------------------------
class TestDefaultLocalReRanker:
"""Test cases for DefaultLocalReRanker."""
def test_init_without_query(self):
"""Test initialization fails without query."""
with pytest.raises(
ValueError, match="Query is required for DefaultLocalReRanker"
):
DefaultLocalReRanker(rerank_field="content")
def test_init_with_empty_query(self):
"""Test initialization fails with empty query."""
with pytest.raises(
ValueError, match="Query is required for DefaultLocalReRanker"
):
DefaultLocalReRanker(query="", rerank_field="content")
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
def test_init_success(self, mock_require_module):
"""Test successful initialization with mocked model."""
# Mock sentence_transformers module
mock_st = MagicMock()
mock_model = MagicMock()
mock_model.predict = MagicMock() # Cross-encoder has predict method
mock_model.device = "cpu"
mock_st.CrossEncoder.return_value = mock_model
mock_require_module.return_value = mock_st
reranker = DefaultLocalReRanker(
query="test query",
rerank_field="content",
model_name="cross-encoder/ms-marco-MiniLM-L6-v2",
)
assert reranker.query == "test query"
assert reranker.rerank_field == "content"
assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L6-v2"
assert reranker.model_source == "huggingface"
assert reranker.batch_size == 32
@pytest.mark.skipif(
not RUN_INTEGRATION_TESTS,
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
)
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
def test_init_with_custom_params(self, mock_require_module):
"""Test initialization with custom parameters."""
mock_st = MagicMock()
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_model.device = "cuda"
mock_st.CrossEncoder.return_value = mock_model
mock_require_module.return_value = mock_st
reranker = DefaultLocalReRanker(
query="custom query",
rerank_field="title",
model_name="cross-encoder/ms-marco-MiniLM-L12-v2",
model_source="modelscope",
device="cuda",
batch_size=64,
)
assert reranker.query == "custom query"
assert reranker.rerank_field == "title"
assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L12-v2"
assert reranker.model_source == "modelscope"
assert reranker.batch_size == 64
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
def test_init_invalid_model(self, mock_require_module):
"""Test initialization fails with non-cross-encoder model."""
# Mock a model without predict method (not a cross-encoder)
mock_st = MagicMock()
mock_model = MagicMock(spec=[]) # No predict method
mock_st.CrossEncoder.return_value = mock_model
mock_require_module.return_value = mock_st
with pytest.raises(ValueError, match="does not appear to be a cross-encoder"):
DefaultLocalReRanker(query="test", rerank_field="content")
def test_query_property(self):
"""Test query property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test query", rerank_field="content")
assert reranker.query == "test query"
def test_rerank_field_property(self):
"""Test rerank_field property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="title")
assert reranker.rerank_field == "title"
def test_batch_size_property(self):
"""Test batch_size property."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", rerank_field="content", batch_size=128
)
assert reranker.batch_size == 128
def test_rerank_empty_results(self):
"""Test rerank with empty query_results."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
results = reranker.rerank({})
assert results == []
def test_rerank_no_valid_documents(self):
"""Test rerank with documents missing rerank_field."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
# Document without the rerank_field
query_results = {"vector1": [Doc(id="1")]}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
def test_rerank_skip_empty_content(self):
"""Test rerank skips documents with empty content."""
mock_model = MagicMock()
mock_model.predict = MagicMock()
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {
"vector1": [
Doc(id="1", fields={"content": ""}),
Doc(id="2", fields={"content": " "}),
]
}
with pytest.raises(ValueError, match="No documents to rerank"):
reranker.rerank(query_results)
def test_rerank_success(self):
"""Test successful rerank with mocked model."""
# Mock standard cross-encoder model
mock_model = MagicMock()
# Mock predict method to return scores
import numpy as np
mock_scores = np.array([0.95, 0.85, 0.75])
mock_model.predict.return_value = mock_scores
mock_model.device = "cpu"
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test query", rerank_field="content")
query_results = {
"vector1": [
Doc(id="1", score=0.8, fields={"content": "Document 1"}),
Doc(id="2", score=0.7, fields={"content": "Document 2"}),
Doc(id="3", score=0.6, fields={"content": "Document 3"}),
]
}
results = reranker.rerank(query_results, topn=3)
# Verify results
assert len(results) == 3
assert results[0].id == "1"
assert results[0].score == 0.95
assert results[1].id == "2"
assert results[1].score == 0.85
assert results[2].id == "3"
assert results[2].score == 0.75
# Verify model.predict was called correctly
assert mock_model.predict.called
call_args = mock_model.predict.call_args
pairs = call_args[0][0]
assert len(pairs) == 3
assert pairs[0] == ["test query", "Document 1"]
assert pairs[1] == ["test query", "Document 2"]
assert pairs[2] == ["test query", "Document 3"]
assert call_args[1]["batch_size"] == 32
assert call_args[1]["show_progress_bar"] is False
def test_rerank_with_topn_limit(self):
"""Test rerank respects topn limit."""
mock_model = MagicMock()
import numpy as np
mock_scores = np.array([0.9, 0.8, 0.7, 0.6, 0.5])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Doc 1"}),
Doc(id="2", fields={"content": "Doc 2"}),
Doc(id="3", fields={"content": "Doc 3"}),
Doc(id="4", fields={"content": "Doc 4"}),
Doc(id="5", fields={"content": "Doc 5"}),
]
}
results = reranker.rerank(query_results, topn=2)
# Should only return top 2
assert len(results) == 2
assert results[0].id == "1"
assert results[0].score == 0.9
assert results[1].id == "2"
assert results[1].score == 0.8
def test_rerank_deduplicate_documents(self):
"""Test rerank deduplicates documents across multiple vectors."""
mock_model = MagicMock()
import numpy as np
mock_scores = np.array([0.95, 0.85])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
# Same document in multiple vector results
doc1 = Doc(id="1", fields={"content": "Document 1"})
doc2 = Doc(id="2", fields={"content": "Document 2"})
query_results = {
"vector1": [doc1, doc2],
"vector2": [doc1], # doc1 appears in both
}
results = reranker.rerank(query_results, topn=5)
# Should only process each document once
assert len(results) == 2
assert mock_model.predict.call_count == 1
call_args = mock_model.predict.call_args
pairs = call_args[0][0]
assert len(pairs) == 2 # Only 2 unique documents
def test_rerank_sorting(self):
"""Test rerank sorts documents by score in descending order."""
mock_model = MagicMock()
import numpy as np
# Return scores in non-sorted order
mock_scores = np.array([0.6, 0.9, 0.7])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Doc 1"}),
Doc(id="2", fields={"content": "Doc 2"}),
Doc(id="3", fields={"content": "Doc 3"}),
]
}
results = reranker.rerank(query_results, topn=3)
# Should be sorted by score (descending)
assert len(results) == 3
assert results[0].id == "2" # score 0.9
assert results[0].score == 0.9
assert results[1].id == "3" # score 0.7
assert results[1].score == 0.7
assert results[2].id == "1" # score 0.6
assert results[2].score == 0.6
def test_rerank_model_error(self):
"""Test rerank handles model prediction errors."""
mock_model = MagicMock()
# Mock predict to raise exception
mock_model.predict.side_effect = Exception("Model inference error")
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
with pytest.raises(RuntimeError, match="Failed to compute rerank scores"):
reranker.rerank(query_results)
def test_rerank_with_custom_batch_size(self):
"""Test rerank uses custom batch_size."""
mock_model = MagicMock()
import numpy as np
mock_scores = np.array([0.9, 0.8])
mock_model.predict.return_value = mock_scores
# Mock sentence_transformers module
mock_st = MagicMock()
mock_st.CrossEncoder.return_value = mock_model
with patch(
"zvec.extension.sentence_transformer_rerank_function.require_module",
return_value=mock_st,
):
reranker = DefaultLocalReRanker(
query="test", rerank_field="content", batch_size=64
)
query_results = {
"vector1": [
Doc(id="1", fields={"content": "Doc 1"}),
Doc(id="2", fields={"content": "Doc 2"}),
]
}
reranker.rerank(query_results)
# Verify batch_size is passed to predict
call_args = mock_model.predict.call_args
assert call_args[1]["batch_size"] == 64
@pytest.mark.skipif(
not RUN_INTEGRATION_TESTS,
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
)
def test_real_sentence_transformer_rerank(self):
"""Integration test with real SentenceTransformer cross-encoder model.
To run this test, set environment variable:
export ZVEC_RUN_INTEGRATION_TESTS=1
Note: This test requires sentence-transformers package and will
download the MS MARCO MiniLM model (~80MB) on first run.
"""
# Create reranker with real model (using default lightweight model)
reranker = DefaultLocalReRanker(
query="What is machine learning?",
rerank_field="content",
)
# Prepare test documents
query_results = {
"vector1": [
Doc(
id="1",
score=0.8,
fields={
"content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data."
},
),
Doc(
id="2",
score=0.7,
fields={
"content": "The weather is nice today with clear skies and sunshine."
},
),
Doc(
id="3",
score=0.75,
fields={
"content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers."
},
),
],
"vector2": [
Doc(
id="4",
score=0.6,
fields={
"content": "Python is a popular programming language for data science and machine learning applications."
},
),
Doc(
id="5",
score=0.65,
fields={
"content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder."
},
),
],
}
# Call real model
results = reranker.rerank(query_results, topn=3)
# Verify results
assert len(results) <= 3, "Should return at most topn documents"
assert len(results) > 0, "Should return at least one document"
# All results should have valid scores
for doc in results:
assert hasattr(doc, "score"), "Each document should have a score"
assert isinstance(doc.score, (int, float)), "Score should be numeric"
# Verify scores are in descending order
scores = [doc.score for doc in results]
assert scores == sorted(scores, reverse=True), (
"Results should be sorted by score in descending order"
)
# Verify relevant documents are ranked higher
# Documents 1, 3, and 4 are about machine learning, should rank higher
result_ids = [doc.id for doc in results]
# At least one of the ML-related documents should be in top results
ml_related_docs = {"1", "3", "4"}
assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (
"ML-related documents should rank higher"
)
# Print results for manual verification (useful during development)
print("\nSentenceTransformer Reranking results:")
for i, doc in enumerate(results, 1):
print(f"{i}. ID={doc.id}, Score={doc.score:.4f}")
if doc.fields:
content = doc.field("content")
if content:
print(f" Content: {content[:80]}...")
+246
View File
@@ -0,0 +1,246 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import pytest
from zvec import (
CollectionSchema,
CollectionStats,
FieldSchema,
VectorSchema,
HnswIndexParam,
InvertIndexParam,
DataType,
IndexType,
MetricType,
)
# ----------------------------
# FieldSchema Test Case
# ----------------------------
class TestFieldSchema:
def test_default(self):
field = FieldSchema("field", data_type=DataType.FLOAT)
assert field.name == "field"
assert field.data_type == DataType.FLOAT
assert field.nullable is False
assert field.index_param is None
def test_custom(self):
field_1 = FieldSchema(
name="float",
data_type=DataType.FLOAT,
nullable=True,
index_param=InvertIndexParam(),
)
assert field_1.name == "float"
assert field_1.data_type == DataType.FLOAT
assert field_1.nullable is True
assert field_1.index_param.enable_range_optimization is False
field_2 = FieldSchema(
name="str",
data_type=DataType.STRING,
nullable=True,
index_param=InvertIndexParam(enable_range_optimization=True),
)
assert field_2.name == "str"
assert field_2.data_type == DataType.STRING
assert field_2.nullable is True
assert field_2.index_param.enable_range_optimization is True
def test_readonly(self):
field = FieldSchema(
name="float",
data_type=DataType.FLOAT,
nullable=True,
index_param=InvertIndexParam(),
)
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
field.index_param = InvertIndexParam(enable_range_optimization=True)
# ----------------------------
# VectorSchema Test Case
# ----------------------------
class TestVectorSchema:
def test_default(self):
field = VectorSchema("vector", data_type=DataType.VECTOR_FP32, dimension=128)
assert field.name == "vector"
assert field.data_type == DataType.VECTOR_FP32
assert field.dimension == 128
assert field.index_param is not None
assert field.index_param.type == IndexType.FLAT
assert field.index_param.metric_type == MetricType.IP
def test_custom(self):
field = VectorSchema(
name="vector",
data_type=DataType.VECTOR_INT8,
dimension=512,
index_param=HnswIndexParam(
metric_type=MetricType.COSINE, m=15, ef_construction=300
),
)
assert field.name == "vector"
assert field.data_type == DataType.VECTOR_INT8
assert field.index_param.metric_type == MetricType.COSINE
assert field.index_param.m == 15
assert field.index_param.ef_construction == 300
def test_readonly(self):
field = VectorSchema(
name="vector",
dimension=128,
data_type=DataType.VECTOR_INT8,
)
import sys
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
field.dimension = 4
# ----------------------------
# CollectionSchema Test Case
# ----------------------------
class TestCollectionSchema:
def test_collection_schema_with_single_field(self):
collection_schema = CollectionSchema(
name="test_collection",
fields=FieldSchema(
name="id",
data_type=DataType.INT64,
index_param=InvertIndexParam(),
nullable=False,
),
vectors=VectorSchema(
name="vector",
data_type=DataType.VECTOR_INT8,
dimension=128,
index_param=HnswIndexParam(),
),
)
assert collection_schema is not None
assert collection_schema.name == "test_collection"
assert len(collection_schema.fields) == 1
assert len(collection_schema.vectors) == 1
field = collection_schema.field("id")
assert field is not None
assert field.name == "id"
assert field.data_type == DataType.INT64
assert not field.nullable
assert field.index_param.type == IndexType.INVERT
assert not field.index_param.enable_range_optimization
vector = collection_schema.vector("vector")
assert vector is not None
assert vector.name == "vector"
assert vector.data_type == DataType.VECTOR_INT8
assert vector.dimension == 128
assert vector.index_param.type == IndexType.HNSW
assert vector.index_param.m == 50
assert vector.index_param.ef_construction == 500
assert vector.index_param.metric_type == MetricType.IP
def test_collection_schema_with_multi_fields(self):
collection_schema = 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.INT32,
nullable=True,
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
VectorSchema(
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
),
],
)
assert collection_schema is not None
assert collection_schema.name == "test_collection"
assert len(collection_schema.fields) == 3
assert len(collection_schema.vectors) == 2
field_id = collection_schema.field("id")
assert field_id is not None
assert field_id.name == "id"
assert field_id.data_type == DataType.INT64
assert not field_id.nullable
assert field_id.index_param.type == IndexType.INVERT
dense = collection_schema.vector("dense")
assert dense is not None
assert dense.name == "dense"
assert dense.data_type == DataType.VECTOR_FP32
assert dense.dimension == 128
assert dense.index_param.type == IndexType.HNSW
sparse = collection_schema.vector("sparse")
assert sparse is not None
assert sparse.name == "sparse"
assert sparse.data_type == DataType.SPARSE_VECTOR_FP32
assert sparse.dimension == 0
assert sparse.index_param.type == IndexType.HNSW
assert str(collection_schema) is not None
# ----------------------------
# CollectionStats Test Case
# ----------------------------
class TestCollectionStats:
"""
The constructor of CollectionStats is not provided.
It can only be obtained through collection.stats()
"""
def test_collection_stats(self):
stats = CollectionStats()
assert stats is not None
+141
View File
@@ -0,0 +1,141 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import pytest
from zvec import (
DataType,
IndexType,
MetricType,
QuantizeType,
Status,
StatusCode,
)
# ----------------------------
# Enum Test Case
# ----------------------------
@pytest.mark.parametrize(
"member, name",
[
(DataType.FLOAT, "FLOAT"),
(IndexType.HNSW, "HNSW"),
(MetricType.COSINE, "COSINE"),
(QuantizeType.INT8, "INT8"),
(StatusCode.OK, "OK"),
],
)
def test_enum_names(member, name):
assert member.name == name
@pytest.mark.parametrize(
"member, value",
[
(DataType.FLOAT, 8),
(IndexType.HNSW, 1),
(MetricType.COSINE, 3),
(QuantizeType.INT8, 2),
(StatusCode.OK, 0),
],
)
def test_enum_values(member, value):
assert member.value == value
@pytest.mark.parametrize("member", ["L2", "IP", "COSINE"])
def test_metric_type_has_member(member):
assert member in MetricType.__members__
@pytest.mark.parametrize(
"member",
[
"STRING",
"BOOL",
"INT32",
"INT64",
"FLOAT",
"DOUBLE",
"UINT32",
"UINT64",
"VECTOR_FP16",
"VECTOR_FP32",
"VECTOR_FP64",
"VECTOR_INT8",
"SPARSE_VECTOR_FP32",
"SPARSE_VECTOR_FP16",
"ARRAY_STRING",
"ARRAY_INT32",
"ARRAY_INT64",
"ARRAY_FLOAT",
"ARRAY_DOUBLE",
"ARRAY_BOOL",
"ARRAY_UINT32",
"ARRAY_UINT64",
],
)
def test_data_type_has_member(member):
assert member in DataType.__members__
@pytest.mark.parametrize("member", ["HNSW", "IVF", "FLAT", "INVERT"])
def test_index_type_has_member(member):
assert member in IndexType.__members__
@pytest.mark.parametrize("member", ["FP16", "INT8", "INT4", "UNDEFINED"])
def test_quantize_type_has_member(member):
assert member in QuantizeType.__members__
@pytest.mark.parametrize(
"member",
[
"OK",
"UNKNOWN",
"NOT_FOUND",
"ALREADY_EXISTS",
"INVALID_ARGUMENT",
"PERMISSION_DENIED",
"FAILED_PRECONDITION",
"RESOURCE_EXHAUSTED",
"UNAVAILABLE",
"INTERNAL_ERROR",
"NOT_SUPPORTED",
],
)
def test_status_code_has_member(member):
assert member in StatusCode.__members__
# ----------------------------
# Status Test Case
# ----------------------------
class TestStatus:
def test_status_code(self):
status = Status(StatusCode.OK)
assert status.code() == StatusCode.OK
def test_status_message(self):
status = Status(StatusCode.OK, "OK")
assert status.message() == "OK"
status = Status(StatusCode.NOT_FOUND, "Not Found")
assert status.message() == "Not Found"
def test_status_ok(self):
status = Status(StatusCode.OK)
assert status.ok()
+89
View File
@@ -0,0 +1,89 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from unittest.mock import MagicMock, patch
import pytest
from zvec import require_module
# ----------------------------
# require_module func Test Case
# ----------------------------
def test_require_module_success():
module = require_module("os")
assert module is not None
assert hasattr(module, "path")
def test_require_module_with_submodule_success():
module = require_module("os.path")
assert module is not None
assert hasattr(module, "join")
def test_require_module_import_error():
with pytest.raises(ImportError) as exc_info:
require_module("nonexistent_module")
exception_msg = str(exc_info.value)
assert "Required package 'nonexistent_module' is not installed." in exception_msg
def test_require_module_with_mitigation_import_error():
with pytest.raises(ImportError) as exc_info:
require_module("nonexistent_module.submodule", mitigation="custom_package")
exception_msg = str(exc_info.value)
assert "Required package 'custom_package' is not installed." in exception_msg
assert (
"Module 'nonexistent_module.submodule' is part of 'nonexistent_module'"
in exception_msg
)
assert "please pip install 'custom_package'." in exception_msg
def test_require_module_submodule_import_error():
with pytest.raises(ImportError) as exc_info:
require_module("os.nonexistent_submodule")
exception_msg = str(exc_info.value)
assert (
"Required package 'os.nonexistent_submodule' is not installed." in exception_msg
)
assert "Module 'os.nonexistent_submodule' is part of 'os'" in exception_msg
assert "please pip install 'os'." in exception_msg
@patch("importlib.import_module")
def test_require_module_wraps_original_exception(mock_import_module):
original_exception = ImportError("Original error")
mock_import_module.side_effect = original_exception
with pytest.raises(ImportError) as exc_info:
require_module("some_module")
assert exc_info.value.__cause__ is original_exception
@patch("importlib.import_module")
def test_require_module_calls_importlib(mock_import_module):
mock_module = MagicMock()
mock_import_module.return_value = mock_module
result = require_module("test_module")
mock_import_module.assert_called_once_with("test_module")
assert result is mock_module
+501
View File
@@ -0,0 +1,501 @@
# 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.
"""
Tests for the Python entry point of the Vamana (DiskANN) dense vector index.
Mirrors the structure of ``test_hnsw_contiguous_memory.py`` (the closest
hnsw dense reference), and is split into two parts:
1. **Surface tests** — verify that ``VamanaIndexParam`` / ``VamanaQueryParam``
are correctly bound: construction defaults, readonly properties,
``to_dict``, ``__repr__``, pickle round-trip, and that they appear in the
public ``zvec`` namespace with the expected ``IndexType.VAMANA`` value.
2. **End-to-end tests** — build a collection that uses Vamana on a dense
FP32 column, insert deterministic documents, then run a top-k query
through ``VamanaQueryParam`` on both the writer segment and the
persisted (post-``optimize()``) segment.
"""
from __future__ import annotations
import pickle
import sys
import numpy as np
import pytest
import zvec
from zvec import (
Collection,
CollectionOption,
CollectionSchema,
Doc,
FieldSchema,
InvertIndexParam,
VamanaIndexParam,
VamanaQueryParam,
Query,
VectorSchema,
)
from zvec.typing import DataType, IndexType, MetricType, QuantizeType
DIMENSION = 32
NUM_DOCS = 128
TOPK = 5
# Defaults pulled from src/include/zvec/core/interface/constants.h. Keep
# in sync with kDefaultVamana* if the engine defaults ever change.
DEFAULT_MAX_DEGREE = 64
DEFAULT_SEARCH_LIST_SIZE = 100
DEFAULT_ALPHA = 1.2
DEFAULT_EF_SEARCH = 200
DEFAULT_SATURATE_GRAPH = False
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _build_schema(
name: str,
*,
metric_type: MetricType = MetricType.IP,
max_degree: int = 32,
search_list_size: int = 64,
alpha: float = 1.2,
use_contiguous_memory: bool = False,
) -> CollectionSchema:
"""Create a simple schema with a single FP32 Vamana vector column."""
return CollectionSchema(
name=name,
fields=[
FieldSchema(
"id",
DataType.INT64,
nullable=False,
index_param=InvertIndexParam(enable_range_optimization=True),
),
],
vectors=[
VectorSchema(
"dense",
DataType.VECTOR_FP32,
dimension=DIMENSION,
index_param=VamanaIndexParam(
metric_type=metric_type,
max_degree=max_degree,
search_list_size=search_list_size,
alpha=alpha,
use_contiguous_memory=use_contiguous_memory,
),
),
],
)
def _generate_docs(rng: np.random.Generator, num: int = NUM_DOCS) -> list[Doc]:
"""Produce deterministic documents for insertion."""
docs: list[Doc] = []
for i in range(num):
vec = rng.standard_normal(DIMENSION).astype(np.float32)
docs.append(
Doc(
id=str(i),
fields={"id": i},
vectors={"dense": vec.tolist()},
)
)
return docs
def _query_topk(
coll: Collection, query_vec: list[float], *, ef_search: int = 64
) -> list[str]:
"""Run a top-k vector query and return the returned ids in order."""
vector_query = Query(
field_name="dense",
vector=query_vec,
param=VamanaQueryParam(ef_search=ef_search),
)
hits = coll.query(vector_query, topk=TOPK)
assert hits is not None, "query returned None"
assert len(hits) >= 1, f"expected at least one hit, got {hits!r}"
return [doc.id for doc in hits]
# ---------------------------------------------------------------------------
# 1) Surface: construction / property / to_dict / repr / pickle / namespace
# ---------------------------------------------------------------------------
class TestVamanaIndexParamSurface:
"""Verify the Python binding for ``VamanaIndexParam``."""
def test_defaults(self):
param = VamanaIndexParam()
assert param.type == IndexType.VAMANA
assert param.metric_type == MetricType.IP
assert param.max_degree == DEFAULT_MAX_DEGREE
assert param.search_list_size == DEFAULT_SEARCH_LIST_SIZE
assert param.alpha == pytest.approx(DEFAULT_ALPHA)
assert param.saturate_graph is DEFAULT_SATURATE_GRAPH
assert param.use_contiguous_memory is False
assert param.use_id_map is False
assert param.quantize_type == QuantizeType.UNDEFINED
def test_custom_construction(self):
param = VamanaIndexParam(
metric_type=MetricType.COSINE,
max_degree=48,
search_list_size=128,
alpha=1.5,
saturate_graph=True,
use_contiguous_memory=True,
use_id_map=False,
quantize_type=QuantizeType.INT8,
)
assert param.type == IndexType.VAMANA
assert param.metric_type == MetricType.COSINE
assert param.max_degree == 48
assert param.search_list_size == 128
assert param.alpha == pytest.approx(1.5)
assert param.saturate_graph is True
assert param.use_contiguous_memory is True
assert param.use_id_map is False
assert param.quantize_type == QuantizeType.INT8
def test_to_dict_includes_all_fields(self):
param = VamanaIndexParam(
metric_type=MetricType.L2,
max_degree=32,
search_list_size=80,
alpha=1.3,
saturate_graph=True,
use_contiguous_memory=True,
use_id_map=False,
quantize_type=QuantizeType.FP16,
)
data = param.to_dict()
assert data["type"] == "VAMANA"
assert data["metric_type"] == "L2"
assert data["max_degree"] == 32
assert data["search_list_size"] == 80
assert data["alpha"] == pytest.approx(1.3)
assert data["saturate_graph"] is True
assert data["use_contiguous_memory"] is True
assert data["use_id_map"] is False
assert data["quantize_type"] == "FP16"
def test_repr_contains_key_fields(self):
text = repr(
VamanaIndexParam(
metric_type=MetricType.COSINE,
max_degree=24,
search_list_size=72,
alpha=1.4,
saturate_graph=True,
use_contiguous_memory=True,
)
)
# Spot-check the most diagnostic fields are rendered.
assert "VAMANA" in text
assert "COSINE" in text
assert "max_degree" in text and "24" in text
assert "search_list_size" in text and "72" in text
assert "alpha" in text
assert "saturate_graph" in text and "true" in text
assert "use_contiguous_memory" in text and "true" in text
@pytest.mark.parametrize(
"field, kwargs",
[
("max_degree", dict(max_degree=99)),
("search_list_size", dict(search_list_size=99)),
("alpha", dict(alpha=1.7)),
("saturate_graph", dict(saturate_graph=True)),
("use_contiguous_memory", dict(use_contiguous_memory=True)),
("use_id_map", dict(use_id_map=True)),
],
)
def test_readonly_properties(self, field, kwargs):
param = VamanaIndexParam(**kwargs)
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
setattr(param, field, getattr(param, field))
def test_pickle_roundtrip(self):
original = VamanaIndexParam(
metric_type=MetricType.COSINE,
max_degree=48,
search_list_size=120,
alpha=1.4,
saturate_graph=True,
use_contiguous_memory=True,
use_id_map=False,
quantize_type=QuantizeType.INT8,
)
restored = pickle.loads(pickle.dumps(original))
assert restored.type == IndexType.VAMANA
assert restored.metric_type == MetricType.COSINE
assert restored.max_degree == 48
assert restored.search_list_size == 120
assert restored.alpha == pytest.approx(1.4)
assert restored.saturate_graph is True
assert restored.use_contiguous_memory is True
assert restored.use_id_map is False
assert restored.quantize_type == QuantizeType.INT8
# to_dict equality is the strongest end-to-end equivalence we have.
assert restored.to_dict() == original.to_dict()
class TestVamanaQueryParamSurface:
"""Verify the Python binding for ``VamanaQueryParam``."""
def test_defaults(self):
q = VamanaQueryParam()
assert q.type == IndexType.VAMANA
assert q.ef_search == DEFAULT_EF_SEARCH
assert q.radius == pytest.approx(0.0)
assert q.is_linear is False
assert q.is_using_refiner is False
assert q.prefetch_offset == 8
assert q.prefetch_lines == 0
def test_custom_construction(self):
q = VamanaQueryParam(
ef_search=300,
radius=0.5,
is_linear=True,
is_using_refiner=True,
extra_params={
"prefetch_offset": 8,
"prefetch_lines": 2,
},
)
assert q.type == IndexType.VAMANA
assert q.ef_search == 300
assert q.radius == pytest.approx(0.5)
assert q.is_linear is True
assert q.is_using_refiner is True
assert q.prefetch_offset == 8
assert q.prefetch_lines == 2
def test_repr_contains_key_fields(self):
text = repr(VamanaQueryParam(ef_search=128, radius=0.25))
assert "VAMANA" in text
assert "ef_search" in text and "128" in text
assert "radius" in text
def test_readonly_ef_search(self):
q = VamanaQueryParam(ef_search=100)
if sys.version_info >= (3, 11):
match_pattern = r"(can't set attribute|has no setter|readonly attribute)"
else:
match_pattern = r"can't set attribute"
with pytest.raises(AttributeError, match=match_pattern):
q.ef_search = 200 # type: ignore[misc]
def test_pickle_roundtrip(self):
original = VamanaQueryParam(
ef_search=256,
radius=0.3,
is_linear=False,
is_using_refiner=True,
extra_params={
"prefetch_offset": 4,
"prefetch_lines": 3,
},
)
restored = pickle.loads(pickle.dumps(original))
assert restored.type == IndexType.VAMANA
assert restored.ef_search == 256
assert restored.radius == pytest.approx(0.3)
assert restored.is_linear is False
assert restored.is_using_refiner is True
assert restored.prefetch_offset == 4
assert restored.prefetch_lines == 3
class TestVamanaPublicNamespace:
"""The Vamana entry points must be importable from the top-level ``zvec``."""
def test_top_level_exports(self):
assert zvec.VamanaIndexParam is VamanaIndexParam
assert zvec.VamanaQueryParam is VamanaQueryParam
assert "VamanaIndexParam" in zvec.__all__
assert "VamanaQueryParam" in zvec.__all__
def test_index_type_enum_member(self):
# Sanity: the IndexType enum exposes VAMANA and it is what the
# bound params advertise.
assert IndexType.VAMANA is not None
assert VamanaIndexParam().type == IndexType.VAMANA
assert VamanaQueryParam().type == IndexType.VAMANA
# ---------------------------------------------------------------------------
# 2) End-to-end: create collection, insert, query through the writer segment
# ---------------------------------------------------------------------------
@pytest.fixture
def rng() -> np.random.Generator:
return np.random.default_rng(seed=42)
# Mirror the hnsw dense test fixture: only the mmap-backed variant is
# currently usable for vector index construction. BufferPool (enable_mmap=
# False) is intentionally omitted because the same write-path guard in
# ``SegmentImpl::merge_vector_indexer`` rejects that combination.
@pytest.fixture(params=[True], ids=["mmap_on"])
def collection_option(request) -> CollectionOption:
return CollectionOption(read_only=False, enable_mmap=request.param)
class TestVamanaEndToEnd:
"""End-to-end: schema -> create_and_open -> insert -> query works."""
def test_schema_round_trip(self, tmp_path_factory, collection_option):
"""The Vamana index params survive the schema persist path."""
schema = _build_schema(
"vamana_schema_rt",
metric_type=MetricType.COSINE,
max_degree=32,
search_list_size=80,
alpha=1.3,
use_contiguous_memory=True,
)
path = tmp_path_factory.mktemp("zvec") / "vamana_schema_rt"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
vec_schema = coll.schema.vectors[0]
ip = vec_schema.index_param
assert ip.type == IndexType.VAMANA
assert ip.metric_type == MetricType.COSINE
assert ip.max_degree == 32
assert ip.search_list_size == 80
assert ip.alpha == pytest.approx(1.3)
assert ip.use_contiguous_memory is True
finally:
coll.destroy()
def test_insert_and_query_self_recall(
self, tmp_path_factory, collection_option, rng
):
"""Top-1 of a query equal to an inserted vector must be that vector.
Exercises the writer-segment Vamana streamer end-to-end through the
Python entry point: ``VamanaIndexParam`` for build and
``VamanaQueryParam`` for search.
"""
schema = _build_schema("vamana_e2e_recall")
path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_recall"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
docs = _generate_docs(rng)
for r in coll.insert(docs=docs):
assert r.ok(), f"insert failed: code={r.code()}"
assert coll.stats.doc_count == NUM_DOCS
# Self-recall: query with the i-th inserted vector, expect id i
# to be the top result.
for probe in (0, 7, 42, NUM_DOCS - 1):
query_vec = docs[probe].vector("dense")
ids = _query_topk(coll, query_vec)
assert ids[0] == str(probe), (
f"expected self-recall at probe={probe}, got top-1 id={ids[0]} "
f"(top-{TOPK}={ids})"
)
finally:
coll.destroy()
def test_query_param_ef_search_affects_only_quality(
self, tmp_path_factory, collection_option, rng
):
"""``ef_search`` is a search-time knob and must not crash for any
sensible value. Larger ``ef_search`` should be at least as good as
smaller for self-recall."""
schema = _build_schema("vamana_e2e_ef")
path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_ef"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
docs = _generate_docs(rng)
for r in coll.insert(docs=docs):
assert r.ok()
query_vec = docs[3].vector("dense")
ids_small = _query_topk(coll, query_vec, ef_search=16)
ids_large = _query_topk(coll, query_vec, ef_search=256)
# Both should self-recall the probe vector at top-1.
assert ids_small[0] == "3"
assert ids_large[0] == "3"
assert len(ids_small) == TOPK
assert len(ids_large) == TOPK
finally:
coll.destroy()
def test_optimize_then_query(self, tmp_path_factory, collection_option, rng):
"""The persisted Vamana segment built by ``optimize()`` must serve
queries correctly.
Until the cmake fix to force-load ``core_knn_vamana_static`` into the
``_zvec`` pybind module, this path failed at ``VamanaStreamer``
creation because the global factory registration in
``vamana_streamer.cc`` was never linked in. This test pins down the
regression.
"""
schema = _build_schema("vamana_e2e_optimize")
path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_optimize"
coll = zvec.create_and_open(
path=str(path), schema=schema, option=collection_option
)
try:
docs = _generate_docs(rng)
for r in coll.insert(docs=docs):
assert r.ok()
assert coll.stats.doc_count == NUM_DOCS
# Snapshot the writer-segment top-k for a probe vector.
query_vec = docs[5].vector("dense")
ids_pre = _query_topk(coll, query_vec)
assert ids_pre[0] == "5"
# Trigger persisted segment build. Pre-fix this raised
# RuntimeError("Failed to create index").
coll.optimize()
# Persisted segment must still serve queries with the same
# top-1 self-recall guarantee. We do not assert full top-k
# equality with the writer segment because the persisted
# streamer may visit nodes in a different order; top-1 self-
# recall is the strong invariant.
ids_post = _query_topk(coll, query_vec)
assert ids_post[0] == "5", (
f"post-optimize top-1 should still be probe id, got {ids_post}"
)
assert len(ids_post) == TOPK
finally:
coll.destroy()
+238
View File
@@ -0,0 +1,238 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import sys
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from importlib.metadata import PackageNotFoundError
# zvec ships a native C++ extension that is only built and tested for 64-bit
# CPython. A 32-bit interpreter would fail to load the extension with an
# obscure error, so fail fast here with an actionable message.
if sys.maxsize <= 2**32:
raise ImportError(
"zvec requires a 64-bit Python interpreter; "
"the current interpreter is 32-bit and is not supported."
)
# Register the wheel-bundled jieba dict dir so `import zvec` alone makes
# the jieba FTS tokenizer usable. Users can still override via
# zvec.init(jieba_dict_dir=...), zvec.set_default_jieba_dict_dir(...),
# ZVEC_JIEBA_DICT_DIR, or per-field FtsIndexParam.extra_params.
try:
from importlib.resources import files as _resource_files
from zvec._zvec import (
get_default_jieba_dict_dir,
set_default_jieba_dict_dir,
)
set_default_jieba_dict_dir(str(_resource_files("zvec").joinpath("data/jieba_dict")))
except Exception:
# Custom builds without bundled dict; users must configure explicitly.
pass
# ==============================
# Public API — grouped by category
# ==============================
# —— DiskAnn runtime plugin ——
# Re-export the plugin management entry points defined by the C++ extension.
# DiskAnn normally auto-loads on first use; these APIs let tests and
# diagnostic tools preload the plugin and get a clear error if libaio is
# missing or the plugin shared object cannot be located.
from zvec._zvec import (
DISKANN_PLUGIN_DLOPEN_FAILED,
DISKANN_PLUGIN_LIBAIO_MISSING,
DISKANN_PLUGIN_OK,
DISKANN_PLUGIN_UNSUPPORTED_PLATFORM,
is_diskann_plugin_loaded,
is_libaio_available,
load_diskann_plugin,
)
from . import model as model
# —— Extensions ——
from .extension import (
BM25EmbeddingFunction,
DefaultLocalDenseEmbedding,
DefaultLocalReRanker,
DefaultLocalSparseEmbedding,
DenseEmbeddingFunction,
OpenAIDenseEmbedding,
OpenAIFunctionBase,
QwenDenseEmbedding,
QwenFunctionBase,
QwenReRanker,
QwenSparseEmbedding,
ReRanker,
RrfReRanker,
SentenceTransformerFunctionBase,
SparseEmbeddingFunction,
WeightedReRanker,
)
# —— Typing ——
from .model import param as param
from .model import schema as schema
# —— Core data structures ——
from .model.collection import Collection
from .model.doc import Doc, DocList
# —— Query & index parameters ——
# —— FTS params (C++ binding) ——
from .model.param import (
AddColumnOption,
AlterColumnOption,
CollectionOption,
DiskAnnIndexParam,
DiskAnnQueryParam,
FlatIndexParam,
FtsIndexParam,
FtsQueryParam,
HnswIndexParam,
HnswQueryParam,
HnswRabitqIndexParam,
HnswRabitqQueryParam,
IndexOption,
InvertIndexParam,
IVFIndexParam,
IVFQueryParam,
OptimizeOption,
QuantizerParam,
VamanaIndexParam,
VamanaQueryParam,
)
from .model.param.query import Fts, Query, VectorQuery
# —— Schema & field definitions ——
from .model.schema import CollectionSchema, CollectionStats, FieldSchema, VectorSchema
# —— tools ——
from .tool import require_module
from .typing import (
DataType,
IndexType,
MetricType,
QuantizeType,
Status,
StatusCode,
)
from .typing.enum import LogLevel, LogType
# —— lifecycle ——
from .zvec import create_and_open, init, open
# ==============================
# Public interface declaration
# ==============================
__all__ = [
# Zvec functions
"create_and_open",
"init",
"open",
"set_default_jieba_dict_dir",
"get_default_jieba_dict_dir",
# Core classes
"Collection",
"Doc",
"DocList",
# Schema
"CollectionSchema",
"FieldSchema",
"VectorSchema",
"CollectionStats",
# Parameters
"Query",
"VectorQuery",
"Fts",
"FtsIndexParam",
"FtsQueryParam",
"InvertIndexParam",
"HnswIndexParam",
"HnswRabitqIndexParam",
"FlatIndexParam",
"IVFIndexParam",
"DiskAnnIndexParam",
"DiskAnnQueryParam",
"CollectionOption",
"IndexOption",
"OptimizeOption",
"AddColumnOption",
"AlterColumnOption",
"HnswQueryParam",
"HnswRabitqQueryParam",
"IVFQueryParam",
"QuantizerParam",
"VamanaIndexParam",
"VamanaQueryParam",
# Extensions
"DenseEmbeddingFunction",
"SparseEmbeddingFunction",
"QwenFunctionBase",
"OpenAIFunctionBase",
"SentenceTransformerFunctionBase",
"ReRanker",
"DefaultLocalDenseEmbedding",
"DefaultLocalSparseEmbedding",
"BM25EmbeddingFunction",
"OpenAIDenseEmbedding",
"QwenDenseEmbedding",
"QwenSparseEmbedding",
"RrfReRanker",
"WeightedReRanker",
"DefaultLocalReRanker",
"QwenReRanker",
# Typing
"DataType",
"MetricType",
"QuantizeType",
"IndexType",
"LogLevel",
"LogType",
"Status",
"StatusCode",
# Tools
"require_module",
# DiskAnn plugin
"load_diskann_plugin",
"is_diskann_plugin_loaded",
"is_libaio_available",
"DISKANN_PLUGIN_OK",
"DISKANN_PLUGIN_UNSUPPORTED_PLATFORM",
"DISKANN_PLUGIN_LIBAIO_MISSING",
"DISKANN_PLUGIN_DLOPEN_FAILED",
]
# ==============================
# Version handling
# ==============================
__version__: str
try:
from importlib.metadata import version
except ImportError:
from importlib_metadata import version # Python < 3.8
try:
__version__ = version("zvec")
except Exception:
__version__ = "unknown"
+202
View File
@@ -0,0 +1,202 @@
"""
Zvec core module
"""
from __future__ import annotations
import collections
from . import typing
from .extension import ReRanker, RrfReRanker, WeightedReRanker
from .extension.embedding import DenseEmbeddingFunction
from .model import param, schema
from .model.collection import Collection
from .model.doc import Doc, DocList
from .model.param import (
AddColumnOption,
AlterColumnOption,
CollectionOption,
DiskAnnIndexParam,
DiskAnnQueryParam,
FlatIndexParam,
FtsIndexParam,
FtsQueryParam,
HnswIndexParam,
HnswQueryParam,
HnswRabitqIndexParam,
HnswRabitqQueryParam,
IndexOption,
InvertIndexParam,
IVFIndexParam,
IVFQueryParam,
OptimizeOption,
QuantizerParam,
VamanaIndexParam,
VamanaQueryParam,
)
from .model.param.query import Fts, Query, VectorQuery
from .model.schema import CollectionSchema, CollectionStats, FieldSchema, VectorSchema
from .tool import require_module
from .typing import (
DataType,
IndexType,
MetricType,
QuantizeType,
Status,
StatusCode,
)
from .typing.enum import LogLevel, LogType
from .zvec import create_and_open, init, open
__all__: list = [
"AddColumnOption",
"AlterColumnOption",
"Collection",
"CollectionOption",
"CollectionSchema",
"CollectionStats",
"DataType",
"DenseEmbeddingFunction",
"DiskAnnIndexParam",
"DiskAnnQueryParam",
"Doc",
"DocList",
"FieldSchema",
"FlatIndexParam",
"Fts",
"FtsIndexParam",
"FtsQueryParam",
"HnswIndexParam",
"HnswQueryParam",
"HnswRabitqIndexParam",
"HnswRabitqQueryParam",
"IVFIndexParam",
"IVFQueryParam",
"IndexOption",
"IndexType",
"InvertIndexParam",
"LogLevel",
"LogType",
"MetricType",
"OptimizeOption",
"QuantizeType",
"QuantizerParam",
"Query",
"ReRanker",
"RrfReRanker",
"Status",
"StatusCode",
"VamanaIndexParam",
"VamanaQueryParam",
"VectorQuery",
"VectorSchema",
"WeightedReRanker",
"create_and_open",
"init",
"open",
"require_module",
]
class _Collection:
@staticmethod
def CreateAndOpen(
arg0: str, arg1: schema._CollectionSchema, arg2: param.CollectionOption
) -> _Collection: ...
@staticmethod
def Open(arg0: str, arg1: param.CollectionOption) -> _Collection: ...
def AddColumn(
self,
arg0: schema._FieldSchema,
arg1: str,
arg2: param.AddColumnOption,
) -> None: ...
def AlterColumn(
self,
arg0: str,
arg1: str,
arg2: schema._FieldSchema,
arg3: param.AlterColumnOption,
) -> None: ...
def CreateIndex(
self, arg0: str, arg1: param.IndexParam, arg2: param.IndexOption
) -> None: ...
def Delete(self, arg0: collections.abc.Sequence[str]) -> list[typing.Status]: ...
def DeleteByFilter(self, arg0: str) -> None: ...
def Destroy(self) -> None: ...
def DropColumn(self, arg0: str) -> None: ...
def DropIndex(self, arg0: str) -> None: ...
def Fetch(
self,
pks: collections.abc.Sequence[str],
output_fields: list[str] | None = None,
include_vector: bool = True,
) -> dict[str, _Doc]: ...
def Flush(self) -> None: ...
def GroupByQuery(self, arg0: ...) -> list[...]: ...
def Insert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...
def Optimize(self, arg0: param.OptimizeOption) -> None: ...
def Options(self) -> param.CollectionOption: ...
def Path(self) -> str: ...
def Query(self, arg0: param._SearchQuery) -> list[_Doc]: ...
def Schema(self) -> schema._CollectionSchema: ...
def Stats(self) -> schema.CollectionStats: ...
def Update(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...
def Upsert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...
def _debug_hnsw_storage_mode(self, column_name: str) -> str:
"""Debug-only: returns the storage mode of the HNSW entity on the
given vector column. One of 'mmap', 'buffer_pool', 'contiguous'.
Raises KeyError if no HNSW index exists on the column, or
ValueError if the column's index is not an HNSW index. Intended
for introspection and testing only; not part of the stable API."""
def __getstate__(self) -> tuple: ...
def __setstate__(self, arg0: tuple) -> None: ...
class _Doc:
def __getstate__(self) -> bytes: ...
def __init__(self) -> None: ...
def __setstate__(self, arg0: bytes) -> None: ...
def field_names(self) -> list[str]: ...
def get_any(self, arg0: str, arg1: typing.DataType) -> typing.Any: ...
def has_field(self, arg0: str) -> bool: ...
def pk(self) -> str: ...
def score(self) -> float: ...
def set_any(self, arg0: str, arg1: typing.DataType, arg2: typing.Any) -> bool: ...
def set_pk(self, arg0: str) -> None: ...
def set_score(self, arg0: typing.SupportsFloat) -> None: ...
class _DocOp:
"""
Members:
INSERT
UPDATE
DELETE
UPSERT
"""
DELETE: typing.ClassVar[_DocOp] # value = <_DocOp.DELETE: 3>
INSERT: typing.ClassVar[_DocOp] # value = <_DocOp.INSERT: 0>
UPDATE: typing.ClassVar[_DocOp] # value = <_DocOp.UPDATE: 2>
UPSERT: typing.ClassVar[_DocOp] # value = <_DocOp.UPSERT: 1>
__members__: typing.ClassVar[
dict[str, _DocOp]
] # value = {'INSERT': <_DocOp.INSERT: 0>, 'UPDATE': <_DocOp.UPDATE: 2>, 'DELETE': <_DocOp.DELETE: 3>, 'UPSERT': <_DocOp.UPSERT: 1>}
def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __init__(self, value: typing.SupportsInt) -> None: ...
def __int__(self) -> int: ...
def __ne__(self, other: typing.Any) -> bool: ...
def __repr__(self) -> str: ...
def __setstate__(self, state: typing.SupportsInt) -> None: ...
def __str__(self) -> str: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
+18
View File
@@ -0,0 +1,18 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from .constants import DenseVectorType, SparseVectorType, VectorType
__all__ = ["DenseVectorType", "SparseVectorType", "VectorType"]
+33
View File
@@ -0,0 +1,33 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Optional, TypeVar, Union
import numpy as np
# VectorType: DenseVectorType | SparseVectorType
DenseVectorType = Union[list[float], list[int], np.ndarray]
SparseVectorType = dict[int, float]
VectorType = Optional[Union[DenseVectorType, SparseVectorType]]
# Embeddable: Text | Image | Audio
TEXT = str
IMAGE = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array
AUDIO = Union[str, bytes, np.ndarray] # file path, raw bytes, or numpy array
Embeddable = Optional[Union[TEXT, IMAGE, AUDIO]]
# Multimodal Embeddable
MD = TypeVar("MD", bound=Embeddable, contravariant=True)
+24
View File
@@ -0,0 +1,24 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from .query_executor import (
QueryContext,
QueryExecutor,
)
__all__ = [
"QueryContext",
"QueryExecutor",
]
+266
View File
@@ -0,0 +1,266 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Optional, Union
import numpy as np
from zvec._zvec import _Collection, _MultiQuery
from zvec._zvec.param import _Fts, _SearchQuery, _SubQuery
from ..extension import CallbackReRanker, ReRanker, RrfReRanker, WeightedReRanker
from ..model.convert import convert_to_py_doc
from ..model.doc import DocList
from ..model.param.query import Query
from ..model.schema import CollectionSchema
from ..typing import DataType
__all__ = [
"QueryContext",
"QueryExecutor",
]
DTYPE_MAP = {
DataType.VECTOR_FP16.value: np.float16,
DataType.VECTOR_FP32.value: np.float32,
DataType.VECTOR_FP64.value: np.float64,
DataType.VECTOR_INT8.value: np.int8,
}
def convert_to_numpy(vec: Union[list, np.ndarray], dtype: np.dtype) -> np.ndarray:
if isinstance(vec, np.ndarray):
if vec.dtype == dtype and vec.ndim == 1:
return vec
return np.asarray(vec, dtype=dtype).flatten()
try:
arr = np.asarray(vec, dtype=dtype)
if arr.ndim != 1:
arr = arr.flatten()
return arr
except (ValueError, TypeError) as e:
raise TypeError(
f"Cannot convert input to 1D numpy array with dtype={dtype}: {type(vec)}"
) from e
class QueryContext:
def __init__(
self,
topk: int,
filter: Optional[str] = None,
include_vector: bool = False,
queries: Optional[list[Query]] = None,
output_fields: Optional[list[str]] = None,
reranker: Optional[ReRanker] = None,
):
# query param
self._filter = filter
self._queries = queries or []
self._topk = topk
self._include_vector = include_vector
self._output_fields = output_fields
# reranker
self._reranker = reranker
@property
def topk(self):
return self._topk
@property
def queries(self):
return self._queries
@property
def filter(self):
return self._filter
@property
def reranker(self):
return self._reranker
@property
def output_fields(self):
return self._output_fields
@property
def include_vector(self):
return self._include_vector
class QueryExecutor:
"""Unified query executor that routes based on query count and reranker type."""
def __init__(self, schema: CollectionSchema):
self._schema = schema
def _build_queries(
self, ctx: QueryContext, collection: _Collection
) -> list[_SearchQuery]:
"""Build query vector list (no validation, conversion only)."""
if not ctx.queries:
return [self._build_base_search_query(ctx)]
return [
self._build_search_query(ctx, query, collection) for query in ctx.queries
]
def execute(self, ctx: QueryContext, collection: _Collection) -> DocList:
"""Execute a query, routing by query count.
A single (or vector-less) query is sent to C++ as a ``_SearchQuery``;
multiple queries are assembled into a ``_MultiQuery``.
"""
queries = self._build_queries(ctx, collection)
if not queries:
raise ValueError("No query to execute")
if len(queries) == 1:
return self._execute_single_query(queries[0], collection)
return self._execute_multi_query(ctx, queries, collection)
def _execute_single_query(
self, query: _SearchQuery, collection: _Collection
) -> DocList:
"""Single/vector-less query: send a ``_SearchQuery`` to C++."""
docs = collection.Query(query)
return [convert_to_py_doc(doc, self._schema) for doc in docs]
def _execute_multi_query(
self, ctx: QueryContext, queries: list[_SearchQuery], collection: _Collection
) -> DocList:
"""Multiple queries: send a ``_MultiQuery`` to C++.
A Python-only reranker (e.g. a model/API-based one) cannot run inside
the C++ MultiQuery, so each route is executed individually and merged by
the reranker in Python. The built-in RRF/Weighted/Callback rerankers use
the C++ variant-based fast path.
"""
reranker = ctx.reranker
if reranker is None:
raise ValueError(
"A reranker is required to merge results from multiple queries; "
"specify the 'reranker' argument."
)
if not isinstance(reranker, (RrfReRanker, WeightedReRanker, CallbackReRanker)):
docs_list = self._execute_python_pipeline(queries, collection)
return self._merge_and_rerank(ctx, docs_list)
multi_query = self._build_multi_query(ctx, queries)
docs = collection.Query(multi_query)
return [convert_to_py_doc(doc, self._schema) for doc in docs]
def _build_multi_query(
self, ctx: QueryContext, queries: list[_SearchQuery]
) -> _MultiQuery:
"""Assemble a C++ ``_MultiQuery`` from per-route ``_SearchQuery`` objects."""
multi_query = _MultiQuery()
multi_query.queries = [_SubQuery.from_search_query(query) for query in queries]
# num_candidates controls per-sub-query candidate count for reranking pool.
# It must NOT be limited to the final output topk; use at least the C++
# SubQuery default of 10 to ensure sufficient candidates for reranking.
_DEFAULT_NUM_CANDIDATES = 10
for sub in multi_query.queries:
sub.num_candidates = max(ctx.topk, _DEFAULT_NUM_CANDIDATES)
multi_query.topk = ctx.topk
if ctx.filter:
multi_query.filter = ctx.filter
multi_query.include_vector = ctx.include_vector
if ctx.output_fields is not None:
multi_query.output_fields = ctx.output_fields
# Set rerank strategy via the C++ variant-based API.
reranker = ctx.reranker
if isinstance(reranker, RrfReRanker):
multi_query.set_rerank_rrf(reranker.rank_constant)
elif isinstance(reranker, WeightedReRanker):
multi_query.set_rerank_weighted(reranker.weights)
elif isinstance(reranker, CallbackReRanker):
multi_query.set_rerank_callback(reranker._callback)
return multi_query
def _execute_python_pipeline(
self, vectors: list[_SearchQuery], collection: _Collection
) -> list[DocList]:
"""Execute queries serially for the Python-only reranker path."""
return [self._execute_single_query(query, collection) for query in vectors]
def _merge_and_rerank(self, ctx: QueryContext, docs_list: list[DocList]) -> DocList:
"""Merge and rerank results from the Python pipeline path."""
if not docs_list:
raise ValueError("Query results is empty")
if len(docs_list) == 1 and not ctx.reranker:
return docs_list[0]
return ctx.reranker.rerank(docs_list, ctx.topk)
def _build_base_search_query(self, ctx: QueryContext) -> _SearchQuery:
search_query = _SearchQuery()
search_query.topk = ctx.topk
search_query.include_vector = ctx.include_vector
if ctx.filter:
search_query.filter = ctx.filter
if ctx.output_fields is not None:
search_query.output_fields = ctx.output_fields
return search_query
def _apply_fts(self, query: Query, search_query: _SearchQuery) -> None:
"""Set FTS query on search_query if the query has FTS parameters."""
if query.has_fts():
fts = _Fts()
fts.query_string = query.fts.query_string or ""
fts.match_string = query.fts.match_string or ""
search_query.fts = fts
def _build_search_query(
self, ctx: QueryContext, query: Query, collection: _Collection
) -> _SearchQuery:
query._validate()
search_query = self._build_base_search_query(ctx)
search_query.field_name = query.field_name
if query.param:
search_query.query_params = query.param
# set FTS query if provided
self._apply_fts(query, search_query)
vector_schema = None
if query.has_vector() or query.has_id():
vector_schema = (
self._schema.vector(query.field_name)
if query
else self._schema.vectors[0]
)
if vector_schema is None:
raise ValueError("No vector field found")
# set vector
if query.has_vector():
vec_data = query.vector
elif query.has_id():
fetched = collection.Fetch([query.id])
doc = next(iter(fetched.values()), None)
if not doc:
raise ValueError(f"Document with id '{query.id}' not found")
vec_data = doc.get_any(vector_schema.name, vector_schema.data_type)
else:
return search_query
target_dtype = DTYPE_MAP.get(vector_schema.data_type.value)
search_query.set_vector(
vector_schema._get_object(),
convert_to_numpy(vec_data, target_dtype) if target_dtype else vec_data,
)
return search_query
+58
View File
@@ -0,0 +1,58 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from .bm25_embedding_function import BM25EmbeddingFunction
from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction
from .http_embedding_function import HTTPDenseEmbedding
from .jina_embedding_function import JinaDenseEmbedding
from .jina_function import JinaFunctionBase
from .multi_vector_reranker import CallbackReRanker, RrfReRanker, WeightedReRanker
from .openai_embedding_function import OpenAIDenseEmbedding
from .openai_function import OpenAIFunctionBase
from .qwen_embedding_function import QwenDenseEmbedding, QwenSparseEmbedding
from .qwen_function import QwenFunctionBase
from .qwen_rerank_function import QwenReRanker
from .rerank_function import RerankFunction
from .rerank_function import RerankFunction as ReRanker
from .sentence_transformer_embedding_function import (
DefaultLocalDenseEmbedding,
DefaultLocalSparseEmbedding,
)
from .sentence_transformer_function import SentenceTransformerFunctionBase
from .sentence_transformer_rerank_function import DefaultLocalReRanker
__all__ = [
"BM25EmbeddingFunction",
"CallbackReRanker",
"DefaultLocalDenseEmbedding",
"DefaultLocalReRanker",
"DefaultLocalSparseEmbedding",
"DenseEmbeddingFunction",
"HTTPDenseEmbedding",
"JinaDenseEmbedding",
"JinaFunctionBase",
"OpenAIDenseEmbedding",
"OpenAIFunctionBase",
"QwenDenseEmbedding",
"QwenFunctionBase",
"QwenReRanker",
"QwenSparseEmbedding",
"ReRanker",
"RerankFunction",
"RrfReRanker",
"SentenceTransformerFunctionBase",
"SparseEmbeddingFunction",
"WeightedReRanker",
]
@@ -0,0 +1,375 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from functools import lru_cache
from typing import Literal, Optional
from ..common.constants import TEXT, SparseVectorType
from ..tool import require_module
from .embedding_function import SparseEmbeddingFunction
class BM25EmbeddingFunction(SparseEmbeddingFunction[TEXT]):
"""BM25-based sparse embedding function using DashText SDK.
This class provides text-to-sparse-vector embedding capabilities using
the DashText library with BM25 algorithm. BM25 (Best Matching 25) is a
probabilistic retrieval function used for lexical search and document
ranking based on term frequency and inverse document frequency.
BM25 generates sparse vectors where each dimension corresponds to a term in
the vocabulary, and the value represents the BM25 score for that term. It's
particularly effective for:
- Lexical search and keyword matching
- Document ranking and information retrieval
- Combining with dense embeddings for hybrid search
- Traditional IR tasks where exact term matching is important
This implementation uses DashText's SparseVectorEncoder, which provides
efficient BM25 computation for Chinese and English text using either a
built-in encoder or custom corpus training.
Args:
corpus (Optional[list[str]], optional): List of documents to train the
BM25 encoder. If provided, creates a custom encoder trained on this
corpus for better domain-specific accuracy. If ``None``, uses the
built-in encoder. Defaults to ``None``.
encoding_type (Literal["query", "document"], optional): Encoding mode
for text processing. Use ``"query"`` for search queries (default) and
``"document"`` for document indexing. This distinction optimizes the
BM25 scoring for asymmetric retrieval tasks. Defaults to ``"query"``.
language (Literal["zh", "en"], optional): Language for built-in encoder.
Only used when corpus is None. ``"zh"`` for Chinese (trained on Chinese
Wikipedia), ``"en"`` for English. Defaults to ``"zh"``.
b (float, optional): Document length normalization parameter for BM25.
Range [0, 1]. 0 means no normalization, 1 means full normalization.
Only used with custom corpus. Defaults to ``0.75``.
k1 (float, optional): Term frequency saturation parameter for BM25.
Higher values give more weight to term frequency. Only used with
custom corpus. Defaults to ``1.2``.
**kwargs: Additional parameters for DashText encoder customization.
Attributes:
corpus_size (int): Number of documents in the training corpus (0 if using built-in encoder).
encoding_type (str): The encoding type being used ("query" or "document").
language (str): The language of the built-in encoder ("zh" or "en").
Raises:
ValueError: If corpus is provided but empty or contains non-string elements.
TypeError: If input to ``embed()`` is not a string.
RuntimeError: If DashText encoder initialization or training fails.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires the ``dashtext`` package: ``pip install dashtext``
- Two encoder options available:
1. **Built-in encoder** (no corpus needed): Pre-trained models for
Chinese (zh) and English (en), good generalization, works out-of-the-box
2. **Custom encoder** (corpus required): Better accuracy for domain-specific
terminology, requires training on your full corpus with BM25 parameters
- Encoding types:
* ``encoding_type="query"``: Optimized for search queries (shorter text)
* ``encoding_type="document"``: Optimized for document indexing (longer text)
- BM25 parameters (b, k1) only apply to custom encoder training
- Output is sorted by indices (vocabulary term IDs) for consistency
- Results are cached (LRU cache, maxsize=10) to reduce computation
- No API key or network connectivity required (local computation)
Examples:
>>> # Option 1: Using built-in encoder for Chinese (no corpus needed)
>>> from zvec.extension import BM25EmbeddingFunction
>>>
>>> # For query encoding (Chinese)
>>> bm25_query_zh = BM25EmbeddingFunction(language="zh", encoding_type="query")
>>> query_vec = bm25_query_zh.embed("什么是机器学习")
>>> isinstance(query_vec, dict)
True
>>> # query_vec: {1169440797: 0.29, 2045788977: 0.70, ...}
>>> # For document encoding (Chinese)
>>> bm25_doc_zh = BM25EmbeddingFunction(language="zh", encoding_type="document")
>>> doc_vec = bm25_doc_zh.embed("机器学习是人工智能的一个重要分支...")
>>> isinstance(doc_vec, dict)
True
>>> # Using built-in encoder for English
>>> bm25_query_en = BM25EmbeddingFunction(language="en", encoding_type="query")
>>> query_vec_en = bm25_query_en.embed("what is vector search service")
>>> isinstance(query_vec_en, dict)
True
>>> # Option 2: Using custom corpus for domain-specific accuracy
>>> corpus = [
... "机器学习是人工智能的一个重要分支",
... "深度学习使用多层神经网络进行特征提取",
... "自然语言处理技术用于理解和生成人类语言"
... ]
>>> bm25_custom = BM25EmbeddingFunction(
... corpus=corpus,
... encoding_type="query",
... b=0.75,
... k1=1.2
... )
>>> custom_vec = bm25_custom.embed("机器学习算法")
>>> isinstance(custom_vec, dict)
True
>>> # Hybrid search: combining with dense embeddings
>>> from zvec.extension import DefaultLocalDenseEmbedding
>>> dense_emb = DefaultLocalDenseEmbedding()
>>> bm25_emb = BM25EmbeddingFunction(language="zh", encoding_type="query")
>>>
>>> query = "machine learning algorithms"
>>> dense_vec = dense_emb.embed(query) # Semantic similarity
>>> sparse_vec = bm25_emb.embed(query) # Lexical matching
>>> # Combine scores for hybrid retrieval
>>> # Callable interface
>>> sparse_vec = bm25_query_zh("information retrieval")
>>> isinstance(sparse_vec, dict)
True
>>> # Error handling
>>> try:
... bm25_query_zh.embed("") # Empty query
... except ValueError as e:
... print(f"Error: {e}")
Error: Input text cannot be empty or whitespace only
See Also:
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
- ``DefaultLocalSparseEmbedding``: SPLADE-based sparse embedding
- ``QwenSparseEmbedding``: API-based sparse embedding using Qwen
- ``DefaultLocalDenseEmbedding``: Dense embedding for semantic search
References:
- DashText Documentation: https://help.aliyun.com/zh/document_detail/2546039.html
- DashText PyPI: https://pypi.org/project/dashtext/
- BM25 Algorithm: Robertson & Zaragoza (2009)
"""
def __init__(
self,
corpus: Optional[list[str]] = None,
encoding_type: Literal["query", "document"] = "query",
language: Literal["zh", "en"] = "zh",
b: float = 0.75,
k1: float = 1.2,
**kwargs,
):
"""Initialize the BM25 embedding function.
Args:
corpus (Optional[list[str]]): Optional corpus for training custom encoder.
If None, uses built-in encoder. Defaults to None.
encoding_type (Literal["query", "document"]): Text encoding mode.
Use "query" for search queries, "document" for indexing.
Defaults to "query".
language (Literal["zh", "en"]): Language for built-in encoder.
"zh" for Chinese, "en" for English. Defaults to "zh".
b (float): Document length normalization for BM25 [0, 1].
Only used with custom corpus. Defaults to 0.75.
k1 (float): Term frequency saturation for BM25.
Only used with custom corpus. Defaults to 1.2.
**kwargs: Additional DashText encoder parameters.
Raises:
ValueError: If corpus is provided but empty or invalid.
ImportError: If dashtext package is not installed.
RuntimeError: If encoder initialization or training fails.
"""
# Validate corpus if provided
if corpus is not None:
if not corpus or not isinstance(corpus, list):
raise ValueError("Corpus must be a non-empty list of strings")
if not all(isinstance(doc, str) for doc in corpus):
raise ValueError("All corpus documents must be strings")
# Import dashtext
self._dashtext = require_module("dashtext")
self._corpus = corpus
self._encoding_type = encoding_type
self._language = language
self._b = b
self._k1 = k1
self._extra_params = kwargs
# Initialize the BM25 encoder
self._build_encoder()
def _build_encoder(self):
"""Build the BM25 sparse vector encoder.
Creates either a built-in encoder (pre-trained) or a custom encoder
trained on the provided corpus.
Raises:
RuntimeError: If encoder initialization or training fails.
ImportError: If dashtext package is not installed.
"""
try:
if self._corpus is None:
# Use built-in encoder (pre-trained on Wikipedia)
# language: 'zh' for Chinese, 'en' for English
self._encoder = self._dashtext.SparseVectorEncoder.default(
name=self._language
)
else:
# Create custom encoder with BM25 parameters
self._encoder = self._dashtext.SparseVectorEncoder(
b=self._b, k1=self._k1, **self._extra_params
)
# Train encoder with the corpus
self._encoder.train(self._corpus)
except ImportError as e:
raise ImportError(
"dashtext package is required for BM25EmbeddingFunction. "
"Install it with: pip install dashtext"
) from e
except Exception as e:
if isinstance(e, (ValueError, RuntimeError)):
raise
raise RuntimeError(f"Failed to build BM25 encoder: {e!s}") from e
@property
def corpus_size(self) -> int:
"""int: Number of documents in the training corpus (0 if using built-in encoder)."""
return len(self._corpus) if self._corpus is not None else 0
@property
def encoding_type(self) -> str:
"""str: The encoding type being used ("query" or "document")."""
return self._encoding_type
@property
def language(self) -> str:
"""str: The language of the built-in encoder ("zh" or "en")."""
return self._language
@property
def extra_params(self) -> dict:
"""dict: Extra parameters for DashText encoder customization."""
return self._extra_params
def __call__(self, input: TEXT) -> SparseVectorType:
"""Make the embedding function callable.
Args:
input (TEXT): Input text to embed.
Returns:
SparseVectorType: Sparse vector as dictionary.
"""
return self.embed(input)
@lru_cache(maxsize=10)
def embed(self, input: TEXT) -> SparseVectorType:
"""Generate BM25 sparse embedding for the input text.
This method computes BM25 scores for the input text using DashText's
SparseVectorEncoder. The encoding behavior depends on the encoding_type:
- ``encoding_type="query"``: Uses ``encode_queries()`` for search queries
- ``encoding_type="document"``: Uses ``encode_documents()`` for documents
The result is a sparse vector where keys are term indices in the
vocabulary and values are BM25 scores.
Args:
input (TEXT): Input text string to embed. Must be non-empty after
stripping whitespace.
Returns:
SparseVectorType: A dictionary mapping vocabulary term index to BM25 score.
Only non-zero scores are included. The dictionary is sorted by indices
(keys) in ascending order for consistent output.
Example: ``{1169440797: 0.29, 2045788977: 0.70, ...}``
Raises:
TypeError: If ``input`` is not a string.
ValueError: If input is empty or whitespace-only.
RuntimeError: If BM25 encoding fails.
Examples:
>>> bm25 = BM25EmbeddingFunction(language="zh", encoding_type="query")
>>> sparse_vec = bm25.embed("query text")
>>> isinstance(sparse_vec, dict)
True
>>> all(isinstance(k, int) and isinstance(v, float) for k, v in sparse_vec.items())
True
>>> # Verify sorted output
>>> keys = list(sparse_vec.keys())
>>> keys == sorted(keys)
True
>>> # Error: empty input
>>> bm25.embed(" ")
ValueError: Input text cannot be empty or whitespace only
>>> # Error: non-string input
>>> bm25.embed(123)
TypeError: Expected 'input' to be str, got int
Note:
- BM25 scores are relative to the vocabulary statistics
- Output dictionary is always sorted by indices for consistency
- Terms not in the vocabulary will have zero scores (not included)
- This method is cached (maxsize=10) for performance
- DashText automatically handles Chinese/English text segmentation
"""
if not isinstance(input, str):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
try:
# Encode based on encoding_type
if self._encoding_type == "query":
sparse_vector = self._encoder.encode_queries(input)
else: # encoding_type == "document"
sparse_vector = self._encoder.encode_documents(input)
# DashText returns dict with int/long keys and float values
# Convert to standard format: {int: float}
sparse_dict: dict[int, float] = {}
for key, value in sparse_vector.items():
try:
idx = int(key)
val = float(value)
if val > 0:
sparse_dict[idx] = val
except (ValueError, TypeError):
# Skip invalid entries
continue
# Sort by indices (keys) to ensure consistent ordering
return dict(sorted(sparse_dict.items()))
except Exception as e:
if isinstance(e, (TypeError, ValueError)):
raise
raise RuntimeError(f"Failed to generate BM25 embedding: {e!s}") from e
+147
View File
@@ -0,0 +1,147 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from abc import abstractmethod
from typing import Protocol, runtime_checkable
from ..common.constants import MD, DenseVectorType, SparseVectorType
@runtime_checkable
class DenseEmbeddingFunction(Protocol[MD]):
"""Protocol for dense vector embedding functions.
Dense embedding functions map multimodal input (text, image, or audio) to
fixed-length real-valued vectors. This is a Protocol class that defines
the interface - implementations should provide their own initialization
and properties.
Type Parameters:
MD: The type of input data (bound to Embeddable: TEXT, IMAGE, or AUDIO).
Note:
- This is a Protocol class - it only defines the ``embed()`` interface.
- Implementations are free to define their own ``__init__``, properties,
and additional methods as needed.
- The ``embed()`` method is the only required interface.
Examples:
>>> # Custom text embedding implementation
>>> class MyTextEmbedding:
... def __init__(self, dimension: int, model_name: str):
... self.dimension = dimension
... self.model = load_model(model_name)
...
... def embed(self, input: str) -> list[float]:
... return self.model.encode(input).tolist()
>>> # Custom image embedding implementation
>>> class MyImageEmbedding:
... def __init__(self, dimension: int = 512):
... self.dimension = dimension
... self.model = load_image_model()
...
... def embed(self, input: Union[str, bytes, np.ndarray]) -> list[float]:
... if isinstance(input, str):
... image = load_image_from_path(input)
... else:
... image = input
... return self.model.extract_features(image).tolist()
>>> # Using built-in implementations
>>> from zvec.extension import QwenDenseEmbedding
>>> text_emb = QwenDenseEmbedding(dimension=768, api_key="sk-xxx")
>>> vector = text_emb.embed("Hello world")
"""
@abstractmethod
def embed(self, input: MD) -> DenseVectorType:
"""Generate a dense embedding vector for the input data.
Args:
input (MD): Multimodal input data to embed. Can be:
- TEXT (str): Text string
- IMAGE (str | bytes | np.ndarray): Image file path, raw bytes, or array
- AUDIO (str | bytes | np.ndarray): Audio file path, raw bytes, or array
Returns:
DenseVectorType: A dense vector representing the embedding.
Can be list[float], list[int], or np.ndarray.
Length should match the implementation's dimension.
"""
...
@runtime_checkable
class SparseEmbeddingFunction(Protocol[MD]):
"""Abstract base class for sparse vector embedding functions.
Sparse embedding functions map multimodal input (text, image, or audio) to
a dictionary of {index: weight}, where only non-zero dimensions are stored.
You can inherit this class to create custom sparse embedding functions.
Type Parameters:
MD: The type of input data (bound to Embeddable: TEXT, IMAGE, or AUDIO).
Note:
Subclasses must implement the ``embed()`` method.
Examples:
>>> # Using built-in text sparse embedding (e.g., BM25, TF-IDF)
>>> sparse_emb = SomeSparseEmbedding()
>>> vector = sparse_emb.embed("Hello world")
>>> # Returns: {0: 0.5, 42: 1.2, 100: 0.8}
>>> # Custom BM25 sparse embedding function
>>> class MyBM25Embedding(SparseEmbeddingFunction):
... def __init__(self, vocab_size: int = 10000):
... self.vocab_size = vocab_size
... self.tokenizer = MyTokenizer()
...
... def embed(self, input: str) -> dict[int, float]:
... tokens = self.tokenizer.tokenize(input)
... sparse_vector = {}
... for token_id, weight in self._calculate_bm25(tokens):
... if weight > 0:
... sparse_vector[token_id] = weight
... return sparse_vector
...
... def _calculate_bm25(self, tokens):
... # BM25 calculation logic
... pass
>>> # Custom sparse image feature extractor
>>> class MySparseImageEmbedding(SparseEmbeddingFunction):
... def embed(self, input: Union[str, bytes, np.ndarray]) -> dict[int, float]:
... image = self._load_image(input)
... features = self._extract_sparse_features(image)
... return {idx: val for idx, val in enumerate(features) if val != 0}
"""
@abstractmethod
def embed(self, input: MD) -> SparseVectorType:
"""Generate a sparse embedding for the input data.
Args:
input (MD): Multimodal input data to embed. Can be:
- TEXT (str): Text string
- IMAGE (str | bytes | np.ndarray): Image file path, raw bytes, or array
- AUDIO (str | bytes | np.ndarray): Audio file path, raw bytes, or array
Returns:
SparseVectorType: Mapping from dimension index to non-zero weight.
Only dimensions with non-zero values are included.
"""
...
@@ -0,0 +1,162 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import json
import os
import urllib.request
from functools import lru_cache
from typing import Optional
from ..common.constants import TEXT, DenseVectorType
from .embedding_function import DenseEmbeddingFunction
class HTTPDenseEmbedding(DenseEmbeddingFunction[TEXT]):
"""Dense text embedding function using any OpenAI-compatible HTTP endpoint.
This class calls any server that implements the ``/v1/embeddings`` API
(LM Studio, Ollama, vLLM, LocalAI, etc.) using only the Python standard
library — no extra dependencies are required.
The embedding dimension is detected automatically from the first server
response.
Args:
base_url (str, optional): Base URL of the embedding server.
Defaults to ``"http://localhost:1234"`` (LM Studio).
Common values:
- ``"http://localhost:1234"`` — LM Studio
- ``"http://localhost:11434"`` — Ollama
model (str, optional): Model identifier as expected by the server.
Defaults to ``"text-embedding-nomic-embed-text-v1.5@f16"``.
api_key (Optional[str], optional): Bearer token for authenticated
endpoints. Falls back to the ``OPENAI_API_KEY`` environment
variable. Leave as ``None`` for local servers that do not
require authentication.
timeout (int, optional): HTTP request timeout in seconds.
Defaults to 30.
Attributes:
dimension (int): Embedding vector dimensionality (auto-detected).
Raises:
TypeError: If ``embed()`` receives a non-string input.
ValueError: If input is empty/whitespace-only or the server returns
an unexpected response format.
RuntimeError: If the HTTP request fails or the server is unreachable.
Examples:
>>> from zvec.extension import HTTPDenseEmbedding
>>>
>>> # LM Studio (default)
>>> emb = HTTPDenseEmbedding()
>>> vector = emb.embed("Hello, world!")
>>> len(vector)
768
>>>
>>> # Ollama
>>> emb = HTTPDenseEmbedding(
... base_url="http://localhost:11434",
... model="nomic-embed-text",
... )
>>> vector = emb.embed("Semantic search with local models")
See Also:
- ``DenseEmbeddingFunction``: Protocol for dense embeddings.
- ``OpenAIDenseEmbedding``: Cloud embedding via the OpenAI API.
"""
ENDPOINT = "/v1/embeddings"
def __init__(
self,
base_url: str = "http://localhost:1234",
model: str = "text-embedding-nomic-embed-text-v1.5@f16",
api_key: Optional[str] = None,
timeout: int = 30,
) -> None:
self._base_url = base_url.rstrip("/")
self._model = model
self._api_key = api_key or os.environ.get("OPENAI_API_KEY", "")
self._timeout = timeout
self._dimension: Optional[int] = None
@property
def dimension(self) -> int:
"""int: Embedding vector dimensionality (auto-detected on first call)."""
if self._dimension is None:
self._dimension = len(self.embed("dimension probe"))
return self._dimension
def __call__(self, input: TEXT) -> DenseVectorType:
"""Make the embedding function callable."""
return self.embed(input)
@lru_cache(maxsize=256)
def embed(self, input: TEXT) -> DenseVectorType:
"""Generate a dense embedding vector for the input text.
Results are cached (LRU, up to 256 entries) so repeated strings
do not trigger extra HTTP requests.
Args:
input (TEXT): Input text string to embed. Must be non-empty
after stripping whitespace.
Returns:
DenseVectorType: A list of floats representing the embedding.
Raises:
TypeError: If *input* is not a string.
ValueError: If *input* is empty/whitespace-only or the server
returns an unexpected response format.
RuntimeError: If the HTTP request fails.
"""
if not isinstance(input, TEXT):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
url = self._base_url + self.ENDPOINT
payload = json.dumps({"model": self._model, "input": input}).encode()
headers: dict[str, str] = {"Content-Type": "application/json"}
if self._api_key:
headers["Authorization"] = f"Bearer {self._api_key}"
req = urllib.request.Request(url, data=payload, headers=headers, method="POST")
try:
with urllib.request.urlopen(req, timeout=self._timeout) as resp:
body = json.loads(resp.read())
except urllib.error.HTTPError as exc:
raise RuntimeError(
f"Embedding server returned HTTP {exc.code}: {exc.read().decode()}"
) from exc
except OSError as exc:
raise RuntimeError(
f"Could not reach embedding server at {url}: {exc}"
) from exc
try:
vector: list[float] = body["data"][0]["embedding"]
except (KeyError, IndexError) as exc:
raise ValueError(
f"Unexpected response format from embedding server: {body}"
) from exc
return vector
@@ -0,0 +1,240 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from functools import lru_cache
from typing import Optional
from ..common.constants import TEXT, DenseVectorType
from .embedding_function import DenseEmbeddingFunction
from .jina_function import JinaFunctionBase
class JinaDenseEmbedding(JinaFunctionBase, DenseEmbeddingFunction[TEXT]):
"""Dense text embedding function using Jina AI API.
This class provides text-to-vector embedding capabilities using Jina AI's
embedding models. It inherits from ``DenseEmbeddingFunction`` and implements
dense text embedding via the Jina Embeddings API (OpenAI-compatible).
Jina Embeddings v5 models support task-specific embedding through the
``task`` parameter, which optimizes the embedding for different use cases
such as retrieval, text matching, or classification. They also support
Matryoshka Representation Learning, allowing flexible output dimensions.
Args:
model (str, optional): Jina embedding model identifier.
Defaults to ``"jina-embeddings-v5-text-nano"``. Available models:
- ``"jina-embeddings-v5-text-nano"``: 768 dims, 239M params, 8K context
- ``"jina-embeddings-v5-text-small"``: 1024 dims, 677M params, 32K context
dimension (Optional[int], optional): Desired output embedding dimension.
If ``None``, uses model's default dimension. Supports Matryoshka
dimensions: 32, 64, 128, 256, 512, 768 (nano) / 1024 (small).
Defaults to ``None``.
api_key (Optional[str], optional): Jina API authentication key.
If ``None``, reads from ``JINA_API_KEY`` environment variable.
Obtain your key from: https://jina.ai/api-dashboard
task (Optional[str], optional): Task type to optimize embeddings for.
Defaults to ``None``. Valid values:
- ``"retrieval.query"``: For search queries
- ``"retrieval.passage"``: For documents/passages to be searched
- ``"text-matching"``: For symmetric text similarity
- ``"classification"``: For text classification
- ``"separation"``: For clustering/separation tasks
Attributes:
dimension (int): The embedding vector dimension.
data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation.
model (str): The Jina model name being used.
task (Optional[str]): The task type for embedding optimization.
Raises:
ValueError: If API key is not provided and not found in environment,
if task is not a valid task type, or if API returns an error response.
TypeError: If input to ``embed()`` is not a string.
RuntimeError: If network error or Jina service error occurs.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires the ``openai`` package: ``pip install openai``
- Jina API is OpenAI-compatible, so it uses the ``openai`` Python client
- Embedding results are cached (LRU cache, maxsize=10) to reduce API calls
- For retrieval tasks, use ``"retrieval.query"`` for queries and
``"retrieval.passage"`` for documents
- API usage requires a Jina API key from https://jina.ai/api-dashboard
Examples:
>>> # Basic usage with default model
>>> from zvec.extension import JinaDenseEmbedding
>>> import os
>>> os.environ["JINA_API_KEY"] = "jina_..."
>>>
>>> emb_func = JinaDenseEmbedding()
>>> vector = emb_func.embed("Hello, world!")
>>> len(vector)
768
>>> # Retrieval use case: embed queries and documents differently
>>> query_emb = JinaDenseEmbedding(task="retrieval.query")
>>> doc_emb = JinaDenseEmbedding(task="retrieval.passage")
>>>
>>> query_vector = query_emb.embed("What is machine learning?")
>>> doc_vector = doc_emb.embed("Machine learning is a subset of AI...")
>>> # Using larger model with custom dimension (Matryoshka)
>>> emb_func = JinaDenseEmbedding(
... model="jina-embeddings-v5-text-small",
... dimension=256,
... api_key="jina_...",
... task="text-matching",
... )
>>> vector = emb_func.embed("Semantic similarity comparison")
>>> len(vector)
256
>>> # Using with zvec collection
>>> import zvec
>>> emb_func = JinaDenseEmbedding(task="retrieval.passage")
>>> schema = zvec.CollectionSchema(
... name="docs",
... vectors=zvec.VectorSchema(
... "embedding", zvec.DataType.VECTOR_FP32, emb_func.dimension
... ),
... )
>>> collection = zvec.create_and_open(path="./my_docs", schema=schema)
See Also:
- ``DenseEmbeddingFunction``: Base class for dense embeddings
- ``OpenAIDenseEmbedding``: Alternative using OpenAI API
- ``QwenDenseEmbedding``: Alternative using Qwen/DashScope API
- ``DefaultLocalDenseEmbedding``: Local model without API calls
"""
def __init__(
self,
model: str = "jina-embeddings-v5-text-nano",
dimension: Optional[int] = None,
api_key: Optional[str] = None,
task: Optional[str] = None,
**kwargs,
):
"""Initialize the Jina dense embedding function.
Args:
model (str): Jina model name. Defaults to "jina-embeddings-v5-text-nano".
dimension (Optional[int]): Target embedding dimension or None for default.
api_key (Optional[str]): API key or None to use environment variable.
task (Optional[str]): Task type for embedding optimization or None.
**kwargs: Additional parameters for API calls.
Raises:
ValueError: If API key is not provided and not in environment,
or if task is not a valid task type.
"""
# Initialize base class for API connection
JinaFunctionBase.__init__(self, model=model, api_key=api_key, task=task)
# Store dimension configuration
self._custom_dimension = dimension
# Determine actual dimension
if dimension is None:
self._dimension = self._MODEL_DIMENSIONS.get(model, 768)
else:
self._dimension = dimension
# Store extra attributes
self._extra_params = kwargs
@property
def dimension(self) -> int:
"""int: The expected dimensionality of the embedding vector."""
return self._dimension
@property
def extra_params(self) -> dict:
"""dict: Extra parameters for model-specific customization."""
return self._extra_params
def __call__(self, input: TEXT) -> DenseVectorType:
"""Make the embedding function callable."""
return self.embed(input)
@lru_cache(maxsize=10)
def embed(self, input: TEXT) -> DenseVectorType:
"""Generate dense embedding vector for the input text.
This method calls the Jina Embeddings API to convert input text
into a dense vector representation. Results are cached to improve
performance for repeated inputs.
Args:
input (TEXT): Input text string to embed. Must be non-empty after
stripping whitespace. Maximum length depends on model:
8192 tokens for v5-nano, 32768 tokens for v5-small.
Returns:
DenseVectorType: A list of floats representing the embedding vector.
Length equals ``self.dimension``. Example:
``[0.123, -0.456, 0.789, ...]``
Raises:
TypeError: If ``input`` is not a string.
ValueError: If input is empty/whitespace-only, or if the API returns
an error or malformed response.
RuntimeError: If network connectivity issues or Jina service
errors occur.
Examples:
>>> emb = JinaDenseEmbedding(task="retrieval.query")
>>> vector = emb.embed("What is deep learning?")
>>> len(vector)
768
>>> isinstance(vector[0], float)
True
>>> # Error: empty input
>>> emb.embed(" ")
ValueError: Input text cannot be empty or whitespace only
>>> # Error: non-string input
>>> emb.embed(123)
TypeError: Expected 'input' to be str, got int
Note:
- This method is cached (maxsize=10). Identical inputs return cached results.
- The cache is based on exact string match (case-sensitive).
- Task type affects embedding optimization but not caching behavior.
"""
if not isinstance(input, TEXT):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
# Call API
embedding_vector = self._call_text_embedding_api(
input=input,
dimension=self._custom_dimension,
)
# Verify dimension
if len(embedding_vector) != self.dimension:
raise ValueError(
f"Dimension mismatch: expected {self.dimension}, "
f"got {len(embedding_vector)}"
)
return embedding_vector
+182
View File
@@ -0,0 +1,182 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import os
from typing import ClassVar, Optional
from ..common.constants import TEXT
from ..tool import require_module
class JinaFunctionBase:
"""Base class for Jina AI functions.
This base class provides common functionality for calling Jina AI APIs
and handling responses. It supports embeddings (dense) operations via
the OpenAI-compatible Jina Embeddings API.
This class is not meant to be used directly. Use concrete implementations:
- ``JinaDenseEmbedding`` for dense embeddings
Args:
model (str): Jina embedding model identifier.
api_key (Optional[str]): Jina API authentication key.
task (Optional[str]): Task type for the embedding model.
Note:
- This is an internal base class for code reuse across Jina features
- Subclasses should inherit from appropriate Protocol
- Provides unified API connection and response handling
- Jina API is OpenAI-compatible, using the ``openai`` Python client
"""
_BASE_URL: ClassVar[str] = "https://api.jina.ai/v1"
# Model default dimensions
_MODEL_DIMENSIONS: ClassVar[dict[str, int]] = {
"jina-embeddings-v5-text-nano": 768,
"jina-embeddings-v5-text-small": 1024,
}
# Model max tokens
_MODEL_MAX_TOKENS: ClassVar[dict[str, int]] = {
"jina-embeddings-v5-text-nano": 8192,
"jina-embeddings-v5-text-small": 32768,
}
# Valid task types
_VALID_TASKS: ClassVar[tuple[str, ...]] = (
"retrieval.query",
"retrieval.passage",
"text-matching",
"classification",
"separation",
)
def __init__(
self,
model: str,
api_key: Optional[str] = None,
task: Optional[str] = None,
):
"""Initialize the base Jina functionality.
Args:
model (str): Jina model name.
api_key (Optional[str]): API key or None to use environment variable.
task (Optional[str]): Task type for the embedding model.
Valid values: "retrieval.query", "retrieval.passage",
"text-matching", "classification", "separation".
Raises:
ValueError: If API key is not provided and not in environment,
or if task is not a valid task type.
"""
self._model = model
self._api_key = api_key or os.environ.get("JINA_API_KEY")
self._task = task
if not self._api_key:
raise ValueError(
"Jina API key is required. Please provide 'api_key' parameter "
"or set the 'JINA_API_KEY' environment variable. "
"Get your key from: https://jina.ai/api-dashboard"
)
if task is not None and task not in self._VALID_TASKS:
raise ValueError(
f"Invalid task '{task}'. Valid tasks: {', '.join(self._VALID_TASKS)}"
)
@property
def model(self) -> str:
"""str: The Jina model name currently in use."""
return self._model
@property
def task(self) -> Optional[str]:
"""Optional[str]: The task type for the embedding model."""
return self._task
def _get_client(self):
"""Get OpenAI-compatible client instance configured for Jina API.
Returns:
OpenAI: Configured OpenAI client pointing to Jina API.
Raises:
ImportError: If openai package is not installed.
"""
openai = require_module("openai")
return openai.OpenAI(api_key=self._api_key, base_url=self._BASE_URL)
def _call_text_embedding_api(
self,
input: TEXT,
dimension: Optional[int] = None,
) -> list:
"""Call Jina Embeddings API.
Args:
input (TEXT): Input text to embed.
dimension (Optional[int]): Target dimension for Matryoshka embeddings.
Returns:
list: Embedding vector as list of floats.
Raises:
RuntimeError: If API call fails.
ValueError: If API returns error response.
"""
try:
client = self._get_client()
# Prepare embedding parameters
params = {"model": self.model, "input": input}
# Add dimension parameter for Matryoshka support
if dimension is not None:
params["dimensions"] = dimension
# Add task parameter via extra_body
if self._task is not None:
params["extra_body"] = {"task": self._task}
# Call Jina API (OpenAI-compatible)
response = client.embeddings.create(**params)
except Exception as e:
# Check if it's an OpenAI API error
openai = require_module("openai")
if isinstance(e, (openai.APIError, openai.APIConnectionError)):
raise RuntimeError(f"Failed to call Jina API: {e!s}") from e
raise RuntimeError(f"Unexpected error during API call: {e!s}") from e
# Extract embedding from response
try:
if not response.data:
raise ValueError("Invalid API response: no embedding data returned")
embedding_vector = response.data[0].embedding
if not isinstance(embedding_vector, list):
raise ValueError(
"Invalid API response: embedding is not a list of numbers"
)
return embedding_vector
except (AttributeError, IndexError, TypeError) as e:
raise ValueError(f"Failed to parse API response: {e!s}") from e
@@ -0,0 +1,197 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from collections.abc import Callable
from typing import TYPE_CHECKING
from zvec._zvec import (
_CallbackParams,
_Doc,
_reranker_rerank,
_RrfParams,
_WeightedParams,
)
from ..model.doc import Doc, DocList
from .rerank_function import RerankFunction
if TYPE_CHECKING:
from ..model.schema import FieldSchema, VectorSchema
def _to_cpp_doc_lists(
query_results: list[list[Doc]],
) -> tuple[list[list], dict[str, Doc]]:
"""Convert Python Doc lists to C++ _Doc lists for reranker input."""
id_to_doc: dict[str, Doc] = {}
cpp_results: list[list] = []
for query_result in query_results:
cpp_list: list = []
for doc in query_result:
_doc = _Doc()
_doc.set_pk(doc.id)
_doc.set_score(doc.score if doc.score is not None else 0.0)
cpp_list.append(_doc)
if doc.id not in id_to_doc:
id_to_doc[doc.id] = doc
cpp_results.append(cpp_list)
return cpp_results, id_to_doc
def _from_cpp_docs(cpp_docs: list, id_to_doc: dict[str, Doc]) -> DocList:
"""Convert C++ rerank result _Doc list back to Python DocList."""
results: DocList = []
for _doc in cpp_docs:
doc_id = _doc.pk()
new_score = _doc.score()
original = id_to_doc.get(doc_id)
if original is not None:
results.append(original._replace(score=new_score))
else:
results.append(Doc(id=doc_id, score=new_score))
return results
class RrfReRanker(RerankFunction):
"""Re-ranker using Reciprocal Rank Fusion (RRF) for multi-vector search.
RRF combines results from multiple vector queries without requiring
relevance scores. The RRF score for a document at rank r is:
score = 1 / (k + r + 1)
where k is the rank constant.
Args:
rank_constant: RRF smoothing constant (default: 60).
Higher values reduce the influence of rank position.
Example:
>>> reranker = RrfReRanker(rank_constant=60)
>>> merged = reranker.rerank([results_a, results_b], topn=10)
"""
def __init__(self, rank_constant: int = 60):
self._rank_constant = rank_constant
@property
def rank_constant(self) -> int:
"""int: RRF rank constant."""
return self._rank_constant
def _to_cpp_params(self):
return _RrfParams(self._rank_constant)
def rerank(
self,
query_results: list[list[Doc]],
topn: int = 10,
*,
fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002
) -> DocList:
"""Apply RRF to combine multiple query results via C++ reranker."""
cpp_results, id_to_doc = _to_cpp_doc_lists(query_results)
cpp_docs = _reranker_rerank(self._to_cpp_params(), cpp_results, [], topn)
return _from_cpp_docs(cpp_docs, id_to_doc)
class WeightedReRanker(RerankFunction):
"""Re-ranker that combines scores using per-sub-query weights.
Each sub-query's score is normalized by metric type (automatic when used
via collection.multi_query), then multiplied by the corresponding weight.
Args:
weights: Per-sub-query weights. Length must match the number of
sub-queries.
Example:
>>> reranker = WeightedReRanker([0.7, 0.3])
>>> merged = reranker.rerank([results_a, results_b], topn=10,
... fields=field_schemas)
"""
def __init__(self, weights: list[float]):
self._weights = list(weights)
@property
def weights(self) -> list[float]:
"""list[float]: Per-sub-query weights."""
return self._weights
def _to_cpp_params(self):
return _WeightedParams(self._weights)
def rerank(
self,
query_results: list[list[Doc]],
topn: int = 10,
*,
fields: list[FieldSchema | VectorSchema] | None = None,
) -> DocList:
"""Combine scores from multiple sub-queries using weighted sum via C++ reranker.
Args:
query_results: Per-sub-query document lists.
topn: Maximum results to return.
fields: Per-sub-query Python FieldSchema/VectorSchema objects
(required for score normalization by metric type).
Raises:
ValueError: If fields is None (required for normalization).
"""
if not fields:
raise ValueError(
"WeightedReRanker.rerank() requires 'fields' for score normalization. "
"Pass field schemas via fields= parameter."
)
cpp_fields = [f._get_object() for f in fields]
cpp_results, id_to_doc = _to_cpp_doc_lists(query_results)
cpp_docs = _reranker_rerank(
self._to_cpp_params(), cpp_results, cpp_fields, topn
)
return _from_cpp_docs(cpp_docs, id_to_doc)
class CallbackReRanker(RerankFunction):
"""Re-ranker that delegates to a user-provided callback.
The callback receives sub-query results, field schemas, and topn.
Args:
callback: A callable with signature
(results: list[list[Doc]], fields: list, topn: int) -> list[Doc]
Example:
>>> def my_rerank(results, fields, topn):
... # custom logic
... return merged[:topn]
>>> reranker = CallbackReRanker(my_rerank)
>>> merged = reranker.rerank([results_a, results_b], topn=10)
"""
def __init__(self, callback: Callable):
self._callback = callback
def _to_cpp_params(self):
return _CallbackParams(self._callback)
def rerank(
self,
query_results: list[list[Doc]],
topn: int = 10,
*,
fields: list[FieldSchema | VectorSchema] | None = None,
) -> DocList:
"""Invoke the callback to re-rank documents."""
return self._callback(query_results, fields, topn)
@@ -0,0 +1,238 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from functools import lru_cache
from typing import Optional
from ..common.constants import TEXT, DenseVectorType
from .embedding_function import DenseEmbeddingFunction
from .openai_function import OpenAIFunctionBase
class OpenAIDenseEmbedding(OpenAIFunctionBase, DenseEmbeddingFunction[TEXT]):
"""Dense text embedding function using OpenAI API.
This class provides text-to-vector embedding capabilities using OpenAI's
embedding models. It inherits from ``DenseEmbeddingFunction`` and implements
dense text embedding via the OpenAI API.
The implementation supports various OpenAI embedding models with different
dimensions and includes automatic result caching for improved performance.
Args:
model (str, optional): OpenAI embedding model identifier.
Defaults to ``"text-embedding-3-small"``. Common options:
- ``"text-embedding-3-small"``: 1536 dims, cost-efficient, good performance
- ``"text-embedding-3-large"``: 3072 dims, highest quality
- ``"text-embedding-ada-002"``: 1536 dims, legacy model
dimension (Optional[int], optional): Desired output embedding dimension.
If ``None``, uses model's default dimension. For text-embedding-3 models,
you can specify custom dimensions (e.g., 256, 512, 1024, 1536).
Defaults to ``None``.
api_key (Optional[str], optional): OpenAI API authentication key.
If ``None``, reads from ``OPENAI_API_KEY`` environment variable.
Obtain your key from: https://platform.openai.com/api-keys
base_url (Optional[str], optional): Custom API base URL for OpenAI-compatible
services. Defaults to ``None`` (uses official OpenAI endpoint).
Attributes:
dimension (int): The embedding vector dimension.
data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation.
model (str): The OpenAI model name being used.
Raises:
ValueError: If API key is not provided and not found in environment,
or if API returns an error response.
TypeError: If input to ``embed()`` is not a string.
RuntimeError: If network error or OpenAI service error occurs.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires the ``openai`` package: ``pip install openai``
- Embedding results are cached (LRU cache, maxsize=10) to reduce API calls
- Network connectivity to OpenAI API endpoints is required
- API usage incurs costs based on your OpenAI subscription plan
- Rate limits apply based on your OpenAI account tier
Examples:
>>> # Basic usage with default model
>>> from zvec.extension import OpenAIDenseEmbedding
>>> import os
>>> os.environ["OPENAI_API_KEY"] = "sk-..."
>>>
>>> emb_func = OpenAIDenseEmbedding()
>>> vector = emb_func.embed("Hello, world!")
>>> len(vector)
1536
>>> # Using specific model with custom dimension
>>> emb_func = OpenAIDenseEmbedding(
... model="text-embedding-3-large",
... dimension=1024,
... api_key="sk-..."
... )
>>> vector = emb_func.embed("Machine learning is fascinating")
>>> len(vector)
1024
>>> # Using with custom base URL (e.g., Azure OpenAI)
>>> emb_func = OpenAIDenseEmbedding(
... model="text-embedding-ada-002",
... api_key="your-azure-key",
... base_url="https://your-resource.openai.azure.com/"
... )
>>> vector = emb_func("Natural language processing")
>>> isinstance(vector, list)
True
>>> # Batch processing with caching benefit
>>> texts = ["First text", "Second text", "First text"]
>>> vectors = [emb_func.embed(text) for text in texts]
>>> # Third call uses cached result for "First text"
>>> # Error handling
>>> try:
... emb_func.embed("") # Empty string
... except ValueError as e:
... print(f"Error: {e}")
Error: Input text cannot be empty or whitespace only
See Also:
- ``DenseEmbeddingFunction``: Base class for dense embeddings
- ``QwenDenseEmbedding``: Alternative using Qwen/DashScope API
- ``DefaultDenseEmbedding``: Local model without API calls
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
"""
def __init__(
self,
model: str = "text-embedding-3-small",
dimension: Optional[int] = None,
api_key: Optional[str] = None,
base_url: Optional[str] = None,
**kwargs,
):
"""Initialize the OpenAI dense embedding function.
Args:
model (str): OpenAI model name. Defaults to "text-embedding-3-small".
dimension (Optional[int]): Target embedding dimension or None for default.
api_key (Optional[str]): API key or None to use environment variable.
base_url (Optional[str]): Custom API base URL or None for default.
**kwargs: Additional parameters for API calls. Examples:
- ``encoding_format`` (str): Format of embeddings, "float" or "base64".
- ``user`` (str): User identifier for tracking.
Raises:
ValueError: If API key is not provided and not in environment.
"""
# Initialize base class for API connection
OpenAIFunctionBase.__init__(
self, model=model, api_key=api_key, base_url=base_url
)
# Store dimension configuration
self._custom_dimension = dimension
# Determine actual dimension
if dimension is None:
# Use model default dimension
self._dimension = self._MODEL_DIMENSIONS.get(model, 1536)
else:
self._dimension = dimension
# Store dense-specific attributes
self._extra_params = kwargs
@property
def dimension(self) -> int:
"""int: The expected dimensionality of the embedding vector."""
return self._dimension
@property
def extra_params(self) -> dict:
"""dict: Extra parameters for model-specific customization."""
return self._extra_params
def __call__(self, input: TEXT) -> DenseVectorType:
"""Make the embedding function callable."""
return self.embed(input)
@lru_cache(maxsize=10)
def embed(self, input: TEXT) -> DenseVectorType:
"""Generate dense embedding vector for the input text.
This method calls the OpenAI Embeddings API to convert input text
into a dense vector representation. Results are cached to improve
performance for repeated inputs.
Args:
input (TEXT): Input text string to embed. Must be non-empty after
stripping whitespace. Maximum length is 8191 tokens for most models.
Returns:
DenseVectorType: A list of floats representing the embedding vector.
Length equals ``self.dimension``. Example:
``[0.123, -0.456, 0.789, ...]``
Raises:
TypeError: If ``input`` is not a string.
ValueError: If input is empty/whitespace-only, or if the API returns
an error or malformed response.
RuntimeError: If network connectivity issues or OpenAI service
errors occur.
Examples:
>>> emb = OpenAIDenseEmbedding()
>>> vector = emb.embed("Natural language processing")
>>> len(vector)
1536
>>> isinstance(vector[0], float)
True
>>> # Error: empty input
>>> emb.embed(" ")
ValueError: Input text cannot be empty or whitespace only
>>> # Error: non-string input
>>> emb.embed(123)
TypeError: Expected 'input' to be str, got int
Note:
- This method is cached (maxsize=10). Identical inputs return cached results.
- The cache is based on exact string match (case-sensitive).
- Consider pre-processing text (lowercasing, normalization) for better caching.
"""
if not isinstance(input, TEXT):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
# Call API
embedding_vector = self._call_text_embedding_api(
input=input,
dimension=self._custom_dimension,
)
# Verify dimension
if len(embedding_vector) != self.dimension:
raise ValueError(
f"Dimension mismatch: expected {self.dimension}, "
f"got {len(embedding_vector)}"
)
return embedding_vector
+149
View File
@@ -0,0 +1,149 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import os
from typing import ClassVar, Optional
from ..common.constants import TEXT
from ..tool import require_module
class OpenAIFunctionBase:
"""Base class for OpenAI functions.
This base class provides common functionality for calling OpenAI APIs
and handling responses. It supports embeddings (dense) operations.
This class is not meant to be used directly. Use concrete implementations:
- ``OpenAIDenseEmbedding`` for dense embeddings
Args:
model (str): OpenAI model identifier.
api_key (Optional[str]): OpenAI API authentication key.
base_url (Optional[str]): Custom API base URL.
Note:
- This is an internal base class for code reuse across OpenAI features
- Subclasses should inherit from appropriate Protocol
- Provides unified API connection and response handling
"""
# Model default dimensions
_MODEL_DIMENSIONS: ClassVar[dict[str, int]] = {
"text-embedding-3-small": 1536,
"text-embedding-3-large": 3072,
"text-embedding-ada-002": 1536,
}
def __init__(
self,
model: str,
api_key: Optional[str] = None,
base_url: Optional[str] = None,
):
"""Initialize the base OpenAI functionality.
Args:
model (str): OpenAI model name.
api_key (Optional[str]): API key or None to use environment variable.
base_url (Optional[str]): Custom API base URL or None for default.
Raises:
ValueError: If API key is not provided and not in environment.
"""
self._model = model
self._api_key = api_key or os.environ.get("OPENAI_API_KEY")
self._base_url = base_url
if not self._api_key:
raise ValueError(
"OpenAI API key is required. Please provide 'api_key' parameter "
"or set the 'OPENAI_API_KEY' environment variable."
)
@property
def model(self) -> str:
"""str: The OpenAI model name currently in use."""
return self._model
def _get_client(self):
"""Get OpenAI client instance.
Returns:
OpenAI: Configured OpenAI client.
Raises:
ImportError: If openai package is not installed.
"""
openai = require_module("openai")
if self._base_url:
return openai.OpenAI(api_key=self._api_key, base_url=self._base_url)
return openai.OpenAI(api_key=self._api_key)
def _call_text_embedding_api(
self,
input: TEXT,
dimension: Optional[int] = None,
) -> list:
"""Call OpenAI Embeddings API.
Args:
input (TEXT): Input text to embed.
dimension (Optional[int]): Target dimension (for models that support it).
Returns:
list: Embedding vector as list of floats.
Raises:
RuntimeError: If API call fails.
ValueError: If API returns error response.
"""
try:
client = self._get_client()
# Prepare embedding parameters
params = {"model": self.model, "input": input}
# Add dimension parameter for models that support it
if dimension is not None:
params["dimensions"] = dimension
# Call OpenAI API
response = client.embeddings.create(**params)
except Exception as e:
# Check if it's an OpenAI API error
openai = require_module("openai")
if isinstance(e, (openai.APIError, openai.APIConnectionError)):
raise RuntimeError(f"Failed to call OpenAI API: {e!s}") from e
raise RuntimeError(f"Unexpected error during API call: {e!s}") from e
# Extract embedding from response
try:
if not response.data:
raise ValueError("Invalid API response: no embedding data returned")
embedding_vector = response.data[0].embedding
if not isinstance(embedding_vector, list):
raise ValueError(
"Invalid API response: embedding is not a list of numbers"
)
return embedding_vector
except (AttributeError, IndexError, TypeError) as e:
raise ValueError(f"Failed to parse API response: {e!s}") from e
@@ -0,0 +1,537 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from functools import lru_cache
from typing import Optional
from ..common.constants import TEXT, DenseVectorType, SparseVectorType
from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction
from .qwen_function import QwenFunctionBase
class QwenDenseEmbedding(QwenFunctionBase, DenseEmbeddingFunction[TEXT]):
"""Dense text embedding function using Qwen (DashScope) API.
This class provides text-to-vector embedding capabilities using Alibaba Cloud's
DashScope service and Qwen embedding models. It inherits from
``DenseEmbeddingFunction`` and implements dense text embedding.
The implementation supports various Qwen embedding models with configurable
dimensions and includes automatic result caching for improved performance.
Args:
dimension (int): Desired output embedding dimension. Common values:
- 512: Balanced performance and accuracy
- 1024: Higher accuracy, larger storage
- 1536: Maximum accuracy for supported models
model (str, optional): DashScope embedding model identifier.
Defaults to ``"text-embedding-v4"``. Other options include:
- ``"text-embedding-v3"``
- ``"text-embedding-v2"``
- ``"text-embedding-v1"``
api_key (Optional[str], optional): DashScope API authentication key.
If ``None``, reads from ``DASHSCOPE_API_KEY`` environment variable.
Obtain your key from: https://dashscope.console.aliyun.com/
**kwargs: Additional DashScope API parameters. Supported options:
- ``text_type`` (str): Specifies the text role in retrieval tasks.
Options: ``"query"`` (search query) or ``"document"`` (indexed content).
This parameter optimizes embeddings for asymmetric search scenarios.
Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
Attributes:
dimension (int): The embedding vector dimension.
data_type (DataType): Always ``DataType.VECTOR_FP32`` for this implementation.
model (str): The DashScope model name being used.
Raises:
ValueError: If API key is not provided and not found in environment,
or if API returns an error response.
TypeError: If input to ``embed()`` is not a string.
RuntimeError: If network error or DashScope service error occurs.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires the ``dashscope`` package: ``pip install dashscope``
- Embedding results are cached (LRU cache, maxsize=10) to reduce API calls
- Network connectivity to DashScope API endpoints is required
- API usage may incur costs based on your DashScope subscription plan
**Parameter Guidelines:**
- Use ``text_type="query"`` for search queries and ``text_type="document"``
for indexed content to optimize asymmetric retrieval tasks.
- For detailed API specifications and parameter usage, refer to:
https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
Examples:
>>> # Basic usage with default model
>>> from zvec.extension import QwenDenseEmbedding
>>> import os
>>> os.environ["DASHSCOPE_API_KEY"] = "your-api-key"
>>>
>>> emb_func = QwenDenseEmbedding(dimension=1024)
>>> vector = emb_func.embed("Hello, world!")
>>> len(vector)
1024
>>> # Using specific model with explicit API key
>>> emb_func = QwenDenseEmbedding(
... dimension=512,
... model="text-embedding-v3",
... api_key="sk-xxxxx"
... )
>>> vector = emb_func("Machine learning is fascinating")
>>> isinstance(vector, list)
True
>>> # Using with custom parameters (text_type)
>>> # For search queries - optimize for query-document matching
>>> emb_func = QwenDenseEmbedding(
... dimension=1024,
... text_type="query"
... )
>>> query_vector = emb_func.embed("What is machine learning?")
>>>
>>> # For document embeddings - optimize for being matched by queries
>>> doc_emb_func = QwenDenseEmbedding(
... dimension=1024,
... text_type="document"
... )
>>> doc_vector = doc_emb_func.embed(
... "Machine learning is a subset of artificial intelligence..."
... )
>>> # Batch processing with caching benefit
>>> texts = ["First text", "Second text", "First text"]
>>> vectors = [emb_func.embed(text) for text in texts]
>>> # Third call uses cached result for "First text"
>>> # Error handling
>>> try:
... emb_func.embed("") # Empty string
... except ValueError as e:
... print(f"Error: {e}")
Error: Input text cannot be empty or whitespace only
See Also:
- ``DenseEmbeddingFunction``: Base class for dense embeddings
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
"""
def __init__(
self,
dimension: int,
model: str = "text-embedding-v4",
api_key: Optional[str] = None,
**kwargs,
):
"""Initialize the Qwen dense embedding function.
Args:
dimension (int): Target embedding dimension.
model (str): DashScope model name. Defaults to "text-embedding-v4".
api_key (Optional[str]): API key or None to use environment variable.
**kwargs: Additional DashScope API parameters. Supported options:
- ``text_type`` (str): Text role in asymmetric retrieval.
* ``"query"``: Optimize for search queries (short, question-like).
* ``"document"``: Optimize for indexed documents (longer content).
Using appropriate text_type improves retrieval accuracy by
optimizing the embedding space for query-document matching.
For detailed API documentation, see:
https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
Raises:
ValueError: If API key is not provided and not in environment.
"""
# Initialize base class for API connection
QwenFunctionBase.__init__(self, model=model, api_key=api_key)
# Store dense-specific attributes
self._dimension = dimension
self._extra_params = kwargs
@property
def dimension(self) -> int:
"""int: The expected dimensionality of the embedding vector."""
return self._dimension
@property
def extra_params(self) -> dict:
"""dict: Extra parameters for model-specific customization."""
return self._extra_params
def __call__(self, input: TEXT) -> DenseVectorType:
"""Make the embedding function callable."""
return self.embed(input)
@lru_cache(maxsize=10)
def embed(self, input: TEXT) -> DenseVectorType:
"""Generate dense embedding vector for the input text.
This method calls the DashScope TextEmbedding API to convert input text
into a dense vector representation. Results are cached to improve
performance for repeated inputs.
Args:
input (TEXT): Input text string to embed. Must be non-empty after
stripping whitespace. Maximum length depends on the model used
(typically 2048-8192 tokens).
Returns:
DenseVectorType: A list of floats representing the embedding vector.
Length equals ``self.dimension``. Example:
``[0.123, -0.456, 0.789, ...]``
Raises:
TypeError: If ``input`` is not a string.
ValueError: If input is empty/whitespace-only, or if the API returns
an error or malformed response.
RuntimeError: If network connectivity issues or DashScope service
errors occur.
Examples:
>>> emb = QwenDenseEmbedding(dimension=1024)
>>> vector = emb.embed("Natural language processing")
>>> len(vector)
1024
>>> isinstance(vector[0], float)
True
>>> # Error: empty input
>>> emb.embed(" ")
ValueError: Input text cannot be empty or whitespace only
>>> # Error: non-string input
>>> emb.embed(123)
TypeError: Expected 'input' to be str, got int
Note:
- This method is cached (maxsize=10). Identical inputs return cached results.
- The cache is based on exact string match (case-sensitive).
- Consider pre-processing text (lowercasing, normalization) for better caching.
"""
if not isinstance(input, TEXT):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
# Call API with dense output type
output = self._call_text_embedding_api(
input=input,
dimension=self.dimension,
output_type="dense",
text_type=self.extra_params.get("text_type"),
)
embeddings = output.get("embeddings")
if not isinstance(embeddings, list):
raise ValueError(
"Invalid API response: 'embeddings' field is missing or not a list"
)
if len(embeddings) != 1:
raise ValueError(
f"Expected exactly 1 embedding in response, got {len(embeddings)}"
)
first_emb = embeddings[0]
if not isinstance(first_emb, dict):
raise ValueError("Invalid API response: embedding item is not a dictionary")
embedding_vector = first_emb.get("embedding")
if not isinstance(embedding_vector, list):
raise ValueError(
"Invalid API response: 'embedding' field is missing or not a list"
)
if len(embedding_vector) != self.dimension:
raise ValueError(
f"Dimension mismatch: expected {self.dimension}, "
f"got {len(embedding_vector)}"
)
return list(embedding_vector)
class QwenSparseEmbedding(QwenFunctionBase, SparseEmbeddingFunction[TEXT]):
"""Sparse text embedding function using Qwen (DashScope) API.
This class provides text-to-sparse-vector embedding capabilities using
Alibaba Cloud's DashScope service and Qwen embedding models. It generates
sparse keyword-weighted vectors suitable for lexical matching and BM25-style
retrieval scenarios.
Sparse embeddings are particularly useful for:
- Keyword-based search and exact matching
- Hybrid retrieval (combining with dense embeddings)
- Interpretable search results (weights show term importance)
Args:
dimension (int): Desired output embedding dimension. Common values:
- 512: Balanced performance and accuracy
- 1024: Higher accuracy, larger storage
- 1536: Maximum accuracy for supported models
model (str, optional): DashScope embedding model identifier.
Defaults to ``"text-embedding-v4"``. Other options include:
- ``"text-embedding-v3"``
- ``"text-embedding-v2"``
api_key (Optional[str], optional): DashScope API authentication key.
If ``None``, reads from ``DASHSCOPE_API_KEY`` environment variable.
Obtain your key from: https://dashscope.console.aliyun.com/
**kwargs: Additional DashScope API parameters. Supported options:
- ``encoding_type`` (Literal["query", "document"]): Encoding type.
* ``"query"``: Optimize for search queries (default).
* ``"document"``: Optimize for indexed documents.
This distinction is important for asymmetric retrieval tasks.
Attributes:
model (str): The DashScope model name being used.
encoding_type (str): The encoding type ("query" or "document").
Raises:
ValueError: If API key is not provided and not found in environment,
or if API returns an error response.
TypeError: If input to ``embed()`` is not a string.
RuntimeError: If network error or DashScope service error occurs.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires the ``dashscope`` package: ``pip install dashscope``
- Embedding results are cached (LRU cache, maxsize=10) to reduce API calls
- Network connectivity to DashScope API endpoints is required
- API usage may incur costs based on your DashScope subscription plan
- Sparse vectors have only non-zero dimensions stored as dict
- Output is sorted by indices (keys) in ascending order
**Parameter Guidelines:**
- Use ``encoding_type="query"`` for search queries and
``encoding_type="document"`` for indexed content to optimize
asymmetric retrieval tasks.
- For detailed API specifications, refer to:
https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
Examples:
>>> # Basic usage for query embedding
>>> from zvec.extension import QwenSparseEmbedding
>>> import os
>>> os.environ["DASHSCOPE_API_KEY"] = "your-api-key"
>>>
>>> query_emb = QwenSparseEmbedding(dimension=1024, encoding_type="query")
>>> query_vec = query_emb.embed("machine learning")
>>> type(query_vec)
<class 'dict'>
>>> len(query_vec) # Only non-zero dimensions
156
>>> # Document embedding
>>> doc_emb = QwenSparseEmbedding(dimension=1024, encoding_type="document")
>>> doc_vec = doc_emb.embed("Machine learning is a subset of AI")
>>> isinstance(doc_vec, dict)
True
>>> # Asymmetric retrieval example
>>> query_vec = query_emb.embed("what causes aging fast")
>>> doc_vec = doc_emb.embed(
... "UV-A light causes tanning, skin aging, and cataracts..."
... )
>>>
>>> # Calculate similarity (dot product for sparse vectors)
>>> similarity = sum(
... query_vec.get(k, 0) * doc_vec.get(k, 0)
... for k in set(query_vec) | set(doc_vec)
... )
>>> # Output is sorted by indices
>>> list(query_vec.items())[:5] # First 5 dimensions (by index)
[(10, 0.45), (23, 0.87), (56, 0.32), (89, 1.12), (120, 0.65)]
>>> # Hybrid retrieval (combining dense + sparse)
>>> from zvec.extension import QwenDenseEmbedding
>>> dense_emb = QwenDenseEmbedding(dimension=1024)
>>> sparse_emb = QwenSparseEmbedding(dimension=1024)
>>>
>>> query = "deep learning neural networks"
>>> dense_vec = dense_emb.embed(query) # [0.1, -0.3, 0.5, ...]
>>> sparse_vec = sparse_emb.embed(query) # {12: 0.8, 45: 1.2, ...}
>>> # Error handling
>>> try:
... sparse_emb.embed("") # Empty string
... except ValueError as e:
... print(f"Error: {e}")
Error: Input text cannot be empty or whitespace only
See Also:
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
- ``QwenDenseEmbedding``: Dense embedding using Qwen API
- ``DefaultSparseEmbedding``: Sparse embedding with SPLADE model
"""
def __init__(
self,
dimension: int,
model: str = "text-embedding-v4",
api_key: Optional[str] = None,
**kwargs,
):
"""Initialize the Qwen sparse embedding function.
Args:
dimension (int): Target embedding dimension.
model (str): DashScope model name. Defaults to "text-embedding-v4".
api_key (Optional[str]): API key or None to use environment variable.
**kwargs: Additional DashScope API parameters. Supported options:
- ``encoding_type`` (Literal["query", "document"]): Encoding type.
* ``"query"``: Optimize for search queries (default).
* ``"document"``: Optimize for indexed documents.
This distinction is important for asymmetric retrieval tasks.
Raises:
ValueError: If API key is not provided and not in environment.
"""
# Initialize base class for API connection
QwenFunctionBase.__init__(self, model=model, api_key=api_key)
self._dimension = dimension
self._extra_params = kwargs
@property
def extra_params(self) -> dict:
"""dict: Extra parameters for model-specific customization."""
return self._extra_params
def __call__(self, input: TEXT) -> SparseVectorType:
"""Make the embedding function callable."""
return self.embed(input)
@lru_cache(maxsize=10)
def embed(self, input: TEXT) -> SparseVectorType:
"""Generate sparse embedding vector for the input text.
This method calls the DashScope TextEmbedding API with sparse output type
to convert input text into a sparse vector representation. The result is
a dictionary where keys are dimension indices and values are importance
weights (only non-zero values included).
The embedding is optimized based on the ``encoding_type`` specified during
initialization: "query" for search queries or "document" for indexed content.
Args:
input (TEXT): Input text string to embed. Must be non-empty after
stripping whitespace. Maximum length depends on the model used
(typically 2048-8192 tokens).
Returns:
SparseVectorType: A dictionary mapping dimension index to weight.
Only non-zero dimensions are included. The dictionary is sorted
by indices (keys) in ascending order for consistent output.
Example: ``{10: 0.5, 245: 0.8, 1023: 1.2, 5678: 0.5}``
Raises:
TypeError: If ``input`` is not a string.
ValueError: If input is empty/whitespace-only, or if the API returns
an error or malformed response.
RuntimeError: If network connectivity issues or DashScope service
errors occur.
Examples:
>>> emb = QwenSparseEmbedding(dimension=1024, encoding_type="query")
>>> sparse_vec = emb.embed("machine learning")
>>> isinstance(sparse_vec, dict)
True
>>>
>>> # Verify sorted output
>>> keys = list(sparse_vec.keys())
>>> keys == sorted(keys)
True
>>> # Error: empty input
>>> emb.embed(" ")
ValueError: Input text cannot be empty or whitespace only
>>> # Error: non-string input
>>> emb.embed(123)
TypeError: Expected 'input' to be str, got int
Note:
- This method is cached (maxsize=10). Identical inputs return cached results.
- The cache is based on exact string match (case-sensitive).
- Output dictionary is always sorted by indices for consistency.
"""
if not isinstance(input, TEXT):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
# Call API with sparse output type
output = self._call_text_embedding_api(
input=input,
dimension=self._dimension,
output_type="sparse",
text_type=self.extra_params.get("encoding_type", "query"),
)
embeddings = output.get("embeddings")
if not isinstance(embeddings, list):
raise ValueError(
"Invalid API response: 'embeddings' field is missing or not a list"
)
if len(embeddings) != 1:
raise ValueError(
f"Expected exactly 1 embedding in response, got {len(embeddings)}"
)
first_emb = embeddings[0]
if not isinstance(first_emb, dict):
raise ValueError("Invalid API response: embedding item is not a dictionary")
sparse_embedding = first_emb.get("sparse_embedding")
if not isinstance(sparse_embedding, list):
raise ValueError(
"Invalid API response: 'sparse_embedding' field is missing or not a list"
)
# Parse sparse embedding: convert array of {index, value, token} to dict
sparse_dict = {}
for item in sparse_embedding:
if not isinstance(item, dict):
raise ValueError(
"Invalid API response: sparse_embedding item is not a dictionary"
)
index = item.get("index")
value = item.get("value")
if index is None or value is None:
raise ValueError(
"Invalid API response: sparse_embedding item missing 'index' or 'value'"
)
# Convert to int and float, filter positive values
idx = int(index)
val = float(value)
if val > 0:
sparse_dict[idx] = val
# Sort by indices (keys) to ensure consistent ordering
return dict(sorted(sparse_dict.items()))
+186
View File
@@ -0,0 +1,186 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import os
from http import HTTPStatus
from typing import Optional
from ..common.constants import TEXT
from ..tool import require_module
class QwenFunctionBase:
"""Base class for Qwen (DashScope) functions.
This base class provides common functionality for calling DashScope APIs
and handling responses. It supports embeddings (dense and sparse) and
re-ranking operations.
This class is not meant to be used directly. Use concrete implementations:
- ``QwenDenseEmbedding`` for dense embeddings
- ``QwenSparseEmbedding`` for sparse embeddings
- ``QwenReRanker`` for semantic re-ranking
Args:
model (str): DashScope model identifier.
api_key (Optional[str]): DashScope API authentication key.
Note:
- This is an internal base class for code reuse across Qwen features
- Subclasses should inherit from appropriate Protocol/ABC
- Provides unified API connection and response handling
"""
def __init__(
self,
model: str,
api_key: Optional[str] = None,
):
"""Initialize the base Qwen embedding functionality.
Args:
model (str): DashScope model name.
api_key (Optional[str]): API key or None to use environment variable.
Raises:
ValueError: If API key is not provided and not in environment.
"""
self._model = model
self._api_key = api_key or os.environ.get("DASHSCOPE_API_KEY")
if not self._api_key:
raise ValueError(
"DashScope API key is required. Please provide 'api_key' parameter "
"or set the 'DASHSCOPE_API_KEY' environment variable."
)
@property
def model(self) -> str:
"""str: The DashScope embedding model name currently in use."""
return self._model
def _get_connection(self):
"""Establish connection to DashScope API.
Returns:
module: The dashscope module with API key configured.
Raises:
ImportError: If dashscope package is not installed.
"""
dashscope = require_module("dashscope")
dashscope.api_key = self._api_key
return dashscope
def _call_text_embedding_api(
self,
input: TEXT,
dimension: int,
output_type: str,
text_type: Optional[str] = None,
) -> dict:
"""Call DashScope TextEmbedding API.
Args:
input (TEXT): Input text to embed.
dimension (int): Target embedding dimension.
output_type (str): Output type ("dense" or "sparse").
text_type (Optional[str]): Text type ("query" or "document").
Returns:
dict: API response output field.
Raises:
RuntimeError: If API call fails.
ValueError: If API returns error response.
"""
try:
# Prepare API call parameters
call_params = {
"model": self.model,
"input": input,
"dimension": dimension,
"output_type": output_type,
}
# Add optional text_type parameter if provided
if text_type is not None:
call_params["text_type"] = text_type
resp = self._get_connection().TextEmbedding.call(**call_params)
except Exception as e:
raise RuntimeError(f"Failed to call DashScope API: {e!s}") from e
if resp.status_code != HTTPStatus.OK:
error_msg = getattr(resp, "message", "Unknown error")
error_code = getattr(resp, "code", "N/A")
raise ValueError(
f"DashScope API error: [Code={error_code}, "
f"Status={resp.status_code}] {error_msg}"
)
output = getattr(resp, "output", None)
if not isinstance(output, dict):
raise ValueError(
"Invalid API response: missing or malformed 'output' field"
)
return output
def _call_rerank_api(
self,
query: str,
documents: list[str],
top_n: int,
) -> dict:
"""Call DashScope TextReRank API.
Args:
query (str): Query text for semantic matching.
documents (list[str]): List of document texts to re-rank.
top_n (int): Maximum number of documents to return.
Returns:
dict: API response output field containing re-ranked results.
Raises:
RuntimeError: If API call fails.
ValueError: If API returns error response.
"""
try:
resp = self._get_connection().TextReRank.call(
model=self.model,
query=query,
documents=documents,
top_n=top_n,
return_documents=False,
)
except Exception as e:
raise RuntimeError(f"Failed to call DashScope API: {e!s}") from e
if resp.status_code != HTTPStatus.OK:
error_msg = getattr(resp, "message", "Unknown error")
error_code = getattr(resp, "code", "N/A")
raise ValueError(
f"DashScope API error: [Code={error_code}, "
f"Status={resp.status_code}] {error_msg}"
)
output = getattr(resp, "output", None)
if not isinstance(output, dict):
raise ValueError(
"Invalid API response: missing or malformed 'output' field"
)
return output
@@ -0,0 +1,177 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
from ..model.doc import Doc, DocList
from .qwen_function import QwenFunctionBase
from .rerank_function import RerankFunction
if TYPE_CHECKING:
from ..model.schema import FieldSchema, VectorSchema
class QwenReRanker(QwenFunctionBase, RerankFunction):
"""Re-ranker using Qwen (DashScope) cross-encoder API for semantic re-ranking.
This re-ranker leverages DashScope's TextReRank service to perform
cross-encoder style re-ranking. It sends query and document pairs to the
API and receives relevance scores based on deep semantic understanding.
The re-ranker is suitable for single-vector or multi-vector search scenarios
where semantic relevance to a specific query is required.
Args:
query (str): Query text for semantic re-ranking. **Required**.
rerank_field (str): Document field name to use as re-ranking input text.
**Required** (e.g., "content", "title", "body").
model (str, optional): DashScope re-ranking model identifier.
Defaults to ``"gte-rerank-v2"``.
api_key (Optional[str], optional): DashScope API authentication key.
If not provided, reads from ``DASHSCOPE_API_KEY`` environment variable.
Raises:
ValueError: If ``query`` is empty/None, ``rerank_field`` is None,
or API key is not available.
Note:
- Requires ``dashscope`` Python package installed
- Documents without valid content in ``rerank_field`` are skipped
- API rate limits and quotas apply per DashScope subscription
Example:
>>> reranker = QwenReRanker(
... query="machine learning algorithms",
... rerank_field="content",
... model="gte-rerank-v2",
... api_key="your-api-key"
... )
>>> # Use in collection.query(reranker=reranker)
"""
def __init__(
self,
query: Optional[str] = None,
rerank_field: Optional[str] = None,
model: str = "gte-rerank-v2",
api_key: Optional[str] = None,
):
"""Initialize QwenReRanker with query and configuration.
Args:
query (Optional[str]): Query text for semantic matching. Required.
rerank_field (Optional[str]): Document field for re-ranking input.
model (str): DashScope model name.
api_key (Optional[str]): API key or None to use environment variable.
Raises:
ValueError: If query is empty or API key is unavailable.
"""
QwenFunctionBase.__init__(self, model=model, api_key=api_key)
self._rerank_field = rerank_field
if not query:
raise ValueError("Query is required for QwenReRanker")
self._query = query
@property
def rerank_field(self) -> Optional[str]:
"""Optional[str]: Field name used as re-ranking input."""
return self._rerank_field
@property
def query(self) -> str:
"""str: Query text used for semantic re-ranking."""
return self._query
def rerank(
self,
query_results: list[list[Doc]],
topn: int = 10,
*,
fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002
) -> DocList:
"""Re-rank documents using Qwen's TextReRank API.
Sends document texts to DashScope TextReRank service along with the query.
Returns documents sorted by relevance scores from the cross-encoder model.
Args:
query_results (list[list[Doc]]): Per-sub-query lists of retrieved
documents. Documents from all lists are deduplicated and
re-ranked together.
topn (int): Maximum number of documents to return.
fields: Unused; present for interface compatibility.
Returns:
list[Doc]: Re-ranked documents (up to ``topn``) with updated ``score``
fields containing relevance scores from the API.
Raises:
ValueError: If no valid documents are found or API call fails.
Note:
- Duplicate documents (same ID) across lists are processed once
- Documents with empty/missing ``rerank_field`` content are skipped
- Returned scores are relevance scores from the cross-encoder model
"""
if not query_results:
return []
# Accept both dict (legacy) and list formats
if isinstance(query_results, dict):
query_results = list(query_results.values())
# Collect and deduplicate documents
id_to_doc: dict[str, Doc] = {}
doc_ids: list[str] = []
contents: list[str] = []
for query_result in query_results:
for doc in query_result:
doc_id = doc.id
if doc_id in id_to_doc:
continue
# Extract text content from specified field
field_value = doc.field(self.rerank_field)
rank_content = str(field_value).strip() if field_value else ""
if not rank_content:
continue
id_to_doc[doc_id] = doc
doc_ids.append(doc_id)
contents.append(rank_content)
if not contents:
raise ValueError("No documents to rerank")
# Call DashScope TextReRank API
output = self._call_rerank_api(
query=self.query,
documents=contents,
top_n=topn,
)
# Build result list with updated scores
results: DocList = []
for item in output["results"]:
idx = item["index"]
doc_id = doc_ids[idx]
doc = id_to_doc[doc_id]
new_doc = doc._replace(score=item["relevance_score"])
results.append(new_doc)
return results
+56
View File
@@ -0,0 +1,56 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
from ..model.doc import Doc, DocList
if TYPE_CHECKING:
from ..model.schema import FieldSchema, VectorSchema
class RerankFunction(ABC):
"""Abstract base class for reranker parameter containers.
Subclasses define rerank parameters and implement _to_cpp_params()
for conversion to C++ parameter structs (used by collection fast path).
Each subclass also provides a standalone rerank() implementation.
"""
def _to_cpp_params(self):
"""Return C++ reranker params. Override in subclasses that use C++ path."""
raise NotImplementedError
@abstractmethod
def rerank(
self,
query_results: list[list[Doc]],
topn: int = 10,
*,
fields: list[FieldSchema | VectorSchema] | None = None,
) -> DocList:
"""Execute rerank on sub-query results.
Args:
query_results: List of per-sub-query document lists.
topn: Maximum number of results to return.
fields: Per-sub-query Python FieldSchema/VectorSchema objects
(required for WeightedReRanker score normalization).
Returns:
Re-ranked document list.
"""
...
@@ -0,0 +1,839 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import ClassVar, Literal, Optional
import numpy as np
from ..common.constants import TEXT, DenseVectorType, SparseVectorType
from .embedding_function import DenseEmbeddingFunction, SparseEmbeddingFunction
from .sentence_transformer_function import SentenceTransformerFunctionBase
class DefaultLocalDenseEmbedding(
SentenceTransformerFunctionBase, DenseEmbeddingFunction[TEXT]
):
"""Default local dense embedding using all-MiniLM-L6-v2 model.
This is the default implementation for dense text embedding that uses the
``all-MiniLM-L6-v2`` model from Hugging Face by default. This model provides
a good balance between speed and quality for general-purpose text embedding.
The class provides text-to-vector dense embedding capabilities using the
sentence-transformers library. It supports models from Hugging Face Hub and
ModelScope, runs locally without API calls, and supports CPU/GPU acceleration.
The model produces 384-dimensional embeddings and is optimized for semantic
similarity tasks. It runs locally without requiring API keys.
Args:
model_source (Literal["huggingface", "modelscope"], optional): Model source.
- ``"huggingface"``: Use Hugging Face Hub (default, for international users)
- ``"modelscope"``: Use ModelScope (recommended for users in China)
Defaults to ``"huggingface"``.
device (Optional[str], optional): Device to run the model on.
Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None``
for automatic detection. Defaults to ``None``.
normalize_embeddings (bool, optional): Whether to normalize embeddings to
unit length (L2 normalization). Useful for cosine similarity.
Defaults to ``True``.
batch_size (int, optional): Batch size for encoding. Defaults to ``32``.
**kwargs: Additional parameters for future extension.
Attributes:
dimension (int): Always 384 for both models.
model_name (str): "all-MiniLM-L6-v2" (HF) or "iic/nlp_gte_sentence-embedding_chinese-small" (MS).
model_source (str): The model source being used.
device (str): The device the model is running on.
Raises:
ValueError: If the model cannot be loaded or input is invalid.
TypeError: If input to ``embed()`` is not a string.
RuntimeError: If model inference fails.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires the ``sentence-transformers`` package:
``pip install sentence-transformers``
- For ModelScope, also requires: ``pip install modelscope``
- First run downloads the model (~50-80MB) from chosen source
- Hugging Face cache: ``~/.cache/torch/sentence_transformers/``
- ModelScope cache: ``~/.cache/modelscope/hub/``
- No API keys or network required after initial download
- Inference speed: ~1000 sentences/sec on CPU, ~10000 on GPU
**For users in China:**
If you encounter Hugging Face access issues, use ModelScope instead:
.. code-block:: python
# Recommended for users in China
emb = DefaultLocalDenseEmbedding(model_source="modelscope")
Alternatively, use Hugging Face mirror:
.. code-block:: bash
export HF_ENDPOINT=https://hf-mirror.com
# Then use default Hugging Face mode
Examples:
>>> # Basic usage with Hugging Face (default)
>>> from zvec.extension import DefaultLocalDenseEmbedding
>>>
>>> emb_func = DefaultLocalDenseEmbedding()
>>> vector = emb_func.embed("Hello, world!")
>>> len(vector)
384
>>> isinstance(vector, list)
True
>>> # Recommended for users in China (uses ModelScope)
>>> emb_func = DefaultLocalDenseEmbedding(model_source="modelscope")
>>> vector = emb_func.embed("你好,世界!") # Works well with Chinese text
>>> len(vector)
384
>>> # Alternative for China users: Use Hugging Face mirror
>>> import os
>>> os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
>>> emb_func = DefaultLocalDenseEmbedding() # Uses HF mirror
>>> vector = emb_func.embed("Hello, world!")
>>> # Using GPU for faster inference
>>> emb_func = DefaultLocalDenseEmbedding(device="cuda")
>>> vector = emb_func("Machine learning is fascinating")
>>> # Normalized vector has unit length
>>> import numpy as np
>>> np.linalg.norm(vector)
1.0
>>> # Batch processing
>>> texts = ["First text", "Second text", "Third text"]
>>> vectors = [emb_func.embed(text) for text in texts]
>>> len(vectors)
3
>>> all(len(v) == 384 for v in vectors)
True
>>> # Semantic similarity
>>> v1 = emb_func.embed("The cat sits on the mat")
>>> v2 = emb_func.embed("A feline rests on a rug")
>>> v3 = emb_func.embed("Python programming")
>>> similarity_high = np.dot(v1, v2) # Similar sentences
>>> similarity_low = np.dot(v1, v3) # Different topics
>>> similarity_high > similarity_low
True
>>> # Error handling
>>> try:
... emb_func.embed("") # Empty string
... except ValueError as e:
... print(f"Error: {e}")
Error: Input text cannot be empty or whitespace only
See Also:
- ``DenseEmbeddingFunction``: Base class for dense embeddings
- ``DefaultLocalSparseEmbedding``: Sparse embedding with SPLADE
- ``QwenDenseEmbedding``: Alternative using Qwen API
"""
def __init__(
self,
model_source: Literal["huggingface", "modelscope"] = "huggingface",
device: Optional[str] = None,
normalize_embeddings: bool = True,
batch_size: int = 32,
**kwargs,
):
"""Initialize with all-MiniLM-L6-v2 model.
Args:
model_source (Literal["huggingface", "modelscope"]): Model source.
Defaults to "huggingface".
device (Optional[str]): Target device ("cpu", "cuda", "mps", or None).
Defaults to None (automatic detection).
normalize_embeddings (bool): Whether to L2-normalize output vectors.
Defaults to True.
batch_size (int): Batch size for encoding. Defaults to 32.
**kwargs: Additional parameters for future extension.
Raises:
ImportError: If sentence-transformers or modelscope is not installed.
ValueError: If model cannot be loaded.
"""
# Use different models based on source
if model_source == "modelscope":
# Use Chinese-optimized model for ModelScope (better for Chinese text)
model_name = "iic/nlp_gte_sentence-embedding_chinese-small"
else:
model_name = "all-MiniLM-L6-v2"
# Initialize base class for model loading
SentenceTransformerFunctionBase.__init__(
self, model_name=model_name, model_source=model_source, device=device
)
self._normalize_embeddings = normalize_embeddings
self._batch_size = batch_size
# Load model and get dimension
model = self._get_model()
self._dimension = model.get_sentence_embedding_dimension()
# Store extra parameters
self._extra_params = kwargs
@property
def dimension(self) -> int:
"""int: The expected dimensionality of the embedding vector."""
return self._dimension
@property
def extra_params(self) -> dict:
"""dict: Extra parameters for model-specific customization."""
return self._extra_params
def __call__(self, input: str) -> DenseVectorType:
"""Make the embedding function callable."""
return self.embed(input)
def embed(self, input: str) -> DenseVectorType:
"""Generate dense embedding vector for the input text.
This method uses the Sentence Transformer model to convert input text
into a dense vector representation. The model runs locally without
requiring API calls.
Args:
input (str): Input text string to embed. Must be non-empty after
stripping whitespace. Maximum length depends on the model used
(typically 128-512 tokens for most models).
Returns:
DenseVectorType: A list of floats representing the embedding vector.
Length equals ``self.dimension``. If ``normalize_embeddings=True``,
the vector has unit length. Example:
``[0.123, -0.456, 0.789, ...]``
Raises:
TypeError: If ``input`` is not a string.
ValueError: If input is empty or whitespace-only.
RuntimeError: If model inference fails.
Examples:
>>> emb = DefaultLocalDenseEmbedding()
>>> vector = emb.embed("Natural language processing")
>>> len(vector)
384
>>> isinstance(vector[0], float)
True
>>> # Normalized vectors have unit length
>>> import numpy as np
>>> emb = DefaultLocalDenseEmbedding(normalize_embeddings=True)
>>> vector = emb.embed("Test sentence")
>>> np.linalg.norm(vector)
1.0
>>> # Error: empty input
>>> emb.embed(" ")
ValueError: Input text cannot be empty or whitespace only
>>> # Error: non-string input
>>> emb.embed(123)
TypeError: Expected 'input' to be str, got int
>>> # Semantic similarity example
>>> v1 = emb.embed("The cat sits on the mat")
>>> v2 = emb.embed("A feline rests on a rug")
>>> similarity = np.dot(v1, v2) # High similarity due to semantic meaning
>>> similarity > 0.7
True
Note:
- First call may be slower due to model loading
- Subsequent calls are much faster as the model stays in memory
- For batch processing, consider encoding multiple texts together
(though this method handles single texts only)
- GPU acceleration provides 5-10x speedup over CPU
"""
if not isinstance(input, str):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
try:
model = self._get_model()
embedding = model.encode(
input,
convert_to_numpy=True,
normalize_embeddings=self._normalize_embeddings,
batch_size=self._batch_size,
)
# Convert numpy array to list
if isinstance(embedding, np.ndarray):
embedding_list = embedding.tolist()
else:
embedding_list = list(embedding)
# Validate dimension
if len(embedding_list) != self.dimension:
raise ValueError(
f"Dimension mismatch: expected {self.dimension}, "
f"got {len(embedding_list)}"
)
return embedding_list
except Exception as e:
if isinstance(e, (TypeError, ValueError)):
raise
raise RuntimeError(f"Failed to generate embedding: {e!s}") from e
class DefaultLocalSparseEmbedding(
SentenceTransformerFunctionBase, SparseEmbeddingFunction[TEXT]
):
"""Default local sparse embedding using SPLADE model.
This class provides sparse vector embedding using the SPLADE (SParse Lexical
AnD Expansion) model. SPLADE generates sparse, interpretable representations
where each dimension corresponds to a vocabulary term with learned importance
weights. It's ideal for lexical matching, BM25-style retrieval, and hybrid
search scenarios.
The default model is ``naver/splade-cocondenser-ensembledistil``, which is
publicly available without authentication. It produces sparse vectors with
thousands of dimensions but only hundreds of non-zero values, making them
efficient for storage and retrieval while maintaining strong lexical matching.
**Model Caching:**
This class uses class-level caching to share the SPLADE model across all instances
with the same configuration (model_source, device). This significantly reduces
memory usage when creating multiple instances for different encoding types
(query vs document).
**Cache Management:**
The class provides methods to manage the model cache:
- ``clear_cache()``: Clear all cached models to free memory
- ``get_cache_info()``: Get information about cached models
- ``remove_from_cache(model_source, device)``: Remove a specific model from cache
.. note::
**Why not use splade-v3?**
The newer ``naver/splade-v3`` model is gated (requires access approval).
We use ``naver/splade-cocondenser-ensembledistil`` instead.
**To use splade-v3 (if you have access):**
1. Request access at https://huggingface.co/naver/splade-v3
2. Get your Hugging Face token from https://huggingface.co/settings/tokens
3. Set environment variable:
.. code-block:: bash
export HF_TOKEN="your_huggingface_token"
4. Or login programmatically:
.. code-block:: python
from huggingface_hub import login
login(token="your_huggingface_token")
5. To use a custom SPLADE model, you can subclass this class and override
the model_name in ``__init__``, or create your own implementation
inheriting from ``SentenceTransformerFunctionBase`` and
``SparseEmbeddingFunction``.
Args:
model_source (Literal["huggingface", "modelscope"], optional): Model source.
Defaults to ``"huggingface"``. ModelScope support may vary for SPLADE models.
device (Optional[str], optional): Device to run the model on.
Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None``
for automatic detection. Defaults to ``None``.
encoding_type (Literal["query", "document"], optional): Encoding type.
- ``"query"``: Optimize for search queries (default)
- ``"document"``: Optimize for indexed documents
**kwargs: Additional parameters (currently unused, for future extension).
Attributes:
model_name (str): Model identifier.
model_source (str): The model source being used.
device (str): The device the model is running on.
Raises:
ValueError: If the model cannot be loaded or input is invalid.
TypeError: If input to ``embed()`` is not a string.
RuntimeError: If model inference fails.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires the ``sentence-transformers`` package:
``pip install sentence-transformers``
- First run downloads the model (~100MB) from Hugging Face
- Cache location: ``~/.cache/torch/sentence_transformers/``
- No API keys or authentication required
- Sparse vectors have ~30k dimensions but only ~100-200 non-zero values
- Best combined with dense embeddings for hybrid retrieval
**SPLADE vs Dense Embeddings:**
- **Dense**: Continuous semantic vectors, good for semantic similarity
- **Sparse**: Lexical keyword-based, interpretable, good for exact matching
- **Hybrid**: Combine both for best retrieval performance
Examples:
>>> # Memory-efficient: both instances share the same model (~200MB)
>>> from zvec.extension import DefaultLocalSparseEmbedding
>>>
>>> # Query embedding
>>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query")
>>> query_vec = query_emb.embed("machine learning algorithms")
>>> type(query_vec)
<class 'dict'>
>>> len(query_vec) # Only non-zero dimensions
156
>>> # Document embedding (shares model with query_emb)
>>> doc_emb = DefaultLocalSparseEmbedding(encoding_type="document")
>>> doc_vec = doc_emb.embed("Machine learning is a subset of AI")
>>> # Total memory: ~200MB (not 400MB) thanks to model caching
>>> # Asymmetric retrieval example
>>> query_vec = query_emb.embed("what causes aging fast")
>>> doc_vec = doc_emb.embed(
... "UV-A light causes tanning, skin aging, and cataracts..."
... )
>>>
>>> # Calculate similarity (dot product for sparse vectors)
>>> similarity = sum(
... query_vec.get(k, 0) * doc_vec.get(k, 0)
... for k in set(query_vec) | set(doc_vec)
... )
>>> # Batch processing
>>> queries = ["query 1", "query 2", "query 3"]
>>> query_vecs = [query_emb.embed(q) for q in queries]
>>>
>>> documents = ["doc 1", "doc 2", "doc 3"]
>>> doc_vecs = [doc_emb.embed(d) for d in documents]
>>> # Inspecting sparse dimensions (output is sorted by indices)
>>> query_vec = query_emb.embed("machine learning")
>>> list(query_vec.items())[:5] # First 5 dimensions (by index)
[(10, 0.45), (23, 0.87), (56, 0.32), (89, 1.12), (120, 0.65)]
>>>
>>> # Sort by weight to find most important terms
>>> sorted_by_weight = sorted(query_vec.items(), key=lambda x: x[1], reverse=True)
>>> top_5 = sorted_by_weight[:5] # Top 5 most important terms
>>> top_5
[(1023, 1.45), (245, 1.23), (8901, 0.98), (5678, 0.87), (12034, 0.76)]
>>> # Using GPU for faster inference
>>> sparse_emb = DefaultLocalSparseEmbedding(device="cuda")
>>> vector = sparse_emb.embed("natural language processing")
>>> # Hybrid retrieval example (combining dense + sparse)
>>> from zvec.extension import DefaultDenseEmbedding
>>> dense_emb = DefaultDenseEmbedding()
>>> sparse_emb = DefaultLocalSparseEmbedding()
>>>
>>> query = "deep learning neural networks"
>>> dense_vec = dense_emb.embed(query) # [0.1, -0.3, 0.5, ...]
>>> sparse_vec = sparse_emb.embed(query) # {12: 0.8, 45: 1.2, ...}
>>> # Error handling
>>> try:
... sparse_emb.embed("") # Empty string
... except ValueError as e:
... print(f"Error: {e}")
Error: Input text cannot be empty or whitespace only
>>> # Cache management
>>> # Check cache status
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 1
>>>
>>> # Clear cache to free memory
>>> DefaultLocalSparseEmbedding.clear_cache()
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 0
>>>
>>> # Remove specific model from cache
>>> query_emb = DefaultLocalSparseEmbedding() # Creates CPU model
>>> cuda_emb = DefaultLocalSparseEmbedding(device="cuda") # Creates CUDA model
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 2
>>>
>>> # Remove only CPU model
>>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device=None)
>>> print(f"Removed: {removed}")
True
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 1
See Also:
- ``SparseEmbeddingFunction``: Base class for sparse embeddings
- ``DefaultDenseEmbedding``: Dense embedding with all-MiniLM-L6-v2
- ``QwenDenseEmbedding``: Alternative using Qwen API
References:
- SPLADE Paper: https://arxiv.org/abs/2109.10086
- Model: https://huggingface.co/naver/splade-cocondenser-ensembledistil
"""
# Class-level model cache: {(model_name, model_source, device): model}
# Shared across all DefaultLocalSparseEmbedding instances to save memory
_model_cache: ClassVar[dict] = {}
@classmethod
def clear_cache(cls) -> None:
"""Clear all cached SPLADE models from memory.
This is useful for:
- Freeing memory when models are no longer needed
- Forcing a fresh model reload
- Testing and debugging
Examples:
>>> # Clear cache to free memory
>>> DefaultLocalSparseEmbedding.clear_cache()
>>> # Or in tests to ensure fresh model loading
>>> def test_something():
... DefaultLocalSparseEmbedding.clear_cache()
... emb = DefaultLocalSparseEmbedding()
... # Test with fresh model
"""
cls._model_cache.clear()
@classmethod
def get_cache_info(cls) -> dict:
"""Get information about currently cached models.
Returns:
dict: Dictionary with cache statistics:
- cached_models (int): Number of cached model instances
- cache_keys (list): List of cache keys (model_name, model_source, device)
Examples:
>>> info = DefaultLocalSparseEmbedding.get_cache_info()
>>> print(f"Cached models: {info['cached_models']}")
Cached models: 2
>>> print(f"Cache keys: {info['cache_keys']}")
Cache keys: [('naver/splade-cocondenser-ensembledistil', 'huggingface', None),
('naver/splade-cocondenser-ensembledistil', 'huggingface', 'cuda')]
"""
return {
"cached_models": len(cls._model_cache),
"cache_keys": list(cls._model_cache.keys()),
}
@classmethod
def remove_from_cache(
cls, model_source: str = "huggingface", device: Optional[str] = None
) -> bool:
"""Remove a specific model from cache.
Args:
model_source (str): Model source ("huggingface" or "modelscope").
Defaults to "huggingface".
device (Optional[str]): Device identifier. Defaults to None.
Returns:
bool: True if model was found and removed, False otherwise.
Examples:
>>> # Remove CPU model from cache
>>> removed = DefaultLocalSparseEmbedding.remove_from_cache()
>>> print(f"Removed: {removed}")
True
>>> # Remove CUDA model from cache
>>> removed = DefaultLocalSparseEmbedding.remove_from_cache(device="cuda")
>>> print(f"Removed: {removed}")
True
"""
model_name = "naver/splade-cocondenser-ensembledistil"
cache_key = (model_name, model_source, device)
if cache_key in cls._model_cache:
del cls._model_cache[cache_key]
return True
return False
def __init__(
self,
model_source: Literal["huggingface", "modelscope"] = "huggingface",
device: Optional[str] = None,
encoding_type: Literal["query", "document"] = "query",
**kwargs,
):
"""Initialize with SPLADE model.
Args:
model_source (Literal["huggingface", "modelscope"]): Model source.
Defaults to "huggingface".
device (Optional[str]): Target device ("cpu", "cuda", "mps", or None).
Defaults to None (automatic detection).
encoding_type (Literal["query", "document"]): Encoding type for embeddings.
- "query": Optimize for search queries (default)
- "document": Optimize for indexed documents
This distinction is important for asymmetric retrieval tasks.
**kwargs: Additional parameters (reserved for future use).
Raises:
ImportError: If sentence-transformers is not installed.
ValueError: If model cannot be loaded.
Note:
Multiple instances with the same (model_source, device) configuration
will share the same underlying model to save memory. Different
instances can use different encoding_type settings while sharing
the model.
**Model Selection:**
Uses ``naver/splade-cocondenser-ensembledistil`` instead of the newer
``naver/splade-v3`` because splade-v3 is a gated model requiring
Hugging Face authentication. The cocondenser-ensembledistil variant:
- Does not require authentication or API tokens
- Is immediately available for all users
- Provides comparable retrieval performance (~2% difference)
- Avoids "Access to model is restricted" errors
If you need splade-v3 and have obtained access, you can subclass
this class and override the model_name parameter.
Examples:
>>> # Both instances share the same model (saves memory)
>>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query")
>>> doc_emb = DefaultLocalSparseEmbedding(encoding_type="document")
>>> # Only one model is loaded in memory
"""
# Use publicly available SPLADE model (no gated access required)
# Note: naver/splade-v3 requires authentication, so we use the
# cocondenser-ensembledistil variant which is publicly accessible
model_name = "naver/splade-cocondenser-ensembledistil"
# Initialize base class for model loading
SentenceTransformerFunctionBase.__init__(
self, model_name=model_name, model_source=model_source, device=device
)
self._encoding_type = encoding_type
self._extra_params = kwargs
# Create cache key for this model configuration
self._cache_key = (model_name, model_source, device)
# Load model to ensure it's available (will use cache if exists)
self._get_model()
@property
def extra_params(self) -> dict:
"""dict: Extra parameters for model-specific customization."""
return self._extra_params
def __call__(self, input: str) -> SparseVectorType:
"""Make the embedding function callable."""
return self.embed(input)
def embed(self, input: str) -> SparseVectorType:
"""Generate sparse embedding vector for the input text.
This method uses the SPLADE model to convert input text into a sparse
vector representation. The result is a dictionary where keys are dimension
indices and values are importance weights (only non-zero values included).
The embedding is optimized based on the ``encoding_type`` specified during
initialization: "query" for search queries or "document" for indexed content.
Args:
input (str): Input text string to embed. Must be non-empty after
stripping whitespace.
Returns:
SparseVectorType: A dictionary mapping dimension index to weight.
Only non-zero dimensions are included. The dictionary is sorted
by indices (keys) in ascending order for consistent output.
Example: ``{10: 0.5, 245: 0.8, 1023: 1.2, 5678: 0.5}``
Raises:
TypeError: If ``input`` is not a string.
ValueError: If input is empty or whitespace-only.
RuntimeError: If model inference fails.
Examples:
>>> # Query embedding
>>> query_emb = DefaultLocalSparseEmbedding(encoding_type="query")
>>> query_vec = query_emb.embed("machine learning")
>>> isinstance(query_vec, dict)
True
Note:
- First call may be slower due to model loading
- Subsequent calls are much faster as the model stays in memory
- GPU acceleration provides significant speedup
- Sparse vectors are memory-efficient (only store non-zero values)
"""
if not isinstance(input, str):
raise TypeError(f"Expected 'input' to be str, got {type(input).__name__}")
input = input.strip()
if not input:
raise ValueError("Input text cannot be empty or whitespace only")
try:
model = self._get_model()
# Use appropriate encoding method based on type
if self._encoding_type == "document" and hasattr(model, "encode_document"):
# Use document encoding
sparse_matrix = model.encode_document([input])
elif hasattr(model, "encode_query"):
# Use query encoding (default)
sparse_matrix = model.encode_query([input])
else:
# Fallback: manual implementation for older sentence-transformers
return self._manual_sparse_encode(input)
# Convert sparse matrix to dictionary
# SPLADE returns shape [1, vocab_size] for single input
# Check if it's a sparse matrix (duck typing - has toarray method)
if hasattr(sparse_matrix, "toarray"):
# Sparse matrix (CSR/CSC/etc.) - convert to dense array
sparse_array = sparse_matrix[0].toarray().flatten()
sparse_dict = {
int(idx): float(val)
for idx, val in enumerate(sparse_array)
if val > 0
}
else:
# Dense array format (numpy array or similar)
if isinstance(sparse_matrix, np.ndarray):
sparse_array = sparse_matrix[0]
else:
sparse_array = sparse_matrix
sparse_dict = {
int(idx): float(val)
for idx, val in enumerate(sparse_array)
if val > 0
}
# Sort by indices (keys) to ensure consistent ordering
return dict(sorted(sparse_dict.items()))
except Exception as e:
if isinstance(e, (TypeError, ValueError)):
raise
raise RuntimeError(f"Failed to generate sparse embedding: {e!s}") from e
def _manual_sparse_encode(self, input: str) -> SparseVectorType:
"""Fallback manual SPLADE encoding for older sentence-transformers.
Args:
input (str): Input text to encode.
Returns:
SparseVectorType: Sparse vector as dictionary.
"""
import torch
model = self._get_model()
# Tokenize input
features = model.tokenize([input])
# Move to correct device
features = {k: v.to(model.device) for k, v in features.items()}
# Forward pass with no gradient
with torch.no_grad():
embeddings = model.forward(features)
# Get logits from model output
# SPLADE models typically output 'token_embeddings'
if isinstance(embeddings, dict) and "token_embeddings" in embeddings:
logits = embeddings["token_embeddings"][0] # First batch item
elif hasattr(embeddings, "token_embeddings"):
logits = embeddings.token_embeddings[0]
# Fallback: try to get first value
elif isinstance(embeddings, dict):
logits = next(iter(embeddings.values()))[0]
else:
logits = embeddings[0]
# Apply SPLADE activation: log(1 + relu(x))
relu_log = torch.log(1 + torch.relu(logits))
# Max pooling over token dimension (reduce to vocab size)
if relu_log.dim() > 1:
sparse_vec, _ = torch.max(relu_log, dim=0)
else:
sparse_vec = relu_log
# Convert to sparse dictionary (only non-zero values)
sparse_vec_np = sparse_vec.cpu().numpy()
sparse_dict = {
int(idx): float(val) for idx, val in enumerate(sparse_vec_np) if val > 0
}
# Sort by indices (keys) to ensure consistent ordering
return dict(sorted(sparse_dict.items()))
def _get_model(self):
"""Load or retrieve the SPLADE model from class-level cache.
Returns:
SentenceTransformer: The loaded SPLADE model instance.
Raises:
ImportError: If required packages are not installed.
ValueError: If model cannot be loaded.
Note:
Models are cached at class level and shared across all instances
with the same (model_name, model_source, device) configuration.
This allows memory-efficient usage when creating multiple instances
with different encoding_type settings.
"""
# Check class-level cache first
if self._cache_key in self._model_cache:
return self._model_cache[self._cache_key]
# Use parent class method to load model
model = super()._get_model()
# Cache the model at class level
self._model_cache[self._cache_key] = model
return model
@@ -0,0 +1,150 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Literal, Optional
from ..tool import require_module
class SentenceTransformerFunctionBase:
"""Base class for Sentence Transformer functions (both dense and sparse).
This base class provides common functionality for loading and managing
sentence-transformers models from Hugging Face or ModelScope. It supports
both dense models (e.g., all-MiniLM-L6-v2) and sparse models (e.g., SPLADE).
This class is not meant to be used directly. Use concrete implementations:
- ``SentenceTransformerEmbeddingFunction`` for dense embeddings
- ``SentenceTransformerSparseEmbeddingFunction`` for sparse embeddings
- ``DefaultDenseEmbedding`` for default dense embeddings
- ``DefaultSparseEmbedding`` for default sparse embeddings
Args:
model_name (str): Model identifier or local path.
model_source (Literal["huggingface", "modelscope"]): Model source.
device (Optional[str]): Device to run the model on.
Note:
- This is an internal base class for code reuse
- Subclasses should inherit from appropriate Protocol (Dense/Sparse)
- Provides model loading and management functionality
"""
def __init__(
self,
model_name: str,
model_source: Literal["huggingface", "modelscope"] = "huggingface",
device: Optional[str] = None,
):
"""Initialize the base Sentence Transformer functionality.
Args:
model_name (str): Model identifier or local path.
model_source (Literal["huggingface", "modelscope"]): Model source.
device (Optional[str]): Device to run the model on.
Raises:
ValueError: If model_source is invalid.
"""
# Validate model_source
if model_source not in ("huggingface", "modelscope"):
raise ValueError(
f"Invalid model_source: '{model_source}'. "
"Must be 'huggingface' or 'modelscope'."
)
self._model_name = model_name
self._model_source = model_source
self._device = device
self._model = None
@property
def model_name(self) -> str:
"""str: The Sentence Transformer model name currently in use."""
return self._model_name
@property
def model_source(self) -> str:
"""str: The model source being used ("huggingface" or "modelscope")."""
return self._model_source
@property
def device(self) -> str:
"""str: The device the model is running on."""
model = self._get_model()
if model is not None:
return str(model.device)
return self._device or "cpu"
def _get_model(self):
"""Load or retrieve the Sentence Transformer model.
Returns:
SentenceTransformer or SparseEncoder: The loaded model instance.
Raises:
ImportError: If required packages are not installed.
ValueError: If model cannot be loaded.
"""
# Return cached model if exists
if self._model is not None:
return self._model
# Load model
try:
sentence_transformers = require_module("sentence_transformers")
if self._model_source == "modelscope":
# Load from ModelScope
require_module("modelscope")
from modelscope.hub.snapshot_download import snapshot_download
# Download model to cache
model_dir = snapshot_download(self._model_name)
# Load from local path
self._model = sentence_transformers.SentenceTransformer(
model_dir, device=self._device, trust_remote_code=True
)
else:
# Load from Hugging Face (default)
self._model = sentence_transformers.SentenceTransformer(
self._model_name, device=self._device, trust_remote_code=True
)
return self._model
except ImportError as e:
if "modelscope" in str(e) and self._model_source == "modelscope":
raise ImportError(
"ModelScope support requires the 'modelscope' package. "
"Please install it with: pip install modelscope"
) from e
raise
except Exception as e:
raise ValueError(
f"Failed to load Sentence Transformer model '{self._model_name}' "
f"from {self._model_source}: {e!s}"
) from e
def _is_sparse_model(self) -> bool:
"""Check if the loaded model is a sparse encoder (e.g., SPLADE).
Returns:
bool: True if model supports sparse encoding.
"""
model = self._get_model()
# Check if model has sparse encoding methods
return hasattr(model, "encode_query") or hasattr(model, "encode_document")
@@ -0,0 +1,396 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import TYPE_CHECKING, Literal, Optional
from ..model.doc import Doc, DocList
from ..tool import require_module
from .rerank_function import RerankFunction
from .sentence_transformer_function import SentenceTransformerFunctionBase
if TYPE_CHECKING:
from ..model.schema import FieldSchema, VectorSchema
class DefaultLocalReRanker(SentenceTransformerFunctionBase, RerankFunction):
"""Re-ranker using Sentence Transformer cross-encoder models for semantic re-ranking.
This re-ranker leverages pre-trained cross-encoder models to perform deep semantic
re-ranking of search results. It runs locally without API calls, supports GPU
acceleration, and works with models from Hugging Face or ModelScope.
Cross-encoder models evaluate query-document pairs jointly, providing more
accurate relevance scores than bi-encoder (embedding-based) similarity.
Args:
query (str): Query text for semantic re-ranking. **Required**.
rerank_field (Optional[str], optional): Document field name to use as
re-ranking input text. **Required** (e.g., "content", "title", "body").
model_name (str, optional): Cross-encoder model identifier or local path.
Defaults to ``"cross-encoder/ms-marco-MiniLM-L6-v2"`` (MS MARCO MiniLM).
Common options:
- ``"cross-encoder/ms-marco-MiniLM-L6-v2"``: Lightweight, fast (~80MB, recommended)
- ``"cross-encoder/ms-marco-MiniLM-L12-v2"``: Better accuracy (~120MB)
- ``"BAAI/bge-reranker-base"``: BGE Reranker Base (~280MB)
- ``"BAAI/bge-reranker-large"``: BGE Reranker Large (highest quality, ~560MB)
model_source (Literal["huggingface", "modelscope"], optional): Model source.
Defaults to ``"huggingface"``.
- ``"huggingface"``: Load from Hugging Face Hub
- ``"modelscope"``: Load from ModelScope (recommended for users in China)
device (Optional[str], optional): Device to run the model on.
Options: ``"cpu"``, ``"cuda"``, ``"mps"`` (for Apple Silicon), or ``None``
for automatic detection. Defaults to ``None``.
batch_size (int, optional): Batch size for processing query-document pairs.
Larger values speed up processing but use more memory. Defaults to ``32``.
Attributes:
query (str): The query text used for re-ranking.
rerank_field (Optional[str]): Field name used for re-ranking input.
model_name (str): The cross-encoder model being used.
model_source (str): The model source ("huggingface" or "modelscope").
device (str): The device the model is running on.
Raises:
ValueError: If ``query`` is empty/None, ``rerank_field`` is None,
or model cannot be loaded.
TypeError: If input types are invalid.
RuntimeError: If model inference fails.
Note:
- Requires Python 3.10, 3.11, or 3.12
- Requires ``sentence-transformers`` package: ``pip install sentence-transformers``
- For ModelScope support, also requires: ``pip install modelscope``
- First run downloads the model (~80-560MB depending on model) from chosen source
- No API keys or network required after initial download
- Cross-encoders are slower than bi-encoders but more accurate
- GPU acceleration provides significant speedup (5-10x)
**MS MARCO MiniLM-L6-v2 Model (Default):**
The default model ``cross-encoder/ms-marco-MiniLM-L6-v2`` is a lightweight and
efficient cross-encoder trained on MS MARCO dataset. It provides:
- Fast inference speed (suitable for real-time applications)
- Small model size (~80MB, quick to download)
- Good balance between speed and accuracy
- Trained on 500K+ query-document pairs
- Public availability without authentication
**For users in China:**
If you encounter Hugging Face access issues, use ModelScope instead:
.. code-block:: python
# Recommended for users in China
reranker = SentenceTransformerReRanker(
query="机器学习算法",
rerank_field="content",
model_source="modelscope"
)
Alternatively, use Hugging Face mirror:
.. code-block:: bash
export HF_ENDPOINT=https://hf-mirror.com
Examples:
>>> # Basic usage with default MS MARCO MiniLM model
>>> from zvec.extension import SentenceTransformerReRanker
>>>
>>> reranker = SentenceTransformerReRanker(
... query="machine learning algorithms",
... rerank_field="content"
... )
>>>
>>> # Use in collection.query()
>>> results = collection.query(
... data={"vector_field": query_vector},
... reranker=reranker,
... topk=20
... )
>>> # Using ModelScope for users in China
>>> reranker = SentenceTransformerReRanker(
... query="深度学习",
... rerank_field="content",
... model_source="modelscope"
... )
>>> # Using larger model for better quality
>>> reranker = SentenceTransformerReRanker(
... query="neural networks",
... rerank_field="content",
... model_name="BAAI/bge-reranker-large",
... device="cuda",
... batch_size=64
... )
>>> # Direct rerank call (for testing)
>>> query_results = {
... "vector1": [
... Doc(id="1", score=0.9, fields={"content": "Machine learning is..."}),
... Doc(id="2", score=0.8, fields={"content": "Deep learning is..."}),
... ]
... }
>>> reranked = reranker.rerank(query_results)
>>> for doc in reranked:
... print(f"ID: {doc.id}, Score: {doc.score:.4f}")
ID: 2, Score: 0.9234
ID: 1, Score: 0.8567
See Also:
- ``RerankFunction``: Abstract base class for re-rankers
- ``QwenReRanker``: Re-ranker using Qwen API
- ``RrfReRanker``: Multi-vector re-ranker using RRF
- ``WeightedReRanker``: Multi-vector re-ranker using weighted scores
References:
- MS MARCO Cross-Encoder: https://huggingface.co/cross-encoder/ms-marco-MiniLM-L6-v2
- BGE Reranker: https://huggingface.co/BAAI/bge-reranker-base
- Cross-Encoder vs Bi-Encoder: https://www.sbert.net/examples/applications/cross-encoder/README.html
"""
def __init__(
self,
query: Optional[str] = None,
rerank_field: Optional[str] = None,
model_name: str = "cross-encoder/ms-marco-MiniLM-L6-v2",
model_source: Literal["huggingface", "modelscope"] = "huggingface",
device: Optional[str] = None,
batch_size: int = 32,
):
"""Initialize SentenceTransformerReRanker with query and configuration.
Args:
query (Optional[str]): Query text for semantic matching. Required.
rerank_field (Optional[str]): Document field for re-ranking input.
model_name (str): Cross-encoder model identifier.
model_source (Literal["huggingface", "modelscope"]): Model source.
device (Optional[str]): Target device ("cpu", "cuda", "mps", or None).
batch_size (int): Batch size for processing query-document pairs.
Raises:
ValueError: If query is empty or model cannot be loaded.
"""
# Initialize base class for model loading
SentenceTransformerFunctionBase.__init__(
self, model_name=model_name, model_source=model_source, device=device
)
# Initialize rerank parameters
self._rerank_field = rerank_field
# Validate query
if not query:
raise ValueError("Query is required for DefaultLocalReRanker")
self._query = query
self._batch_size = batch_size
# Load and validate cross-encoder model
model = self._get_model()
if not hasattr(model, "predict"):
raise ValueError(
f"Model '{model_name}' does not appear to be a cross-encoder model. "
"Cross-encoder models should have a 'predict' method."
)
self._model = model
def _get_model(self):
"""Load or retrieve the CrossEncoder model.
This overrides the base class method to load CrossEncoder instead of
SentenceTransformer, as reranking requires cross-encoder models.
Returns:
CrossEncoder: The loaded cross-encoder model instance.
Raises:
ImportError: If required packages are not installed.
ValueError: If model cannot be loaded.
"""
# Return cached model if exists
if self._model is not None:
return self._model
# Load cross-encoder model
try:
sentence_transformers = require_module("sentence_transformers")
if self._model_source == "modelscope":
# Load from ModelScope
require_module("modelscope")
from modelscope.hub.snapshot_download import snapshot_download
# Download model to cache
model_dir = snapshot_download(self._model_name)
# Load CrossEncoder from local path
model = sentence_transformers.CrossEncoder(
model_dir, device=self._device
)
else:
# Load CrossEncoder from Hugging Face (default)
model = sentence_transformers.CrossEncoder(
self._model_name, device=self._device
)
return model
except ImportError as e:
if "modelscope" in str(e) and self._model_source == "modelscope":
raise ImportError(
"ModelScope support requires the 'modelscope' package. "
"Please install it with: pip install modelscope"
) from e
raise
except Exception as e:
raise ValueError(
f"Failed to load CrossEncoder model '{self._model_name}' "
f"from {self._model_source}: {e!s}"
) from e
@property
def rerank_field(self) -> Optional[str]:
"""Optional[str]: Field name used as re-ranking input."""
return self._rerank_field
@property
def query(self) -> str:
"""str: Query text used for semantic re-ranking."""
return self._query
@property
def batch_size(self) -> int:
"""int: Batch size for processing query-document pairs."""
return self._batch_size
def rerank(
self,
query_results: list[list[Doc]],
topn: int = 10,
*,
fields: list[FieldSchema | VectorSchema] | None = None, # noqa: ARG002
) -> DocList:
"""Re-rank documents using Sentence Transformer cross-encoder model.
Evaluates each query-document pair using the cross-encoder model to compute
relevance scores. Documents are then sorted by these scores and the top-k
results are returned.
Args:
query_results (list[list[Doc]]): Per-sub-query lists of retrieved
documents. Documents from all lists are deduplicated and
re-ranked together.
topn (int): Maximum number of documents to return.
fields: Unused; present for interface compatibility.
Returns:
list[Doc]: Re-ranked documents (up to ``topn``) with updated ``score``
fields containing relevance scores from the cross-encoder model.
Raises:
ValueError: If no valid documents are found or model inference fails.
Note:
- Duplicate documents (same ID) across fields are processed once
- Documents with empty/missing ``rerank_field`` content are skipped
- Returned scores are logits from the cross-encoder model
- Higher scores indicate higher relevance
- Processing time is O(n) where n is the number of documents
Examples:
>>> reranker = SentenceTransformerReRanker(
... query="machine learning",
... topn=3,
... rerank_field="content"
... )
>>> query_results = {
... "vector1": [
... Doc(id="1", score=0.9, fields={"content": "ML basics"}),
... Doc(id="2", score=0.8, fields={"content": "DL tutorial"}),
... ]
... }
>>> reranked = reranker.rerank(query_results)
>>> len(reranked) <= 3
True
"""
if not query_results:
return []
# Accept both dict (legacy) and list formats
if isinstance(query_results, dict):
query_results = list(query_results.values())
# Collect and deduplicate documents
id_to_doc: dict[str, Doc] = {}
doc_ids: list[str] = []
contents: list[str] = []
for query_result in query_results:
for doc in query_result:
doc_id = doc.id
if doc_id in id_to_doc:
continue
# Extract text content from specified field
field_value = doc.field(self.rerank_field)
rank_content = str(field_value).strip() if field_value else ""
if not rank_content:
continue
id_to_doc[doc_id] = doc
doc_ids.append(doc_id)
contents.append(rank_content)
if not contents:
raise ValueError("No documents to rerank")
try:
# Use standard cross-encoder predict method
pairs = [[self.query, content] for content in contents]
scores = self._model.predict(
pairs,
batch_size=self.batch_size,
show_progress_bar=False,
convert_to_numpy=True,
)
# Convert to float list if needed
if hasattr(scores, "tolist"):
scores = scores.tolist()
else:
scores = [float(s) for s in scores]
except Exception as e:
raise RuntimeError(f"Failed to compute rerank scores: {e!s}") from e
# Create scored documents
scored_docs = [
(doc_ids[i], id_to_doc[doc_ids[i]], scores[i]) for i in range(len(doc_ids))
]
# Sort by score (descending) and take top-k
scored_docs.sort(key=lambda x: x[2], reverse=True)
top_scored_docs = scored_docs[:topn]
# Build result list with updated scores
results: DocList = []
for _, doc, score in top_scored_docs:
new_doc = doc._replace(score=score)
results.append(new_doc)
return results
+30
View File
@@ -0,0 +1,30 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from .collection import Collection
from .doc import Doc
from .param.query import Fts, Query, VectorQuery
from .schema.collection_schema import CollectionSchema
from .schema.field_schema import FieldSchema
__all__ = [
"Collection",
"CollectionSchema",
"Doc",
"FieldSchema",
"Fts",
"Query",
"VectorQuery",
]
+439
View File
@@ -0,0 +1,439 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import warnings
from typing import Optional, Union, overload
from zvec._zvec import _Collection
from ..executor import QueryContext, QueryExecutor
from ..extension import ReRanker
from ..typing import Status
from .convert import convert_to_cpp_doc, convert_to_py_doc
from .doc import Doc, DocList
from .param import (
AddColumnOption,
AlterColumnOption,
CollectionOption,
FlatIndexParam,
FtsIndexParam,
HnswIndexParam,
HnswRabitqIndexParam,
IndexOption,
InvertIndexParam,
IVFIndexParam,
OptimizeOption,
)
from .param.query import Query
from .schema import CollectionSchema, CollectionStats, FieldSchema
__all__ = ["Collection"]
class Collection:
"""Represents an opened collection in Zvec.
A `Collection` provides methods for data definition (DDL), data manipulation (DML),
and querying (DQL). It is obtained via `create_and_open()` or `open()`.
This class is not meant to be instantiated directly; use factory functions instead.
"""
def __init__(self, obj: _Collection):
self._obj = obj
self._schema = None
self._querier = None
@classmethod
def _from_core(cls, core_collection: _Collection) -> Collection:
if not core_collection:
raise ValueError("Collection is None")
inst = cls.__new__(cls)
inst._obj = core_collection
schema = CollectionSchema._from_core(core_collection.Schema())
inst._schema = schema
inst._querier = QueryExecutor(schema)
return inst
@property
def path(self) -> str:
"""str: The filesystem path of the collection."""
return self._obj.Path()
@property
def option(self) -> CollectionOption:
"""CollectionOption: The options used to open the collection."""
return self._obj.Options()
@property
def schema(self) -> CollectionSchema:
"""CollectionSchema: The schema defining the structure of the collection."""
return self._schema
@property
def stats(self) -> CollectionStats:
"""CollectionStats: Runtime statistics about the collection (e.g., doc count, size)."""
return self._obj.Stats()
# ========== Collection DDL Methods ==========
def destroy(self) -> None:
"""Permanently delete the collection from disk.
Warning:
This operation is irreversible. All data will be lost.
"""
self._obj.Destroy()
def flush(self) -> None:
"""Force all pending writes to disk.
Ensures durability of recent inserts/updates.
"""
self._obj.Flush()
# ========== Index DDL Methods ==========
def create_index(
self,
field_name: str,
index_param: Union[
HnswIndexParam,
HnswRabitqIndexParam,
IVFIndexParam,
FlatIndexParam,
InvertIndexParam,
FtsIndexParam,
],
option: IndexOption = IndexOption(),
) -> None:
"""Create an index on a field.
Vector index types (HNSW, IVF, FLAT) can only be applied to vector fields.
Inverted index (`InvertIndexParam`) is for scalar fields.
FTS index (`FtsIndexParam`) is for full-text search on STRING fields.
Args:
field_name (str): Name of the field to index.
index_param (Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam, FtsIndexParam]):
Index configuration.
option (Optional[IndexOption], optional): Index creation options.
Defaults to ``IndexOption()``.
"""
self._obj.CreateIndex(field_name, index_param, option)
self._schema = CollectionSchema._from_core(self._obj.Schema())
self._querier._schema = self._schema
def drop_index(self, field_name: str) -> None:
"""Remove the index from a field.
Args:
field_name (str): Name of the indexed field.
"""
self._obj.DropIndex(field_name)
self._schema = CollectionSchema._from_core(self._obj.Schema())
self._querier._schema = self._schema
def optimize(self, option: OptimizeOption = OptimizeOption()) -> None:
"""Optimize the collection (e.g., merge segments, rebuild index).
Args:
option (Optional[OptimizeOption], optional): Optimization options.
Defaults to ``OptimizeOption()``.
"""
self._obj.Optimize(option)
# ========== COLUMN DDL Methods ==========
def add_column(
self,
field_schema: FieldSchema,
expression: str = "",
option: AddColumnOption = AddColumnOption(),
) -> None:
"""Add a new column to the collection.
The column is populated using the provided expression (e.g., SQL-like formula).
Args:
field_schema (FieldSchema): Schema definition for the new column.
expression (str): Expression to compute values for existing documents.
option (Optional[AddColumnOption], optional): Options for the operation.
Defaults to ``AddColumnOption()``.
"""
self._obj.AddColumn(field_schema._get_object(), expression, option)
self._schema = CollectionSchema._from_core(self._obj.Schema())
self._querier._schema = self._schema
def drop_column(self, field_name: str) -> None:
"""Remove a column from the collection.
Args:
field_name (str): Name of the column to drop.
"""
self._obj.DropColumn(field_name)
self._schema = CollectionSchema._from_core(self._obj.Schema())
self._querier._schema = self._schema
def alter_column(
self,
old_name: str,
new_name: Optional[str] = None,
field_schema: Optional[FieldSchema] = None,
option: AlterColumnOption = AlterColumnOption(),
) -> None:
"""Rename a column, update its schema.
This method supports three atomic operations:
1. Rename only (when `field_schema` is None).
2. Modify schema only (when `new_name` is None or empty string).
Args:
old_name (str): The current name of the column to be altered.
new_name (Optional[str]): The new name for the column.
- If provided and non-empty, the column will be renamed.
- If `None` or empty string, no rename occurs.
field_schema (Optional[FieldSchema]): The new schema definition.
- If provided, the column's type, dimension, or other properties will be updated.
- If `None`, only renaming (if requested) is performed.
option (AlterColumnOption, optional): Options controlling the alteration behavior.
Defaults to ``AlterColumnOption()``.
**Limitation**: This operation **only supports scalar numeric columns**. such as:
- `DOUBLE`, `FLOAT`,
- `INT32`, `INT64`, `UINT32`, `UINT64`
Note:
- Schema modification may trigger data migration or index rebuild.
Examples:
>>> # Rename column only
>>> results = collection.alter_column(old_name="id", new_name="doc_id")
>>> # Modify schema only
>>> new_schema = FieldSchema(name="doc_id", dtype=DataType.INT64)
>>> collection.alter_column("id", field_schema=new_schema)
"""
self._obj.AlterColumn(
old_name,
new_name or "",
field_schema._get_object() if field_schema else None,
option,
)
self._schema = CollectionSchema._from_core(self._obj.Schema())
self._querier._schema = self._schema
# ========== Collection DDL Methods ==========
@overload
def insert(self, docs: Doc) -> Status:
pass
@overload
def insert(self, docs: list[Doc]) -> list[Status]:
pass
def insert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:
"""Insert new documents into the collection.
Documents must have unique IDs and conform to the schema.
Args:
docs (Union[Doc, list[Doc]]): One or more documents to insert.
Returns:
Union[Status, list[Status]]: If a single Doc was given, returns its Status;
if a list was given, returns a list of Status objects.
"""
is_single = isinstance(docs, Doc)
doc_list = [docs] if is_single else docs
results = self._obj.Insert(
[convert_to_cpp_doc(doc, self.schema) for doc in doc_list]
)
return results[0] if is_single else results
@overload
def upsert(self, docs: Doc) -> Status:
pass
@overload
def upsert(self, docs: list[Doc]) -> list[Status]:
pass
def upsert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:
"""Insert new documents or update existing ones by ID.
Args:
docs (Union[Doc, list[Doc]]): Documents to upsert.
Returns:
Union[Status, list[Status]]: If a single Doc was given, returns its Status;
if a list was given, returns a list of Status objects.
"""
is_single = isinstance(docs, Doc)
doc_list = [docs] if is_single else docs
results = self._obj.Upsert(
[convert_to_cpp_doc(doc, self.schema) for doc in doc_list]
)
return results[0] if is_single else results
@overload
def update(self, docs: Doc) -> Status:
pass
@overload
def update(self, docs: list[Doc]) -> list[Status]:
pass
def update(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:
"""Update existing documents by ID.
Only specified fields are updated; others remain unchanged.
Args:
docs (Union[Doc, list[Doc]]): Documents containing updated fields.
Returns:
Union[Status, list[Status]]: If a single Doc was given, returns its Status;
if a list was given, returns a list of Status objects.
"""
is_single = isinstance(docs, Doc)
doc_list = [docs] if is_single else docs
results = self._obj.Update(
[convert_to_cpp_doc(doc, self.schema) for doc in doc_list]
)
return results[0] if is_single else results
@overload
def delete(self, ids: str) -> Status:
pass
@overload
def delete(self, ids: list[str]) -> list[Status]:
pass
def delete(self, ids: Union[str, list[str]]) -> Union[Status, list[Status]]:
"""Delete documents by ID.
Args:
ids (Union[str, list[str]]): One or more document IDs to delete.
Returns:
Union[Status, list[Status]]: If a single id was given, returns its Status;
if a list was given, returns a list of Status objects.
"""
is_single = isinstance(ids, str)
id_list = [ids] if isinstance(ids, str) else ids
results = self._obj.Delete(id_list)
return results[0] if is_single else results
def delete_by_filter(self, filter: str) -> None:
"""Delete documents matching a filter expression.
Args:
filter (str): Boolean expression (e.g., ``"age > 30"``).
"""
self._obj.DeleteByFilter(filter)
# ========== Collection DQL-fetch Methods ==========
def fetch(
self,
ids: Union[str, list[str]],
*,
output_fields: Optional[list[str]] = None,
include_vector: bool = True,
) -> dict[str, Doc]:
"""Retrieve documents by ID.
Args:
ids (Union[str, list[str]]): Document IDs to fetch.
output_fields (Optional[list[str]], optional): Scalar fields to
include. If None, all fields are returned. Defaults to None.
include_vector (bool, optional): Whether to include vector data in
results. Defaults to True.
Returns:
dict[str, Doc]: Mapping from ID to document. Missing IDs are omitted.
"""
ids = [ids] if isinstance(ids, str) else ids
docs = self._obj.Fetch(ids, output_fields, include_vector)
return {
doc_id: py_doc
for doc_id, core_doc in docs.items()
if (py_doc := convert_to_py_doc(core_doc, self.schema)) is not None
}
# ========== Collection DQL-Query Methods ==========
def query(
self,
queries: Optional[Union[Query, list[Query]]] = None,
*,
vectors: Optional[Union[Query, list[Query]]] = None,
topk: int = 10,
filter: Optional[str] = None,
include_vector: bool = False,
output_fields: Optional[list[str]] = None,
reranker: Optional[ReRanker] = None,
) -> DocList:
"""Perform vector similarity search with optional filtering and re-ranking.
At least one `Query` must be provided via `queries`.
Args:
queries (Optional[Union[Query, list[Query]]], optional):
One or more vector queries. Defaults to None.
vectors (Optional[Union[Query, list[Query]]], optional):
Deprecated. Use `queries` instead.
topk (int, optional): Number of nearest neighbors to return.
Defaults to 10.
filter (Optional[str], optional): Boolean expression to pre-filter candidates.
Defaults to None.
include_vector (bool, optional): Whether to include vector data in results.
Defaults to False.
output_fields (Optional[list[str]], optional): Scalar fields to include.
If None, all fields are returned. Defaults to None.
reranker (Optional[ReRanker], optional): Re-ranker to refine results.
Defaults to None.
Returns:
DocList: Top-k matching documents, sorted by relevance score.
Examples:
>>> from zvec import Query
>>> results = collection.query(
... queries=Query(field_name="embedding", vector=[0.1, 0.2]),
... topk=5,
... filter="category == 'tech'",
... output_fields=["title", "url"]
... )
"""
if vectors is not None:
warnings.warn(
"The 'vectors' parameter is deprecated and will be removed in a future version. "
"Use 'queries' instead.",
DeprecationWarning,
stacklevel=2,
)
if queries is not None:
raise ValueError("Cannot specify both 'queries' and 'vectors'.")
queries = vectors
ctx = QueryContext(
topk=topk,
filter=filter,
queries=[queries] if isinstance(queries, Query) else queries,
include_vector=include_vector,
output_fields=output_fields,
reranker=reranker,
)
return self._querier.execute(ctx, self._obj)
+54
View File
@@ -0,0 +1,54 @@
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from zvec._zvec import _Doc
from .doc import Doc
from .schema import CollectionSchema
def convert_to_cpp_doc(doc: Doc, collection_schema: CollectionSchema) -> _Doc:
if not doc or not collection_schema:
return None
_doc = _Doc()
# set pk
_doc.set_pk(doc.id)
# set scalar fields
for k, v in doc.fields.items():
field_schema = collection_schema.field(k)
if not field_schema:
raise ValueError(
f"schema validate failed: {k} not found in collection schema"
)
_doc.set_any(k, field_schema._get_object(), v)
# set vector fields
for k, v in doc.vectors.items():
vector_schema = collection_schema.vector(k)
if not vector_schema:
raise ValueError(
f"schema validate failed: {k} not found in collection schema"
)
_doc.set_any(k, vector_schema._get_object(), v)
return _doc
def convert_to_py_doc(doc: _Doc, collection_schema: CollectionSchema) -> Doc:
if not doc or not collection_schema:
return None
data_tuple = doc.get_all(collection_schema._get_object())
return Doc._from_tuple(data_tuple)
+178
View File
@@ -0,0 +1,178 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import json
from typing import Any, Optional
from ..common import VectorType
__all__ = [
"Doc",
"DocList",
]
class Doc:
"""Represents a retrieved document with optional metadata, fields, and vectors.
This immutable data class encapsulates the result of a search or retrieval
operation. It includes the document ID, relevance score (if applicable),
scalar fields, and vector embeddings.
During initialization, any `numpy.ndarray` in `vectors` is automatically
converted to a plain Python list for JSON serialization and immutability.
Attributes:
id (str): Unique identifier of the document.
score (Optional[float], optional): Relevance score from search.
Defaults to None.
vectors (Optional[dict[str, VectorType]], optional): Named vector
embeddings associated with the document. Values are converted to
lists if originally `np.ndarray`. Defaults to None.
fields (Optional[dict[str, Any]], optional): Scalar metadata fields
(e.g., title, timestamp). Defaults to None.
Examples:
>>> import numpy as np
>>> import zvec
>>> doc = zvec.Doc(
... id="doc1",
... score=0.95,
... vectors={"emb": np.array([0.1, 0.2, 0.3])},
... fields={"title": "Hello World"}
... )
>>> print(doc.vector("emb"))
[0.1, 0.2, 0.3]
>>> print(doc.has_field("title"))
True
"""
__slots__ = ("id", "score", "vectors", "fields")
def __init__(
self,
id: str,
score: Optional[float] = None,
vectors: Optional[dict[str, VectorType]] = None,
fields: Optional[dict[str, Any]] = None,
):
self.id = id
self.score = score
self.vectors = vectors or {}
self.fields = fields or {}
def has_field(self, name: str) -> bool:
"""Check if the document contains a scalar field with the given name.
Args:
name (str): Name of the field to check.
Returns:
bool: True if the field exists, False otherwise.
"""
return name in self.fields
def has_vector(self, name: str) -> bool:
"""Check if the document contains a vector with the given name.
Args:
name (str): Name of the vector to check.
Returns:
bool: True if the vector exists, False otherwise.
"""
return name in self.vectors
def vector(self, name: str):
"""Get a vector by name.
Args:
name (str): Name of the vector.
Returns:
Any: The vector (as a list) if it exists, otherwise None.
"""
return self.vectors and self.vectors.get(name)
def field(self, name: str):
"""Get a scalar field by name.
Args:
name (str): Name of the field.
Returns:
Any: The field value if it exists, otherwise None.
"""
return self.fields and self.fields.get(name)
def vector_names(self) -> list[str]:
"""Get the list of all vector names in this document.
Returns:
list[str]: A list of vector field names. Empty if no vectors.
"""
return [] if not self.vectors else list(self.vectors.keys())
def field_names(self) -> list[str]:
"""Get the list of all scalar field names in this document.
Returns:
list[str]: A list of field names. Empty if no fields.
"""
return [] if not self.fields else list(self.fields.keys())
def __repr__(self) -> str:
try:
schema = {
"id": self.id,
"score": self.score,
"fields": self.fields,
"vectors": self.vectors,
}
return json.dumps(schema, indent=2, ensure_ascii=False)
except Exception as e:
return f"<Doc error during repr: {e}>"
def _replace(self, **changes):
new_tuple = (
changes.get("id", self.id),
changes.get("score", self.score),
changes.get("fields", self.fields.copy() if self.fields else None),
changes.get("vectors", self.vectors.copy() if self.vectors else None),
)
return type(self)._from_tuple(new_tuple)
@classmethod
def _from_tuple(
cls, data_tuple: tuple[str, float, dict[str, Any], dict[str, VectorType]]
):
obj = object.__new__(cls)
obj.id = data_tuple[0]
obj.score = data_tuple[1]
obj.fields = data_tuple[2] or {}
vectors = data_tuple[3]
if vectors is not None:
obj.vectors = {
name: (vec.tolist() if hasattr(vec, "tolist") else vec)
for name, vec in vectors.items()
}
else:
obj.vectors = {}
return obj
#: Type alias for query results: a list of documents returned by a single query route.
DocList = list[Doc]
+60
View File
@@ -0,0 +1,60 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from zvec._zvec.param import (
AddColumnOption,
AlterColumnOption,
CollectionOption,
DiskAnnIndexParam,
DiskAnnQueryParam,
FlatIndexParam,
FtsIndexParam,
FtsQueryParam,
HnswIndexParam,
HnswQueryParam,
HnswRabitqIndexParam,
HnswRabitqQueryParam,
IndexOption,
InvertIndexParam,
IVFIndexParam,
IVFQueryParam,
OptimizeOption,
QuantizerParam,
VamanaIndexParam,
VamanaQueryParam,
)
__all__ = [
"AddColumnOption",
"AlterColumnOption",
"CollectionOption",
"DiskAnnIndexParam",
"DiskAnnQueryParam",
"FlatIndexParam",
"FtsIndexParam",
"FtsQueryParam",
"HnswIndexParam",
"HnswQueryParam",
"HnswRabitqIndexParam",
"HnswRabitqQueryParam",
"IVFIndexParam",
"IVFQueryParam",
"IndexOption",
"InvertIndexParam",
"OptimizeOption",
"QuantizerParam",
"VamanaIndexParam",
"VamanaQueryParam",
]
File diff suppressed because it is too large Load Diff
+143
View File
@@ -0,0 +1,143 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import warnings
from dataclasses import dataclass
from typing import Optional, Union
from ...common import VectorType
from . import FtsQueryParam, HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam
__all__ = ["Fts", "Query", "VectorQuery"]
@dataclass(frozen=True)
class Fts:
"""Full-text search query parameters.
Attributes:
query_string (Optional[str]): FTS query expression
(e.g. '+vector -slow "exact phrase"'). Mutually exclusive with match_string.
match_string (Optional[str]): Natural language match string,
tokenized and combined using the default operator.
Mutually exclusive with query_string.
"""
query_string: Optional[str] = None
match_string: Optional[str] = None
@dataclass(frozen=True)
class Query:
"""Represents a search query for a specific field in a collection.
A `Query` can be constructed for either vector search or full-text search,
but not both simultaneously.
For vector search, provide `id` or `vector` (and optionally `param`).
For FTS, provide `fts`.
Attributes:
field_name (str): Name of the field to query.
id (Optional[str], optional): Document ID to fetch vector from. Default is None.
vector (VectorType, optional): Explicit query vector. Default is None.
param (Optional[Union[HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam, FtsQueryParam]], optional):
Index-specific query parameters. Default is None.
fts (Optional[Fts], optional): Full-text search parameters. Default is None.
Examples:
>>> import zvec
>>> # Query by ID
>>> q1 = zvec.Query(field_name="embedding", id="doc123")
>>> # Query by vector
>>> q2 = zvec.Query(
... field_name="embedding",
... vector=[0.1, 0.2, 0.3],
... param=HnswQueryParam(ef=300)
... )
>>> # FTS query
>>> q3 = zvec.Query(
... field_name="content",
... fts=Fts(match_string="machine learning")
... )
>>> # FTS query with custom operator
>>> q4 = zvec.Query(
... field_name="content",
... fts=Fts(match_string="machine learning"),
... param=FtsQueryParam(default_operator="AND")
... )
"""
field_name: str
id: Optional[str] = None
vector: VectorType = None
param: Optional[
Union[HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam, FtsQueryParam]
] = None
fts: Optional[Fts] = None
def has_id(self) -> bool:
"""Check if the query is based on a document ID.
Returns:
bool: True if `id` is set, False otherwise.
"""
return self.id is not None
def has_vector(self) -> bool:
"""Check if the query contains an explicit vector.
Returns:
bool: True if `vector` is non-empty, False otherwise.
"""
return self.vector is not None and len(self.vector) > 0
def has_fts(self) -> bool:
"""Check if the query contains an FTS (full-text search) condition.
Returns:
bool: True if `fts` is set with a query_string or match_string.
"""
if self.fts is not None:
return bool(self.fts.query_string) or bool(self.fts.match_string)
return False
def _validate(self) -> None:
if self.field_name is None:
raise ValueError("Field name cannot be empty")
if self.has_id() and self.has_vector():
raise ValueError("Cannot provide both id and vector")
if self.has_fts() and (self.has_vector() or self.has_id()):
raise ValueError(
"Cannot combine fts with vector search fields (id/vector) in a single Query"
)
if self.fts is not None and self.fts.query_string and self.fts.match_string:
raise ValueError(
"Cannot provide both query_string and match_string in Fts; "
"they are mutually exclusive"
)
class VectorQuery(Query):
"""Deprecated alias for Query. Use Query instead."""
def __new__(cls, *args, **kwargs): # noqa : ARG004
warnings.warn(
"VectorQuery is deprecated and will be removed in a future version. "
"Use Query instead.",
DeprecationWarning,
stacklevel=2,
)
return super().__new__(cls)
+21
View File
@@ -0,0 +1,21 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from zvec._zvec.schema import CollectionStats
from .collection_schema import CollectionSchema
from .field_schema import FieldSchema, VectorSchema
__all__ = ["CollectionSchema", "CollectionStats", "FieldSchema", "VectorSchema"]
+109
View File
@@ -0,0 +1,109 @@
"""
This module contains the schema of Zvec
"""
from __future__ import annotations
import collections.abc
import typing
import zvec._zvec.param
import zvec._zvec.typing
from .collection_schema import CollectionSchema
from .field_schema import FieldSchema, VectorSchema
__all__: list[str] = [
"CollectionSchema",
"CollectionStats",
"FieldSchema",
"VectorSchema",
]
class CollectionStats:
def __init__(self) -> None: ...
def __repr__(self) -> str: ...
@property
def doc_count(self) -> int: ...
@property
def index_completeness(self) -> dict[str, float]: ...
class _CollectionSchema:
__hash__: typing.ClassVar[None] = None
def __eq__(self, arg0: _CollectionSchema) -> bool: ...
def __init__(
self, name: str, fields: collections.abc.Sequence[_FieldSchema]
) -> None:
"""
Construct with name and list of fields
"""
def __ne__(self, arg0: _CollectionSchema) -> bool: ...
def fields(self) -> list[_FieldSchema]:
"""
Return list of all field schemas.
"""
def forward_fields(self) -> list[_FieldSchema]:
"""
Return list of forward-indexed fields.
"""
def get_field(self, field_name: str) -> _FieldSchema:
"""
Get field by name (const pointer), returns None if not found.
"""
def get_forward_field(self, field_name: str) -> _FieldSchema:
"""
Get forward field (used for filtering).
"""
def get_vector_field(self, field_name: str) -> _FieldSchema:
"""
Get vector field by name.
"""
def has_field(self, field_name: str) -> bool:
"""
Check if a field exists.
"""
def vector_fields(self) -> list[_FieldSchema]:
"""
Return list of vector fields.
"""
@property
def name(self) -> str: ...
class _FieldSchema:
__hash__: typing.ClassVar[None] = None
def __eq__(self, arg0: _FieldSchema) -> bool: ...
def __init__(
self,
name: str,
data_type: zvec._zvec.typing.DataType,
nullable: bool = False,
dimension: typing.SupportsInt = 0,
index_param: zvec._zvec.param.IndexParam = None,
) -> None: ...
def __ne__(self, arg0: _FieldSchema) -> bool: ...
@property
def data_type(self) -> zvec._zvec.typing.DataType: ...
@property
def dimension(self) -> int: ...
@property
def index_param(self) -> typing.Any: ...
@property
def index_type(self) -> zvec._zvec.typing.IndexType: ...
@property
def is_dense_vector(self) -> bool: ...
@property
def is_sparse_vector(self) -> bool: ...
@property
def name(self) -> str: ...
@property
def nullable(self) -> bool: ...
@@ -0,0 +1,215 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import json
from typing import Optional, Union
from zvec._zvec.schema import _CollectionSchema, _FieldSchema
from .field_schema import FieldSchema, VectorSchema
__all__ = [
"CollectionSchema",
]
class CollectionSchema:
"""Defines the structure of a collection in Zvec.
A collection schema specifies the name of the collection and its fields,
including both scalar fields (e.g., int, string) and vector fields.
Field names must be unique across both scalar and vector fields.
Args:
name (str): Name of the collection.
fields (Optional[Union[FieldSchema, list[FieldSchema]]], optional):
One or more scalar field definitions. Defaults to None.
vectors (Optional[Union[VectorSchema, list[VectorSchema]]], optional):
One or more vector field definitions. Defaults to None.
Raises:
TypeError: If `fields` or `vectors` are of unsupported types.
ValueError: If any field or vector name is duplicated.
Examples:
>>> from zvec import FieldSchema, VectorSchema, DataType, IndexType
>>> id_field = FieldSchema("id", DataType.INT64, is_primary=True)
>>> emb_field = VectorSchema("embedding", dim=128, data_type=DataType.VECTOR_FP32)
>>> schema = CollectionSchema(
... name="my_collection",
... fields=id_field,
... vectors=emb_field
... )
>>> print(schema.name)
my_collection
"""
def __init__(
self,
name: str,
fields: Optional[Union[FieldSchema, list[FieldSchema]]] = None,
vectors: Optional[Union[VectorSchema, list[VectorSchema]]] = None,
):
if name is None or not isinstance(name, str):
raise ValueError(
f"schema validate failed: collection name must be str, got {type(name).__name__}"
)
# handle fields
_fields_name: list[str] = []
_fields_list: list[_FieldSchema] = []
self._check_fields(fields, _fields_name, _fields_list)
self._check_vectors(vectors, _fields_name, _fields_list)
# init
self._cpp_obj = _CollectionSchema(
name=name,
fields=_fields_list,
)
def _check_fields(
self,
fields: Optional[Union[FieldSchema, list[FieldSchema]]],
_fields_name: list[str],
_fields_list: list[_FieldSchema],
) -> None:
field_items = []
if isinstance(fields, FieldSchema):
field_items = [fields]
elif isinstance(fields, list):
field_items = fields
elif fields is None:
field_items = []
else:
raise TypeError(
f"schema validate failed: invalid 'fields' type, expected FieldSchema or list[FieldSchema], "
f"got {type(fields).__name__}"
)
for idx, field in enumerate(field_items):
if not isinstance(field, FieldSchema):
raise TypeError(
f"schema validate failed: invalid field type in 'fields' list, expected FieldSchema, "
f"got {type(field).__name__} at index {idx}"
)
if field.name in _fields_name:
raise ValueError(
f"schema validate failed: duplicate field name '{field.name}': field names must be unique"
)
_fields_name.append(field.name)
_fields_list.append(field._get_object())
def _check_vectors(
self,
vectors: Optional[Union[VectorSchema, list[VectorSchema]]],
_fields_name: list[str],
_fields_list: list[_FieldSchema],
) -> None:
# handle vector
if isinstance(vectors, VectorSchema):
vectors_items = [vectors]
elif isinstance(vectors, list):
vectors_items = vectors
elif vectors is None:
vectors_items = []
else:
raise TypeError(
f"schema validate failed: invalid 'vectors' type, expected VectorSchema or list[VectorSchema], "
f"got {type(vectors).__name__}"
)
for idx, vector in enumerate(vectors_items):
if not isinstance(vector, VectorSchema):
raise TypeError(
f"schema validate failed: invalid vector type in 'vectors' list, expected VectorSchema, "
f"got {type(vector).__name__} at index {idx}"
)
if vector.name in _fields_name:
raise ValueError(
f"schema validate failed: duplicate vector name '{vector.name}', vector names must be unique "
f"(conflicts with existing field or vector)"
)
_fields_name.append(vector.name)
_fields_list.append(vector._get_object())
@classmethod
def _from_core(cls, core_collection_schema: _CollectionSchema):
inst = cls.__new__(cls)
if not core_collection_schema:
raise ValueError("schema validate failed: schema is null")
inst._cpp_obj = core_collection_schema
return inst
@property
def name(self) -> str:
"""str: The name of the collection."""
return self._cpp_obj.name
def field(self, name: str) -> Optional[FieldSchema]:
"""Retrieve a scalar field by name.
Args:
name (str): Name of the field.
Returns:
Optional[FieldSchema]: The field if found, otherwise None.
"""
_field = self._cpp_obj.get_forward_field(name)
return FieldSchema._from_core(_field) if _field else None
def vector(self, name: str) -> Optional[VectorSchema]:
"""Retrieve a vector field by name.
Args:
name (str): Name of the vector field.
Returns:
Optional[VectorSchema]: The vector field if found, otherwise None.
"""
_field = self._cpp_obj.get_vector_field(name)
return VectorSchema._from_core(_field) if _field else None
@property
def fields(self) -> list[FieldSchema]:
"""list[FieldSchema]: All scalar (non-vector) fields in the schema."""
_fields = self._cpp_obj.forward_fields()
return [FieldSchema._from_core(_field) for _field in _fields]
@property
def vectors(self) -> list[VectorSchema]:
"""list[VectorSchema]: All vector fields in the schema."""
_vectors = self._cpp_obj.vector_fields()
return [VectorSchema._from_core(_vector) for _vector in _vectors]
def _get_object(self) -> _CollectionSchema:
return self._cpp_obj
def __repr__(self) -> str:
try:
schema = {
"name": self.name,
"fields": {field.name: field.__dict__() for field in self.fields},
"vectors": {vector.name: vector.__dict__() for vector in self.vectors},
}
return json.dumps(schema, indent=2, ensure_ascii=False)
except Exception as e:
return f"<CollectionSchema error during repr: {e}>"
def __str__(self) -> str:
return self.__repr__()
+310
View File
@@ -0,0 +1,310 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import json
from typing import Any, Optional, Union
from zvec._zvec.schema import _FieldSchema
from zvec.model.param import (
FlatIndexParam,
FtsIndexParam,
HnswIndexParam,
HnswRabitqIndexParam,
InvertIndexParam,
IVFIndexParam,
)
from zvec.typing import DataType
__all__ = [
"FieldSchema",
"VectorSchema",
]
SUPPORT_VECTOR_DATA_TYPE = [
DataType.VECTOR_FP16,
DataType.VECTOR_FP32,
DataType.VECTOR_FP64,
DataType.VECTOR_INT8,
DataType.SPARSE_VECTOR_FP16,
DataType.SPARSE_VECTOR_FP32,
]
SUPPORT_SCALAR_DATA_TYPE = [
DataType.INT32,
DataType.INT64,
DataType.UINT32,
DataType.UINT64,
DataType.FLOAT,
DataType.DOUBLE,
DataType.STRING,
DataType.BOOL,
DataType.ARRAY_INT32,
DataType.ARRAY_INT64,
DataType.ARRAY_UINT32,
DataType.ARRAY_UINT64,
DataType.ARRAY_FLOAT,
DataType.ARRAY_DOUBLE,
DataType.ARRAY_STRING,
DataType.ARRAY_BOOL,
]
class FieldSchema:
"""Represents a scalar (non-vector) field in a collection schema.
A `FieldSchema` defines the name, data type, nullability, and optional
inverted index configuration for a regular field (e.g., ID, timestamp, category).
Args:
name (str): Name of the field. Must be unique within the collection.
data_type (DataType): Data type of the field (e.g., INT64, STRING).
nullable (bool, optional): Whether the field can contain null values.
Defaults to False.
index_param (Optional[Union[InvertIndexParam, FtsIndexParam]], optional):
Index parameters for this field. Use ``InvertIndexParam`` for scalar
inverted indexing, or ``FtsIndexParam`` for full-text search indexing
on STRING fields. Defaults to None.
Examples:
>>> from zvec.typing import DataType
>>> from zvec.model.param import InvertIndexParam, FtsIndexParam
>>> id_field = FieldSchema(
... name="id",
... data_type=DataType.INT64,
... nullable=False,
... index_param=InvertIndexParam(enable_range_optimization=True)
... )
>>> content_field = FieldSchema(
... name="content",
... data_type=DataType.STRING,
... nullable=False,
... index_param=FtsIndexParam(tokenizer_name="standard")
... )
"""
def __init__(
self,
name: str,
data_type: DataType,
nullable: bool = False,
index_param: Optional[Union[InvertIndexParam, FtsIndexParam]] = None,
):
if name is None or not isinstance(name, str):
raise ValueError(
f"schema validate failed: field name must be str, got {type(name).__name__}"
)
if data_type not in SUPPORT_SCALAR_DATA_TYPE:
raise ValueError(
f"schema validate failed: scalar_field's data_type must be one of "
f"{', '.join(str(dt) for dt in SUPPORT_SCALAR_DATA_TYPE)}, "
f"but field[{name}]'s data_type is {data_type}"
)
self._cpp_obj = _FieldSchema(
name=name,
data_type=data_type,
dimension=0,
nullable=nullable,
index_param=index_param,
)
@classmethod
def _from_core(cls, core_field_schema: _FieldSchema):
if core_field_schema is None:
raise ValueError("schema validate failed: field schema is None")
inst = cls.__new__(cls)
inst._cpp_obj = core_field_schema
return inst
def _get_object(self) -> _FieldSchema:
return self._cpp_obj
@property
def name(self) -> str:
"""str: The name of the field."""
return self._cpp_obj.name
@property
def data_type(self) -> DataType:
"""DataType: The data type of the field (e.g., INT64, STRING)."""
return self._cpp_obj.data_type
@property
def nullable(self) -> bool:
"""bool: Whether the field allows null values."""
return self._cpp_obj.nullable
@property
def index_param(self) -> Optional[Union[InvertIndexParam, FtsIndexParam]]:
"""Optional[Union[InvertIndexParam, FtsIndexParam]]: Index configuration, if any."""
return self._cpp_obj.index_param
def __dict__(self) -> dict[str, Any]:
return {
"name": self.name,
"data_type": (
self.data_type.name
if hasattr(self.data_type, "name")
else str(self.data_type)
),
"nullable": self.nullable,
"index_param": (
self.index_param.to_dict() if self.index_param is not None else None
),
}
def __repr__(self) -> str:
try:
schema = self.__dict__()
return json.dumps(schema, indent=2, ensure_ascii=False)
except Exception as e:
return f"<FieldSchema error during repr: {e}>"
def __str__(self) -> str:
return self.__repr__()
def __eq__(self, other: object) -> bool:
if not isinstance(other, FieldSchema):
return False
return self._cpp_obj == other._cpp_obj
def __hash__(self) -> int:
return hash((self.name, self.data_type, self.nullable))
class VectorSchema:
"""Represents a vector field in a collection schema.
A `VectorSchema` defines the name, data type, dimensionality, and index
configuration for a vector field used in similarity search.
Args:
name (str): Name of the vector field. Must be unique within the collection.
data_type (DataType): Vector data type (e.g., VECTOR_FP32, VECTOR_INT8).
dimension (int, optional): Dimensionality of the vector. Must be > 0 for dense vectors;
may be `None` for sparse vectors.
index_param (Union[HnswIndexParam, IVFIndexParam, FlatIndexParam], optional):
Index configuration for this vector field. Defaults to
``HnswIndexParam()``.
Examples:
>>> from zvec.typing import DataType
>>> from zvec.model.param import HnswIndexParam
>>> emb_field = VectorSchema(
... name="embedding",
... data_type=DataType.VECTOR_FP32,
... dimension=128,
... index_param=HnswIndexParam(ef_construction=200, m=16)
... )
"""
def __init__(
self,
name: str,
data_type: DataType,
dimension: Optional[int] = 0,
index_param: Optional[
Union[HnswIndexParam, HnswRabitqIndexParam, FlatIndexParam, IVFIndexParam]
] = None,
):
if name is None or not isinstance(name, str):
raise ValueError(
f"schema validate failed: field name must be str, got {type(name).__name__}"
)
if not isinstance(dimension, int) or dimension < 0:
raise ValueError("schema validate failed: vector's dimension must be >= 0")
if data_type not in SUPPORT_VECTOR_DATA_TYPE:
raise ValueError(
f"schema validate failed: vector's data_type must be one of "
f"{', '.join(str(dt) for dt in SUPPORT_VECTOR_DATA_TYPE)}, "
f"but field[{name}]'s data_type is {data_type}"
)
if index_param is None:
index_param = FlatIndexParam()
self._cpp_obj = _FieldSchema(
name=name,
data_type=data_type,
dimension=dimension,
nullable=False,
index_param=index_param,
)
@classmethod
def _from_core(cls, core_field_schema: _FieldSchema):
inst = cls.__new__(cls)
inst._cpp_obj = core_field_schema
return inst
def _get_object(self) -> _FieldSchema:
return self._cpp_obj
@property
def name(self) -> str:
"""str: The name of the vector field."""
return self._cpp_obj.name
@property
def data_type(self) -> DataType:
"""DataType: The vector data type (e.g., VECTOR_FP32)."""
return self._cpp_obj.data_type
@property
def dimension(self) -> int:
"""int: The dimensionality of the vector."""
return self._cpp_obj.dimension
@property
def index_param(
self,
) -> Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]:
"""Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]: Index configuration for the vector."""
return self._cpp_obj.index_param
def __dict__(self) -> dict[str, Any]:
return {
"name": self.name,
"data_type": (
self.data_type.name
if hasattr(self.data_type, "name")
else str(self.data_type)
),
"dimension": self.dimension,
"index_param": (
self.index_param.to_dict() if self.index_param is not None else None
),
}
def __repr__(self) -> str:
try:
schema = self.__dict__()
return json.dumps(schema, indent=2, ensure_ascii=False)
except Exception as e:
return f"<FieldSchema error during repr: {e}>"
def __str__(self) -> str:
return self.__repr__()
def __eq__(self, other: object) -> bool:
if not isinstance(other, VectorSchema):
return False
return self._cpp_obj == other._cpp_obj
def __hash__(self) -> int:
return hash((self.name, self.data_type, self.dimension))
View File
+18
View File
@@ -0,0 +1,18 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from .util import require_module
__all__ = ["require_module"]
+63
View File
@@ -0,0 +1,63 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import importlib
from typing import Any, Optional
def require_module(module: str, mitigation: Optional[str] = None) -> Any:
"""Import a Python module and raise a user-friendly error if it is not available.
This utility helps provide actionable error messages when optional dependencies
are missing. It attempts to import the given module and, on failure, suggests
a `pip install` command based on either the module name or an optional
mitigation package name.
Args:
module (str): The full module name to import (e.g., ``"numpy"``, ``"pandas.io.parquet"``).
mitigation (Optional[str], optional): The package name to suggest for installation
if the import fails. If not provided, the top-level package of `module`
will be used (e.g., ``"pandas"`` for ``"pandas.io.parquet"``).
Returns:
Any: The imported module object.
Raises:
ImportError: If the module cannot be imported, with a clear installation hint.
Examples:
>>> import zvec
>>> np = zvec.require_module("numpy")
>>> pq = zvec.require_module("pyarrow.parquet", mitigation="pyarrow")
Note:
This function is intended for lazy-loading optional dependencies
with helpful error messages, not for core dependencies.
"""
try:
return importlib.import_module(module)
except ImportError as e:
package = mitigation or module
msg = f"Required package '{package}' is not installed. "
if "." in module:
top_level = module.split(".", maxsplit=1)[0]
msg += f"Module '{module}' is part of '{top_level}', "
if mitigation:
msg += f"please pip install '{mitigation}'."
else:
msg += f"please pip install '{top_level}'."
else:
msg += f"Please pip install '{package}'."
raise ImportError(msg) from e
+32
View File
@@ -0,0 +1,32 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from zvec._zvec.typing import (
DataType,
IndexType,
MetricType,
QuantizeType,
Status,
StatusCode,
)
__all__ = [
"DataType",
"IndexType",
"MetricType",
"QuantizeType",
"Status",
"StatusCode",
]
+404
View File
@@ -0,0 +1,404 @@
"""
This module contains the basic data types of Zvec
"""
from __future__ import annotations
import typing
__all__: list[str] = [
"DataType",
"IndexType",
"MetricType",
"QuantizeType",
"Status",
"StatusCode",
]
class DataType:
"""
Enumeration of supported data types in Zvec.
Includes scalar types, dense/sparse vector types, and array types.
Examples:
>>> import zvec
>>> print(zvec.DataType.FLOAT)
DataType.FLOAT
>>> print(zvec.DataType.VECTOR_FP32)
DataType.VECTOR_FP32
Members:
STRING
BOOL
INT32
INT64
FLOAT
DOUBLE
UINT32
UINT64
VECTOR_FP16
VECTOR_FP32
VECTOR_FP64
VECTOR_INT8
SPARSE_VECTOR_FP32
SPARSE_VECTOR_FP16
ARRAY_STRING
ARRAY_INT32
ARRAY_INT64
ARRAY_FLOAT
ARRAY_DOUBLE
ARRAY_BOOL
ARRAY_UINT32
ARRAY_UINT64
"""
ARRAY_BOOL: typing.ClassVar[DataType] # value = <DataType.ARRAY_BOOL: 42>
ARRAY_DOUBLE: typing.ClassVar[DataType] # value = <DataType.ARRAY_DOUBLE: 48>
ARRAY_FLOAT: typing.ClassVar[DataType] # value = <DataType.ARRAY_FLOAT: 47>
ARRAY_INT32: typing.ClassVar[DataType] # value = <DataType.ARRAY_INT32: 43>
ARRAY_INT64: typing.ClassVar[DataType] # value = <DataType.ARRAY_INT64: 44>
ARRAY_STRING: typing.ClassVar[DataType] # value = <DataType.ARRAY_STRING: 41>
ARRAY_UINT32: typing.ClassVar[DataType] # value = <DataType.ARRAY_UINT32: 45>
ARRAY_UINT64: typing.ClassVar[DataType] # value = <DataType.ARRAY_UINT64: 46>
BOOL: typing.ClassVar[DataType] # value = <DataType.BOOL: 3>
DOUBLE: typing.ClassVar[DataType] # value = <DataType.DOUBLE: 9>
FLOAT: typing.ClassVar[DataType] # value = <DataType.FLOAT: 8>
INT32: typing.ClassVar[DataType] # value = <DataType.INT32: 4>
INT64: typing.ClassVar[DataType] # value = <DataType.INT64: 5>
SPARSE_VECTOR_FP16: typing.ClassVar[
DataType
] # value = <DataType.SPARSE_VECTOR_FP16: 30>
SPARSE_VECTOR_FP32: typing.ClassVar[
DataType
] # value = <DataType.SPARSE_VECTOR_FP32: 31>
STRING: typing.ClassVar[DataType] # value = <DataType.STRING: 2>
UINT32: typing.ClassVar[DataType] # value = <DataType.UINT32: 6>
UINT64: typing.ClassVar[DataType] # value = <DataType.UINT64: 7>
VECTOR_FP16: typing.ClassVar[DataType] # value = <DataType.VECTOR_FP16: 22>
VECTOR_FP32: typing.ClassVar[DataType] # value = <DataType.VECTOR_FP32: 23>
VECTOR_FP64: typing.ClassVar[DataType] # value = <DataType.VECTOR_FP64: 24>
VECTOR_INT8: typing.ClassVar[DataType] # value = <DataType.VECTOR_INT8: 26>
__members__: typing.ClassVar[
dict[str, DataType]
] # value = {'STRING': <DataType.STRING: 2>, 'BOOL': <DataType.BOOL: 3>, 'INT32': <DataType.INT32: 4>, 'INT64': <DataType.INT64: 5>, 'FLOAT': <DataType.FLOAT: 8>, 'DOUBLE': <DataType.DOUBLE: 9>, 'UINT32': <DataType.UINT32: 6>, 'UINT64': <DataType.UINT64: 7>, 'VECTOR_FP16': <DataType.VECTOR_FP16: 22>, 'VECTOR_FP32': <DataType.VECTOR_FP32: 23>, 'VECTOR_FP64': <DataType.VECTOR_FP64: 24>, 'VECTOR_INT8': <DataType.VECTOR_INT8: 26>, 'SPARSE_VECTOR_FP32': <DataType.SPARSE_VECTOR_FP32: 31>, 'SPARSE_VECTOR_FP16': <DataType.SPARSE_VECTOR_FP16: 30>, 'ARRAY_STRING': <DataType.ARRAY_STRING: 41>, 'ARRAY_INT32': <DataType.ARRAY_INT32: 43>, 'ARRAY_INT64': <DataType.ARRAY_INT64: 44>, 'ARRAY_FLOAT': <DataType.ARRAY_FLOAT: 47>, 'ARRAY_DOUBLE': <DataType.ARRAY_DOUBLE: 48>, 'ARRAY_BOOL': <DataType.ARRAY_BOOL: 42>, 'ARRAY_UINT32': <DataType.ARRAY_UINT32: 45>, 'ARRAY_UINT64': <DataType.ARRAY_UINT64: 46>}
def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __init__(self, value: typing.SupportsInt) -> None: ...
def __int__(self) -> int: ...
def __ne__(self, other: typing.Any) -> bool: ...
def __repr__(self) -> str: ...
def __setstate__(self, state: typing.SupportsInt) -> None: ...
def __str__(self) -> str: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
class IndexType:
"""
Enumeration of supported index types in Zvec.
Examples:
>>> import zvec
>>> print(zvec.IndexType.HNSW)
IndexType.HNSW
Members:
UNDEFINED
HNSW
IVF
FLAT
INVERT
"""
FLAT: typing.ClassVar[IndexType] # value = <IndexType.FLAT: 4>
HNSW: typing.ClassVar[IndexType] # value = <IndexType.HNSW: 1>
INVERT: typing.ClassVar[IndexType] # value = <IndexType.INVERT: 10>
IVF: typing.ClassVar[IndexType] # value = <IndexType.IVF: 3>
UNDEFINED: typing.ClassVar[IndexType] # value = <IndexType.UNDEFINED: 0>
__members__: typing.ClassVar[
dict[str, IndexType]
] # value = {'UNDEFINED': <IndexType.UNDEFINED: 0>, 'HNSW': <IndexType.HNSW: 1>, 'IVF': <IndexType.IVF: 3>, 'FLAT': <IndexType.FLAT: 4>, 'INVERT': <IndexType.INVERT: 10>}
def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __init__(self, value: typing.SupportsInt) -> None: ...
def __int__(self) -> int: ...
def __ne__(self, other: typing.Any) -> bool: ...
def __repr__(self) -> str: ...
def __setstate__(self, state: typing.SupportsInt) -> None: ...
def __str__(self) -> str: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
class MetricType:
"""
Enumeration of supported distance/similarity metrics.
- COSINE: Cosine similarity.
- IP: Inner product (dot product).
- L2: Euclidean distance (L2 norm).
Examples:
>>> import zvec
>>> print(zvec.MetricType.COSINE)
MetricType.COSINE
Members:
COSINE
IP
L2
"""
COSINE: typing.ClassVar[MetricType] # value = <MetricType.COSINE: 3>
IP: typing.ClassVar[MetricType] # value = <MetricType.IP: 2>
L2: typing.ClassVar[MetricType] # value = <MetricType.L2: 1>
__members__: typing.ClassVar[
dict[str, MetricType]
] # value = {'COSINE': <MetricType.COSINE: 3>, 'IP': <MetricType.IP: 2>, 'L2': <MetricType.L2: 1>}
def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __init__(self, value: typing.SupportsInt) -> None: ...
def __int__(self) -> int: ...
def __ne__(self, other: typing.Any) -> bool: ...
def __repr__(self) -> str: ...
def __setstate__(self, state: typing.SupportsInt) -> None: ...
def __str__(self) -> str: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
class QuantizeType:
"""
Enumeration of supported quantization types for vector compression.
Examples:
>>> import zvec
>>> print(zvec.QuantizeType.INT8)
QuantizeType.INT8
Members:
UNDEFINED
FP16
INT8
INT4
"""
FP16: typing.ClassVar[QuantizeType] # value = <QuantizeType.FP16: 1>
INT4: typing.ClassVar[QuantizeType] # value = <QuantizeType.INT4: 3>
INT8: typing.ClassVar[QuantizeType] # value = <QuantizeType.INT8: 2>
UNDEFINED: typing.ClassVar[QuantizeType] # value = <QuantizeType.UNDEFINED: 0>
__members__: typing.ClassVar[
dict[str, QuantizeType]
] # value = {'UNDEFINED': <QuantizeType.UNDEFINED: 0>, 'FP16': <QuantizeType.FP16: 1>, 'INT8': <QuantizeType.INT8: 2>, 'INT4': <QuantizeType.INT4: 3>}
def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __init__(self, value: typing.SupportsInt) -> None: ...
def __int__(self) -> int: ...
def __ne__(self, other: typing.Any) -> bool: ...
def __repr__(self) -> str: ...
def __setstate__(self, state: typing.SupportsInt) -> None: ...
def __str__(self) -> str: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
class Status:
"""
Represents the outcome of a Zvec operation.
A `Status` object is either OK (success) or carries an error code and message.
Examples:
>>> from zvec.typing import Status, StatusCode
>>> s = Status()
>>> print(s.ok())
True
>>> s = Status(StatusCode.INVALID_ARGUMENT, "Field not found")
>>> print(s.code() == StatusCode.INVALID_ARGUMENT)
True
>>> print(s.message())
Field not found
"""
__hash__: typing.ClassVar[None] = None
@staticmethod
def AlreadyExists(message: str) -> Status: ...
@staticmethod
def InternalError(message: str) -> Status: ...
@staticmethod
def InvalidArgument(message: str) -> Status: ...
@staticmethod
def NotFound(message: str) -> Status: ...
@staticmethod
def OK() -> Status:
"""
Create an OK status.
"""
@staticmethod
def PermissionDenied(message: str) -> Status: ...
def __eq__(self, arg0: Status) -> bool: ...
@typing.overload
def __init__(self) -> None: ...
@typing.overload
def __init__(self, code: StatusCode, message: str = "") -> None:
"""
Construct a status with the given code and optional message.
Args:
code (StatusCode): The status code.
message (str, optional): Error message. Defaults to empty string.
"""
def __ne__(self, arg0: Status) -> bool: ...
def __repr__(self) -> str: ...
def code(self) -> StatusCode:
"""
StatusCode: Returns the status code.
"""
def message(self) -> str:
"""
str: Returns the error message (may be empty).
"""
def ok(self) -> bool:
"""
bool: Returns True if the status is OK.
"""
class StatusCode:
"""
Enumeration of possible status codes for Zvec operations.
Used by the `Status` class to indicate success or failure reason.
Members:
OK
NOT_FOUND
ALREADY_EXISTS
INVALID_ARGUMENT
PERMISSION_DENIED
FAILED_PRECONDITION
RESOURCE_EXHAUSTED
UNAVAILABLE
INTERNAL_ERROR
NOT_SUPPORTED
UNKNOWN
"""
ALREADY_EXISTS: typing.ClassVar[
StatusCode
] # value = <StatusCode.ALREADY_EXISTS: 2>
FAILED_PRECONDITION: typing.ClassVar[
StatusCode
] # value = <StatusCode.FAILED_PRECONDITION: 5>
INTERNAL_ERROR: typing.ClassVar[
StatusCode
] # value = <StatusCode.INTERNAL_ERROR: 8>
INVALID_ARGUMENT: typing.ClassVar[
StatusCode
] # value = <StatusCode.INVALID_ARGUMENT: 3>
NOT_FOUND: typing.ClassVar[StatusCode] # value = <StatusCode.NOT_FOUND: 1>
NOT_SUPPORTED: typing.ClassVar[StatusCode] # value = <StatusCode.NOT_SUPPORTED: 9>
OK: typing.ClassVar[StatusCode] # value = <StatusCode.OK: 0>
PERMISSION_DENIED: typing.ClassVar[
StatusCode
] # value = <StatusCode.PERMISSION_DENIED: 4>
RESOURCE_EXHAUSTED: typing.ClassVar[
StatusCode
] # value = <StatusCode.RESOURCE_EXHAUSTED: 6>
UNAVAILABLE: typing.ClassVar[StatusCode] # value = <StatusCode.UNAVAILABLE: 7>
UNKNOWN: typing.ClassVar[StatusCode] # value = <StatusCode.UNKNOWN: 10>
__members__: typing.ClassVar[
dict[str, StatusCode]
] # value = {'OK': <StatusCode.OK: 0>, 'NOT_FOUND': <StatusCode.NOT_FOUND: 1>, 'ALREADY_EXISTS': <StatusCode.ALREADY_EXISTS: 2>, 'INVALID_ARGUMENT': <StatusCode.INVALID_ARGUMENT: 3>, 'PERMISSION_DENIED': <StatusCode.PERMISSION_DENIED: 4>, 'FAILED_PRECONDITION': <StatusCode.FAILED_PRECONDITION: 5>, 'RESOURCE_EXHAUSTED': <StatusCode.RESOURCE_EXHAUSTED: 6>, 'UNAVAILABLE': <StatusCode.UNAVAILABLE: 7>, 'INTERNAL_ERROR': <StatusCode.INTERNAL_ERROR: 8>, 'NOT_SUPPORTED': <StatusCode.NOT_SUPPORTED: 9>, 'UNKNOWN': <StatusCode.UNKNOWN: 10>}
def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __init__(self, value: typing.SupportsInt) -> None: ...
def __int__(self) -> int: ...
def __ne__(self, other: typing.Any) -> bool: ...
def __repr__(self) -> str: ...
def __setstate__(self, state: typing.SupportsInt) -> None: ...
def __str__(self) -> str: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
+62
View File
@@ -0,0 +1,62 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from enum import IntEnum
__all__ = ["LogLevel", "LogType"]
class LogLevel(IntEnum):
"""Enumeration of logging severity levels, ordered from lowest to highest priority.
Used to control verbosity and filtering of log messages. Higher numeric values
indicate more severe conditions.
Note:
``WARNING`` is an alias for ``WARN`` to match Python's built-in :mod:`logging`
module convention.
Attributes:
DEBUG (int): Detailed information, typically of interest only when diagnosing problems.
INFO (int): Confirmation that things are working as expected.
WARN (int): An indication that something unexpected happened, or indicative of
potential future problems. (Alias: ``WARNING``)
WARNING (int): Same as ``WARN``.
ERROR (int): Due to a more serious problem, the software has not been able
to perform some function.
FATAL (int): A serious error, indicating that the program itself may be unable
to continue running.
"""
DEBUG = 0
INFO = 1
WARN = 2
WARNING = 2
ERROR = 3
FATAL = 4
class LogType(IntEnum):
"""Enumeration of supported log output destinations.
Specifies where log messages should be written.
Attributes:
CONSOLE (int): Output logs to standard output/error (e.g., terminal or IDE console).
FILE (int): Write logs to a persistent file on disk.
"""
CONSOLE = 0
FILE = 1
+246
View File
@@ -0,0 +1,246 @@
# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import Optional
from zvec._zvec import Initialize, _Collection
from .model import Collection
from .model.param import CollectionOption
from .model.schema import CollectionSchema
__all__ = ["create_and_open", "init", "open"]
from .typing.enum import LogLevel, LogType
def init(
*,
log_type: Optional[LogType] = None,
log_level: Optional[LogLevel] = None,
log_dir: Optional[str] = "./logs",
log_basename: Optional[str] = "zvec.log",
log_file_size: Optional[int] = 2048,
log_overdue_days: Optional[int] = 7,
query_threads: Optional[int] = None,
optimize_threads: Optional[int] = None,
invert_to_forward_scan_ratio: Optional[float] = None,
brute_force_by_keys_ratio: Optional[float] = None,
fts_brute_force_by_keys_ratio: Optional[float] = None,
memory_limit_mb: Optional[int] = None,
jieba_dict_dir: Optional[str] = None,
) -> None:
"""Initialize Zvec with configuration options.
This function must be called before any other operation.
It can only be called once — subsequent calls raise a ``RuntimeError``.
Parameters set to ``None`` are **omitted** from the configuration and
fall back to Zvec's internal defaults, which may be derived from the runtime
environment (e.g., cgroup CPU/memory limits). Explicitly provided values
always override defaults.
Args:
log_type (Optional[LogType], optional): Logger destination.
- ``LogType.CONSOLE`` (default if omitted or set to this)
- ``LogType.FILE``
- If ``None``, uses internal default (currently ``CONSOLE``).
log_level (Optional[LogLevel], optional): Minimum log severity.
Default: ``LogLevel.WARN``.
Accepted values: ``DEBUG``, ``INFO``, ``WARN``, ``ERROR``, ``FATAL``.
If ``None``, uses internal default (``WARN``).
log_dir (Optional[str], optional):
Directory for log files (only used when ``log_type=FILE``).
Parent directories are **not** created automatically.
Default: ``"./logs"``.
If ``None``, internal default is used.
log_basename (Optional[str], optional):
Base name for rotated log files (e.g., ``zvec.log.1``, ``zvec.log.2``).
Default: ``"zvec.log"``.
log_file_size (Optional[int], optional):
Max size per log file in **MB** before rotation.
Default: ``2048`` MB (2 GB).
log_overdue_days (Optional[int], optional):
Days to retain rotated log files before deletion.
Default: ``7`` days.
query_threads (Optional[int], optional):
Number of threads for query execution.
If ``None`` (default), inferred from available CPU cores (via cgroup).
Must be ≥ 1 if provided.
optimize_threads (Optional[int], optional):
Threads for background tasks (e.g., compaction, indexing).
If ``None``, defaults to same as ``query_threads`` or CPU count.
invert_to_forward_scan_ratio (Optional[float], optional):
Threshold to switch from inverted index to full forward scan.
Range: [0.0, 1.0]. Higher → more aggressive index skipping.
Default: ``0.9`` (if omitted).
brute_force_by_keys_ratio (Optional[float], optional):
Threshold to use brute-force key lookup over index.
Lower → prefer index; higher → prefer brute-force.
Range: [0.0, 1.0]. Default: ``0.1``.
fts_brute_force_by_keys_ratio (Optional[float], optional):
Threshold to switch FTS scan from posting-driven to
candidate-driven (brute-force) when the invert filter is
highly selective. Independent from ``brute_force_by_keys_ratio``
because per-candidate FTS cost is higher.
Range: [0.0, 1.0]. Default: ``0.05``.
memory_limit_mb (Optional[int], optional):
Soft memory cap in MB. Zvec may throttle or fail operations
approaching this limit.
If ``None``, inferred from cgroup memory limit * 0.8 (e.g., in Docker).
Must be > 0 if provided.
jieba_dict_dir (Optional[str], optional):
Override the default directory containing ``jieba.dict.utf8`` and
``hmm_model.utf8`` for the jieba FTS tokenizer. When ``None``, the
value previously registered by ``zvec.set_default_jieba_dict_dir``
(called automatically on ``import zvec`` to point at the wheel's
bundled dict) is preserved. JiebaTokenizer also honors the
``ZVEC_JIEBA_DICT_DIR`` environment variable and per-field
``FtsIndexParam.extra_params.jieba_dict_dir`` ahead of this value.
Raises:
RuntimeError: If Zvec is already initialized.
ValueError: On invalid values (e.g., negative thread count, log level out of range).
TypeError: If a value has incorrect type (e.g., string for ``query_threads``).
Note:
- All ``None`` arguments are **excluded** from the configuration payload,
allowing the core library to apply environment-aware defaults.
- This design ensures container-friendliness: in Kubernetes/Docker,
omitting ``memory_limit_mb`` and thread counts lets Zvec auto-adapt.
Examples:
Initialize with defaults (log to console, auto-detect resources):
>>> import zvec
>>> zvec.init()
Customize logging to file with rotation:
>>> zvec.init(
... log_type=LogType.FILE,
... log_dir="/var/log/zvec",
... log_file_size=1024,
... log_overdue_days=30
... )
Limit resources explicitly:
>>> zvec.init(
... memory_limit_mb=2048,
... query_threads=4,
... optimize_threads=2
... )
Fine-tune query heuristics:
>>> zvec.init(
... invert_to_forward_scan_ratio=0.95,
... brute_force_by_keys_ratio=0.05
... )
"""
# Build config dict, skipping None values
config_dict = {}
if log_type is not None:
if not isinstance(log_type, LogType):
raise TypeError("log_type must be LogType")
config_dict["log_type"] = log_type.name
if log_level is not None:
if not isinstance(log_level, LogLevel):
raise TypeError("log_level must be LogLevel")
config_dict["log_level"] = log_level.name
if log_dir is not None:
config_dict["log_dir"] = log_dir
if log_basename is not None:
config_dict["log_basename"] = log_basename
if log_file_size is not None:
config_dict["log_file_size"] = log_file_size
if log_overdue_days is not None:
config_dict["log_overdue_days"] = log_overdue_days
if query_threads is not None:
config_dict["query_threads"] = query_threads
if optimize_threads is not None:
config_dict["optimize_threads"] = optimize_threads
if invert_to_forward_scan_ratio is not None:
config_dict["invert_to_forward_scan_ratio"] = invert_to_forward_scan_ratio
if brute_force_by_keys_ratio is not None:
config_dict["brute_force_by_keys_ratio"] = brute_force_by_keys_ratio
if fts_brute_force_by_keys_ratio is not None:
config_dict["fts_brute_force_by_keys_ratio"] = fts_brute_force_by_keys_ratio
if memory_limit_mb is not None:
config_dict["memory_limit_mb"] = memory_limit_mb
if jieba_dict_dir is not None:
config_dict["jieba_dict_dir"] = jieba_dict_dir
Initialize(config_dict)
def create_and_open(
path: str,
schema: CollectionSchema,
option: Optional[CollectionOption] = None,
) -> Collection:
"""Create a new collection and open it for use.
If a collection already exists at the given path, it may raise an error
depending on the underlying implementation.
Args:
path (str): Path or name of the collection to create.
schema (CollectionSchema): Schema defining the structure of the collection.
option (Optional[CollectionOption]): Configuration options
for opening the collection. Defaults to a default-constructed
``CollectionOption()`` if not provided.
Returns:
Collection: An opened collection instance ready for operations.
Examples:
>>> import zvec
>>> schema = zvec.CollectionSchema(
... name="my_collection",
... fields=[zvec.FieldSchema("id", zvec.DataType.INT64, nullable=True)]
... )
>>> coll = create_and_open("./my_collection", schema)
"""
if not isinstance(path, str):
raise TypeError("path must be a string")
if not isinstance(schema, CollectionSchema):
raise TypeError("schema must be a CollectionSchema")
option = option or CollectionOption()
if not isinstance(option, CollectionOption):
raise TypeError("option must be a CollectionOption")
_collection = _Collection.CreateAndOpen(path, schema._get_object(), option)
return Collection._from_core(_collection)
def open(path: str, option: CollectionOption = CollectionOption()) -> Collection:
"""Open an existing collection from disk.
The collection must have been previously created with ``create_and_open``.
Args:
path (str): Path or name of the existing collection.
option (CollectionOption, optional): Configuration options
for opening the collection. Defaults to a default-constructed
``CollectionOption()`` if not provided.
Returns:
Collection: An opened collection instance.
Examples:
>>> import zvec
>>> coll = zvec.open("./my_collection")
"""
_collection = _Collection.Open(path, option)
return Collection._from_core(_collection)