Cuiys
c46efe1241
* refactor: change rerank interface from map-based to vector-based (#452) - Define QueryResult = list[Doc] type alias in doc.py - Change C++ Reranker::rerank() signature from map<string, DocPtrList> to vector<DocPtrList> - Extend bind_schema() to accept field_names for index-based field lookup - Update ScoreBasedReranker/WeightedReranker/CallbackReranker implementations - Adapt collection.cc MultiQuery path to use vector<DocPtrList> - Update Python binding to expose rerank() and use vector<double> weights - Refactor Python RerankFunction interface to list[QueryResult] -> QueryResult - Remove Python-layer rerank logic from RrfReRanker/WeightedReRanker (delegate to C++) - Update query_executor to return list[list[Doc]] instead of dict - Update all related unit tests (C++ and Python) * refactor: replace list[Doc] with QueryResult type alias in executor and rerank functions * refactor: replace list[list[Doc]] with list[QueryResult] in query_executor * fix: remove unused Doc import in rerank_function.py (ruff F401) * refactor(query_executor): merge duplicate rerank return paths * refactor: RrfReRanker/WeightedReRanker.rerank() directly call C++ reranker * refactor: simplify QueryExecutor into unified class, remove Factory/subclasses/validation/concurrency * refactor: rename _VectorQuery to _SearchQuery, from_vector_query to from_search_query * refactor(query_executor): split execute into single/multi paths, rename core_vector to search_query, drop unused core_vectors * style: apply ruff formatter to test_reranker.py and query_executor.py * refactor: make rescore() private in ScoreBasedReranker hierarchy * style: apply clang-format to reranker.h * style: apply clang-format to all modified C++ files * refactor: rename private methods in QueryExecutor for clearer semantics * refactor: rename mvq to multi_query for clarity * fix: make BasicRRF test order-independent for equal scores * fix: update collection_test to use vector-based reranker interface * fix: update reranker tests to expect TypeError instead of NotImplementedError * refactor: remove PendingQuery wrapper, use SearchQuery directly in MultiQuery path * refactor: simplify MultiQuery path - remove seen_fields, merge field_names into main loop * fix: address review comments - defensive checks and remove fields param from C API - ScoreBasedReranker::rerank(): early return empty list when topn <= 0 - WeightedReranker::rescore(): null-check schema_ before use - CallbackReranker::rerank(): check callback_ is not empty before invoke - C API zvec_reranker_create_weighted(): remove unused fields parameter * fix: remove duplicate field name test (check was intentionally removed) * fix: address egolearner review comments - Rename QueryResult to DocList for clarity (见名知义) - Change docstring to #: comment for type alias - Fix output_fields check: use 'is not None' instead of truthy check (None means unset, [] means explicit empty list - different semantics) - Raise ValueError when search-by-id finds no document * refactor: remove redundant output_fields assignment in _build_search_query * refactor: address egolearner review comments (C++ refactoring) - c_api.cc: simplify weighted reranker creation with inline vector ctor - python_reranker.cc: refactor unwrap_rerank_result - take by value, early error return, move semantics - Rename C API functions for consistent naming: zvec_reranker_create_rrf -> zvec_create_rrf_reranker zvec_reranker_create_weighted -> zvec_create_weighted_reranker zvec_reranker_destroy -> zvec_destroy_reranker zvec_reranker_get_rank_constant -> zvec_get_reranker_rank_constant - reranker.h/cc: bind_schema returns Result<void>, caches vector<const FieldSchema*> to avoid repeated schema lookups in rescore - python_param.cc: rename py::arg vector_query to search_query * revert: rollback bind_schema refactoring due to thread-safety concern The field_schemas_ caching approach introduces a data race when the same WeightedReranker instance is shared across concurrent queries: bind_schema() writes field_schemas_ while rerank() reads it concurrently. Revert to storing schema_ + field_names_ and looking up fields in rescore(). Add @note thread-safety warning to WeightedReranker class documentation. * fix: unify error message format in collection.cc Change 'Vector field not found: X' to 'Invalid query: field X not found' for consistent error formatting as suggested by zhourrr. * fix: sort __all__ and remove duplicates in __init__.pyi Fix RUF022 lint error: sort __all__ alphabetically and remove duplicate entries (DenseEmbeddingFunction, ReRanker). * style: format query_executor.py with ruff formatter * fix: resolve Python test failures after FTS rebase integration - test_query_executor.py: update method names to match refactored API (_do_build -> _build_queries, _do_merge_rerank_results -> _merge_and_rerank) - test_reranker.py: fix expected exception type (TypeError from pybind11) - test_collection_fts.py: update error message match patterns - test_collection_fts_vector_hybrid.py: remove obsolete 'metrics' param, update weights from dict to positional list, adapt validation tests for multi-vector queries (now supported with reranker) - test_collection_dql.py: remove 'metrics' param, update weights format - collection.cc: distinguish FTS vs vector fields in MultiQuery path using get_fts_clause() to route field lookup correctly - reranker.cc: use get_field() instead of get_vector_field() in rescore to support FTS+vector hybrid weighted reranking * refactor: pass topn as rerank() parameter, move rerank_field to model rerankers * fix: address review comments - rename test functions and restore duplicate field check * refactor: simplify MultiQuery field lookup, let validate_and_sanitize handle type check
437 lines
14 KiB
Python
437 lines
14 KiB
Python
# 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,
|
|
DataType,
|
|
VectorSchema,
|
|
)
|
|
|
|
from _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 == 0
|
|
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
|
|
|
|
def test_custom(self):
|
|
param = HnswQueryParam(ef=10, is_using_refiner=True, radius=30, is_linear=True)
|
|
assert param.ef == 10
|
|
assert param.is_using_refiner == True
|
|
assert param.radius == 30
|
|
assert param.is_linear == True
|
|
|
|
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()
|
|
|
|
|
|
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)
|