xiaoxue_admin/zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: change rerank interface from map-based to vector-based (#458)

Browse Source 
* 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
This commit is contained in:
Cuiys
2026-06-04 16:03:28 +08:00
committed by GitHub
parent f562bdd636
commit c46efe1241
29 changed files with 704 additions and 1098 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_query_executor.py
+64 -145
View File
@@ -14,20 +14,16 @@
from __future__ import annotations
from typing import Dict, Union
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
import numpy as np
import math
from _zvec.param import _VectorQuery
from _zvec.param import _SearchQuery
import pytest
from zvec.executor.query_executor import (
MultiVectorQueryExecutor,
NoVectorQueryExecutor,
QueryContext,
QueryExecutor,
QueryExecutorFactory,
SingleVectorQueryExecutor,
)
from zvec import (
RrfReRanker,
@@ -43,21 +39,6 @@ from zvec import (
from zvec.extension.multi_vector_reranker import CallbackReRanker
# ----------------------------
# Mock Vector Schema
# ----------------------------
class MockVectorSchema(VectorSchema):
def __init__(self, name="test_vector"):
self._name = name
@property
def name(self):
return self._name
def _get_object(self):
return MagicMock()
# ----------------------------
# Mock Collection Schema
# ----------------------------
@@ -110,7 +91,7 @@ class TestQuery:
assert query.has_vector()
def test_validate_dense_fp16_convert(self):
v = _VectorQuery()
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)
@@ -118,7 +99,7 @@ class TestQuery:
assert np.array_equal(vec, ret)
def test_validate_dense_fp32_convert(self):
v = _VectorQuery()
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)
@@ -126,7 +107,7 @@ class TestQuery:
assert np.array_equal(vec, ret)
def test_validate_dense_fp64_convert(self):
v = _VectorQuery()
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)
@@ -134,7 +115,7 @@ class TestQuery:
assert np.array_equal(vec, ret)
def test_validate_dense_int8_convert(self):
v = _VectorQuery()
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)
@@ -142,7 +123,7 @@ class TestQuery:
assert np.array_equal(vec, ret)
def test_validate_sparse_fp32_convert(self):
v = _VectorQuery()
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)
@@ -151,7 +132,7 @@ class TestQuery:
assert math.isclose(vec[k], ret[k], abs_tol=1e-6)
def test_validate_sparse_fp16_convert(self):
v = _VectorQuery()
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)
@@ -189,7 +170,6 @@ class TestQueryContext:
assert ctx.reranker is None
assert ctx.output_fields is None
assert ctx.include_vector is False
assert ctx.core_vectors == []
def test_properties(self):
queries = [Query(field_name="test")]
@@ -215,9 +195,7 @@ class TestQueryContext:
def test_properties_with_weighted_reranker(self):
queries = [Query(field_name="test")]
reranker = WeightedReRanker(
topn=10,
metrics={"test": MetricType.L2},
weights={"test": 1.0},
weights=[1.0],
)
ctx = QueryContext(
@@ -227,13 +205,12 @@ class TestQueryContext:
)
assert ctx.reranker == reranker
assert ctx.reranker.weights == {"test": 1.0}
assert ctx.reranker.metrics == {"test": MetricType.L2}
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, topn=10)
reranker = CallbackReRanker(callback=cb)
ctx = QueryContext(
topk=5,
@@ -243,141 +220,83 @@ class TestQueryContext:
assert ctx.reranker == reranker
def test_core_vectors_setter(self):
ctx = QueryContext(topk=10)
core_vectors = [MagicMock()]
ctx.core_vectors = core_vectors
assert ctx.core_vectors == core_vectors
class TestNoVectorQueryExecutor:
class TestQueryExecutor:
def test_init(self):
schema = MockCollectionSchema()
executor = NoVectorQueryExecutor(schema)
executor = QueryExecutor(schema)
assert isinstance(executor, QueryExecutor)
def test_do_validate_with_queries(self):
def test_do_build_without_queries(self):
# When no queries are given, build a single vector-less query.
schema = MockCollectionSchema()
executor = NoVectorQueryExecutor(schema)
ctx = QueryContext(
topk=10, queries=[Query(field_name="test", vector=[0.1, 0.2, 0.3])]
)
with pytest.raises(
ValueError, match="Collection does not support query with vector or id"
):
executor._do_validate(ctx)
def test_do_validate_without_queries(self):
schema = MockCollectionSchema()
executor = NoVectorQueryExecutor(schema)
ctx = QueryContext(topk=10)
executor._do_validate(ctx)
def test_do_build(self):
schema = MockCollectionSchema()
executor = NoVectorQueryExecutor(schema)
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5, filter="test_filter")
result = executor._do_build(ctx, MagicMock())
result = executor._build_queries(ctx, MagicMock())
assert len(result) == 1
assert result[0].topk == 5
assert result[0].filter == "test_filter"
class TestSingleVectorQueryExecutor:
def test_init(self):
def test_do_build_query_wo_vector(self):
# Vector-less core query should carry the context query params.
schema = MockCollectionSchema()
executor = SingleVectorQueryExecutor(schema)
assert isinstance(executor, NoVectorQueryExecutor)
executor = QueryExecutor(schema)
ctx = QueryContext(topk=7, filter="f", include_vector=True)
def test_do_validate_multiple_queries(self):
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 = SingleVectorQueryExecutor(schema)
queries = [Query(field_name="test1"), Query(field_name="test2")]
ctx = QueryContext(topk=10, queries=queries)
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
docs_list = [["doc1", "doc2"]]
with pytest.raises(
ValueError,
match="Collection has only one vector field, cannot query with multiple vectors",
):
executor._do_validate(ctx)
result = executor._merge_and_rerank(ctx, docs_list)
assert result == ["doc1", "doc2"]
def test_do_build_without_queries(self):
def test_do_merge_rerank_results_empty(self):
# Empty results should raise an error.
schema = MockCollectionSchema()
executor = SingleVectorQueryExecutor(schema)
executor = QueryExecutor(schema)
ctx = QueryContext(topk=5)
result = executor._do_build(ctx, MagicMock())
assert len(result) == 1
assert result[0].topk == 5
with pytest.raises(ValueError, match="Query results is empty"):
executor._merge_and_rerank(ctx, [])
class TestMultiVectorQueryExecutor:
def test_init(self):
def test_do_merge_rerank_results_with_reranker(self):
# Multiple result lists are merged through the reranker.
schema = MockCollectionSchema()
executor = MultiVectorQueryExecutor(schema)
assert isinstance(executor, SingleVectorQueryExecutor)
def test_do_validate_multiple_queries_without_reranker(self):
schema = MockCollectionSchema()
executor = MultiVectorQueryExecutor(schema)
queries = [Query(field_name="test1"), Query(field_name="test2")]
ctx = QueryContext(topk=10, queries=queries)
with pytest.raises(ValueError, match="Reranker is required for multi-query"):
executor._do_validate(ctx)
def test_do_validate_multiple_queries_with_reranker(self):
schema = MockCollectionSchema()
executor = MultiVectorQueryExecutor(schema)
queries = [Query(field_name="test1"), Query(field_name="test2")]
reranker = RrfReRanker()
ctx = QueryContext(topk=10, queries=queries, reranker=reranker)
executor._do_validate(ctx)
def test_do_validate_multiple_queries_with_weighted_reranker(self):
schema = MockCollectionSchema()
executor = MultiVectorQueryExecutor(schema)
queries = [Query(field_name="test1"), Query(field_name="test2")]
reranker = WeightedReRanker(
topn=10,
metrics={"test1": MetricType.L2, "test2": MetricType.L2},
weights={"test1": 0.7, "test2": 0.3},
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,
)
ctx = QueryContext(topk=10, queries=queries, reranker=reranker)
docs_list = [["d1"], ["d2"]]
executor._do_validate(ctx)
result = executor._merge_and_rerank(ctx, docs_list)
assert result == ["merged"]
reranker.rerank.assert_called_once_with(docs_list, ctx.topk)
def test_do_validate_multiple_queries_with_callback_reranker(self):
def test_execute_python_pipeline(self):
# Each query is executed serially and converted into a result list.
schema = MockCollectionSchema()
executor = MultiVectorQueryExecutor(schema)
queries = [Query(field_name="test1"), Query(field_name="test2")]
reranker = CallbackReRanker(
callback=lambda query_results, topn: [],
topn=10,
)
ctx = QueryContext(topk=10, queries=queries, reranker=reranker)
executor = QueryExecutor(schema)
collection = MagicMock()
collection.Query.side_effect = [["raw1"], ["raw2"]]
vectors = [MagicMock(), MagicMock()]
executor._do_validate(ctx)
class TestQueryExecutorFactory:
def test_create_no_vectors(self):
schema = MockCollectionSchema()
executor = QueryExecutorFactory.create(schema)
assert isinstance(executor, NoVectorQueryExecutor)
def test_create_single_vector(self):
schema = MockCollectionSchema(vectors=MockVectorSchema())
executor = QueryExecutorFactory.create(schema)
assert isinstance(executor, SingleVectorQueryExecutor)
def test_create_multiple_vectors(self):
schema = MockCollectionSchema(
vectors={"test1": MockVectorSchema(), "test2": MockVectorSchema()}
)
executor = QueryExecutorFactory.create(schema)
assert isinstance(executor, MultiVectorQueryExecutor)
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