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