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

feat: migrate multi-vector query and reranker logic to C++ (#405)

Browse Source 
* feat: migrate multi-vector query and reranker logic to C++

- Add Reranker base class with RrfReRanker and WeightedReRanker implementations
- Add Collection::MultiQuery interface for multi-vector queries with reranking
- Add MultiVectorQuery struct in doc.h with forward declaration for Reranker
- Add C API bindings for reranker and MultiQuery (zvec_reranker_*, zvec_multi_vector_query_*, zvec_collection_multi_query)
- Add Python binding for reranker classes with py::function bridge for callback
- Validate duplicate field names in multi-vector queries (C++ and Python consistent)
- Remove TODO comment about concurrent execution (SQLEngine is not thread-safe)
- Update collection.h MultiQuery doc comment from concurrently to sequentially
- Add C++ collection tests (6 MultiQuery test cases)
- Add C API tests (reranker functions + multi_vector_query end-to-end)
- Implement Python test cases (11 previously skipped tests now active)
- Simplify Python query_executor validation for unified duplicate field check

* style: format Python files with ruff

* fix: adapt to main branch API changes (VectorQuery->Query rename, validate_and_sanitize)

* fix: multi_vector tests now use multiple same-type vector fields (dense2, sparse2)

* fix: suppress RET501 for intentional default return None in RerankFunction._get_object

* style: ruff format test_collection.py

* refact multi vector query

* format code

* fix(multi-vector): expose SubVectorQuery in Python binding, fix tests

- Register _SubVectorQuery in pybind11 with from_vector_query() factory
- Convert _VectorQuery to _SubVectorQuery in MultiVectorQueryExecutor
- Relax RRF/Weighted score assertion tolerance from 1e-10 to 1e-6
- Fix WeightedReRanker test metric to IP (matching HnswIndexParam default)

* style: ruff format query_executor.py

* fix: define _USE_MATH_DEFINES for M_PI on Windows (MSVC)

* refactor: include reranker.h directly in query.h instead of forward declaration

* refactor(reranker): move topn from member variable to rerank() parameter

* refact code

* style(python): fix ruff UP035/UP037 in multi_vector_reranker

- import Callable from collections.abc instead of typing (UP035)
- remove redundant quotes around MetricType annotations (UP037)

* chore: trigger PR sync

* refact code

* fix(examples): restore CMakeLists.txt formatting broken by clang-format

* refactor(reranker): remove redundant metrics_ map by querying schema directly, and use insert return value to avoid duplicate set lookup

* refactor(reranker): defer schema binding to query time and remove C API callback reranker
This commit is contained in:
lichen2015
2026-05-29 10:10:41 +08:00
committed by GitHub
parent e0ba23179b
commit f539580138
30 changed files with 2740 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/detail/test_collection_dql.py
+4 -5
View File
@@ -843,7 +843,7 @@ class TestCollectionQuery:
)
expected_score = expected_rrf_scores[doc_id]
actual_score = doc.score
assert abs(actual_score - expected_score) < 1e-10, (
assert abs(actual_score - expected_score) < 1e-6, (
f"RRF score mismatch for document {doc_id}: expected {expected_score}, got {actual_score}"
)
assert doc.score <= prev_score, (
@@ -876,9 +876,8 @@ class TestCollectionQuery:
batchdoc_and_check(full_collection, multiple_docs, doc_num, operator="insert")
doc_fields, doc_vectors = generate_vectordict_random(full_collection.schema)
weighted_reranker = WeightedReRanker(
topn=3, weights=weights, metric=MetricType.IP
)
metrics = {field: MetricType.IP for field in weights}
weighted_reranker = WeightedReRanker(topn=3, weights=weights, metrics=metrics)
single_query_results = {}
for k, v in DEFAULT_VECTOR_FIELD_NAME.items():
@@ -911,7 +910,7 @@ class TestCollectionQuery:
)
expected_score = expected_weighted_scores[doc_id]
actual_score = doc.score
assert abs(actual_score - expected_score) < 1e-10, (
assert abs(actual_score - expected_score) < 1e-6, (
f"score mismatch for document {doc_id}: expected {expected_score}, got {actual_score}"
)
assert doc.score <= prev_score, (
python/tests/test_collection.py
+249 -29
View File
@@ -27,11 +27,17 @@ from zvec import (
InvertIndexParam,
LogLevel,
LogType,
MetricType,
OptimizeOption,
StatusCode,
Query,
VectorSchema,
)
from zvec.extension.multi_vector_reranker import (
CallbackReRanker,
RrfReRanker,
WeightedReRanker,
)
# ==================== Common ====================
@@ -60,9 +66,18 @@ def collection_schema():
dimension=128,
index_param=HnswIndexParam(),
),
VectorSchema(
"dense2",
DataType.VECTOR_FP32,
dimension=128,
index_param=HnswIndexParam(),
),
VectorSchema(
"sparse", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
),
VectorSchema(
"sparse2", DataType.SPARSE_VECTOR_FP32, index_param=HnswIndexParam()
),
],
)
@@ -78,7 +93,12 @@ def single_doc():
return Doc(
id=f"{id}",
fields={"id": id, "name": "test", "weight": 80.0, "height": id + 140},
vectors={"dense": [id + 0.1] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}},
vectors={
"dense": [id + 0.1] * 128,
"dense2": [id + 0.2] * 128,
"sparse": {1: 1.0, 2: 2.0, 3: 3.0},
"sparse2": {4: 1.5, 5: 2.5, 6: 3.5},
},
)
@@ -88,7 +108,12 @@ def multiple_docs():
Doc(
id=f"{id}",
fields={"id": id, "name": "test", "weight": 80.0, "height": 210},
vectors={"dense": [id + 0.1] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}},
vectors={
"dense": [id + 0.1] * 128,
"dense2": [id + 0.2] * 128,
"sparse": {1: 1.0, 2: 2.0, 3: 3.0},
"sparse2": {4: 1.5, 5: 2.5, 6: 3.5},
},
)
for id in range(1, 101)
]
@@ -182,9 +207,11 @@ class TestCollectionDDL:
assert test_collection.stats is not None
stats = test_collection.stats
assert stats.doc_count == 0
assert len(stats.index_completeness) == 2
assert len(stats.index_completeness) == 4
assert stats.index_completeness["dense"] == 1
assert stats.index_completeness["dense2"] == 1
assert stats.index_completeness["sparse"] == 1
assert stats.index_completeness["sparse2"] == 1
# ----------------------------
@@ -449,7 +476,12 @@ class TestCollectionInsert:
"id": 1,
"name": "test",
},
vectors={"dense": [1 + 0.1] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}},
vectors={
"dense": [1 + 0.1] * 128,
"dense2": [1 + 0.2] * 128,
"sparse": {1: 1.0, 2: 2.0, 3: 3.0},
"sparse2": {4: 1.5, 5: 2.5, 6: 3.5},
},
)
result = test_collection.insert(doc)
assert bool(result)
@@ -465,7 +497,12 @@ class TestCollectionInsert:
# without id, name
doc = Doc(
id="0",
vectors={"dense": [1 + 0.1] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}},
vectors={
"dense": [1 + 0.1] * 128,
"dense2": [1 + 0.2] * 128,
"sparse": {1: 1.0, 2: 2.0, 3: 3.0},
"sparse2": {4: 1.5, 5: 2.5, 6: 3.5},
},
)
with pytest.raises(ValueError) as e:
# ValueError: Invalid doc: field[id] is required but not provided
@@ -478,7 +515,12 @@ class TestCollectionInsert:
fields={
"id": 1,
},
vectors={"dense": [1 + 0.1] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}},
vectors={
"dense": [1 + 0.1] * 128,
"dense2": [1 + 0.2] * 128,
"sparse": {1: 1.0, 2: 2.0, 3: 3.0},
"sparse2": {4: 1.5, 5: 2.5, 6: 3.5},
},
)
with pytest.raises(ValueError) as e:
test_collection.insert(doc)
@@ -494,7 +536,12 @@ class TestCollectionInsert:
"id": 1,
"name": "test",
},
vectors={"dense": [1 + 0.1] * 128, "sparse": {1: 1.0, 2: 2.0, 3: 3.0}},
vectors={
"dense": [1 + 0.1] * 128,
"dense2": [1 + 0.2] * 128,
"sparse": {1: 1.0, 2: 2.0, 3: 3.0},
"sparse2": {4: 1.5, 5: 2.5, 6: 3.5},
},
)
result = test_collection.insert(doc)
assert bool(result)
@@ -969,70 +1016,243 @@ class TestCollectionQuery:
def test_collection_query_multi_vector_with_same_field(
self, collection_with_multiple_docs: Collection, multiple_docs
):
with pytest.raises(ValueError):
# Multi-vector query on same field without reranker should raise ValueError
with pytest.raises(ValueError, match="Reranker is required"):
collection_with_multiple_docs.query(
[
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
Query(field_name="dense", vector=multiple_docs[1].vector("dense")),
]
)
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
# Same field name with reranker should also raise ValueError
reranker = RrfReRanker(topn=10, rank_constant=60)
with pytest.raises(ValueError, match="appears more than once"):
collection_with_multiple_docs.query(
[
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
Query(field_name="dense", vector=multiple_docs[1].vector("dense")),
],
topk=10,
reranker=reranker,
)
def test_collection_query_by_dense_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
result = collection_with_multiple_docs.query(
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
topk=10,
)
assert len(result) > 0
assert len(result) <= 10
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
def test_collection_query_by_sparse_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
result = collection_with_multiple_docs.query(
Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")),
topk=10,
)
assert len(result) > 0
assert len(result) <= 10
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
def test_collection_query_by_dense_vector_with_filter(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
result = collection_with_multiple_docs.query(
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
topk=10,
filter="id > 50",
)
assert len(result) > 0
assert len(result) <= 10
for doc in result:
assert int(doc.id) > 50
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
def test_collection_query_by_sparse_vector_with_filter(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
result = collection_with_multiple_docs.query(
Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")),
topk=10,
filter="id > 50",
)
assert len(result) > 0
assert len(result) <= 10
for doc in result:
assert int(doc.id) > 50
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
def test_collection_query_with_rrf_reranker_by_multi_dense_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
"""Test multi-vector query with RRF reranker on multiple dense vectors."""
reranker = RrfReRanker(topn=10, rank_constant=60)
result = collection_with_multiple_docs.query(
[
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
Query(field_name="dense2", vector=multiple_docs[0].vector("dense2")),
],
topk=10,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 10
# Results should have RRF-fused scores
for doc in result:
assert hasattr(doc, "score")
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
def test_collection_query_with_rrf_reranker_by_multi_sparse_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
"""Test multi-vector query with RRF reranker on multiple sparse vectors."""
reranker = RrfReRanker(topn=10, rank_constant=60)
result = collection_with_multiple_docs.query(
[
Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")),
Query(
field_name="sparse2",
vector=multiple_docs[0].vector("sparse2"),
),
],
topk=10,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 10
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
def test_collection_query_with_rrf_reranker_by_hybrid_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
"""Test multi-vector query with RRF reranker combining dense + sparse."""
reranker = RrfReRanker(topn=10, rank_constant=60)
result = collection_with_multiple_docs.query(
[
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")),
],
topk=10,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 10
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
def test_collection_query_with_weighted_reranker_by_multi_dense_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
"""Test multi-vector query with Weighted reranker on multiple dense vectors."""
metrics = {"dense": MetricType.IP, "dense2": MetricType.IP}
weights = {"dense": 0.6, "dense2": 0.4}
reranker = WeightedReRanker(topn=10, metrics=metrics, weights=weights)
result = collection_with_multiple_docs.query(
[
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
Query(field_name="dense2", vector=multiple_docs[0].vector("dense2")),
],
topk=10,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 10
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
def test_collection_query_with_weighted_reranker_by_multi_sparse_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
"""Test multi-vector query with Weighted reranker on multiple sparse vectors."""
metrics = {"sparse": MetricType.IP, "sparse2": MetricType.IP}
weights = {"sparse": 0.6, "sparse2": 0.4}
reranker = WeightedReRanker(topn=10, metrics=metrics, weights=weights)
result = collection_with_multiple_docs.query(
[
Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")),
Query(
field_name="sparse2",
vector=multiple_docs[0].vector("sparse2"),
),
],
topk=10,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 10
@pytest.mark.skip(reason="TODO: This test case is pending implementation")
def test_collection_query_with_weighted_reranker_by_hybrid_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
pass
"""Test multi-vector query with Weighted reranker combining dense + sparse."""
metrics = {"dense": MetricType.IP, "sparse": MetricType.IP}
weights = {"dense": 0.7, "sparse": 0.3}
reranker = WeightedReRanker(topn=10, metrics=metrics, weights=weights)
result = collection_with_multiple_docs.query(
[
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")),
],
topk=10,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 10
def test_collection_query_with_callback_reranker_by_multi_dense_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
"""Test multi-vector query with CallbackReRanker (Python callback via C++)."""
callback_invoked = []
def my_rerank_callback(query_results, topn):
callback_invoked.append(True)
all_docs = []
for docs in query_results.values():
all_docs.extend(docs)
seen = set()
unique_docs = []
for doc in all_docs:
if doc.pk() not in seen:
seen.add(doc.pk())
unique_docs.append(doc)
unique_docs.sort(key=lambda d: d.score(), reverse=True)
return unique_docs[:topn]
reranker = CallbackReRanker(callback=my_rerank_callback, topn=10)
result = collection_with_multiple_docs.query(
[
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
Query(field_name="dense2", vector=multiple_docs[0].vector("dense2")),
],
topk=10,
reranker=reranker,
)
assert len(callback_invoked) == 1
assert len(result) > 0
assert len(result) <= 10
def test_collection_query_with_callback_reranker_by_hybrid_vector(
self, collection_with_multiple_docs: Collection, multiple_docs
):
"""Test multi-vector query with CallbackReRanker combining dense + sparse."""
def my_rerank_callback(query_results, topn):
all_docs = []
for docs in query_results.values():
all_docs.extend(docs)
seen = set()
unique_docs = []
for doc in all_docs:
if doc.pk() not in seen:
seen.add(doc.pk())
unique_docs.append(doc)
unique_docs.sort(key=lambda d: d.score(), reverse=True)
return unique_docs[:topn]
reranker = CallbackReRanker(callback=my_rerank_callback, topn=5)
result = collection_with_multiple_docs.query(
[
Query(field_name="dense", vector=multiple_docs[0].vector("dense")),
Query(field_name="sparse", vector=multiple_docs[0].vector("sparse")),
],
topk=5,
reranker=reranker,
)
assert len(result) > 0
assert len(result) <= 5
python/tests/test_query_executor.py
+59
View File
@@ -31,13 +31,16 @@ from zvec.executor.query_executor import (
)
from zvec import (
RrfReRanker,
WeightedReRanker,
HnswQueryParam,
CollectionSchema,
VectorSchema,
DataType,
MetricType,
Query,
VectorQuery,
)
from zvec.extension.multi_vector_reranker import CallbackReRanker
# ----------------------------
@@ -209,6 +212,37 @@ class TestQueryContext:
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(
topn=10,
metrics={"test": MetricType.L2},
weights={"test": 1.0},
)
ctx = QueryContext(
topk=5,
queries=queries,
reranker=reranker,
)
assert ctx.reranker == reranker
assert ctx.reranker.weights == {"test": 1.0}
assert ctx.reranker.metrics == {"test": MetricType.L2}
def test_properties_with_callback_reranker(self):
queries = [Query(field_name="test")]
cb = lambda query_results, topn: []
reranker = CallbackReRanker(callback=cb, topn=10)
ctx = QueryContext(
topk=5,
queries=queries,
reranker=reranker,
)
assert ctx.reranker == reranker
def test_core_vectors_setter(self):
ctx = QueryContext(topk=10)
core_vectors = [MagicMock()]
@@ -304,6 +338,31 @@ class TestMultiVectorQueryExecutor:
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},
)
ctx = QueryContext(topk=10, queries=queries, reranker=reranker)
executor._do_validate(ctx)
def test_do_validate_multiple_queries_with_callback_reranker(self):
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._do_validate(ctx)
class TestQueryExecutorFactory:
def test_create_no_vectors(self):
python/tests/test_reranker.py
+67 -3
View File
@@ -20,6 +20,7 @@ import os
from zvec import Doc, MetricType
from zvec.extension.multi_vector_reranker import (
CallbackReRanker,
RrfReRanker,
WeightedReRanker,
)
@@ -75,16 +76,17 @@ class TestRrfReRanker:
# ----------------------------
class TestWeightedReRanker:
def test_init(self):
metrics = {"vector1": MetricType.L2, "vector2": MetricType.COSINE}
weights = {"vector1": 0.7, "vector2": 0.3}
reranker = WeightedReRanker(
topn=5,
rerank_field="content",
metric=MetricType.L2,
metrics=metrics,
weights=weights,
)
assert reranker.topn == 5
assert reranker.rerank_field == "content"
assert reranker.metric == MetricType.L2
assert reranker.metrics == metrics
assert reranker.weights == weights
def test_normalize_score(self):
@@ -106,8 +108,9 @@ class TestWeightedReRanker:
reranker._normalize_score(1.0, "unsupported_metric")
def test_rerank(self):
metrics = {"vector1": MetricType.L2, "vector2": MetricType.L2}
weights = {"vector1": 0.7, "vector2": 0.3}
reranker = WeightedReRanker(topn=3, weights=weights, metric=MetricType.L2)
reranker = WeightedReRanker(topn=3, weights=weights, metrics=metrics)
doc1 = Doc(id="1", score=0.8)
doc2 = Doc(id="2", score=0.7)
@@ -122,9 +125,70 @@ class TestWeightedReRanker:
for doc in results:
assert hasattr(doc, "score")
def test_rerank_missing_metric_raises(self):
metrics = {"vector1": MetricType.L2}
reranker = WeightedReRanker(topn=3, metrics=metrics)
doc1 = Doc(id="1", score=0.8)
query_results = {"vector1": [doc1], "vector2": [doc1]}
with pytest.raises(ValueError, match="no metric type specified"):
reranker.rerank(query_results)
# ----------------------------
# CallbackReRanker Test Case
# ----------------------------
class TestCallbackReRanker:
def test_init(self):
def my_callback(query_results, topn):
return []
reranker = CallbackReRanker(callback=my_callback, topn=5)
assert reranker.topn == 5
def test_rerank(self):
def my_callback(query_results, topn):
all_docs = []
for docs in query_results.values():
all_docs.extend(docs)
all_docs.sort(key=lambda d: d.score, reverse=True)
return all_docs[:topn]
reranker = CallbackReRanker(callback=my_callback, topn=3)
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 = {"vector1": [doc1, doc2], "vector2": [doc3, doc4]}
results = reranker.rerank(query_results)
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, topn):
received_topn.append(topn)
return []
reranker = CallbackReRanker(callback=my_callback, topn=7)
reranker.rerank({"v1": [Doc(id="1", score=0.5)]})
assert received_topn == [7]
def test_get_object_returns_cpp_reranker(self):
def my_callback(query_results, topn):
return []
reranker = CallbackReRanker(callback=my_callback)
assert reranker._get_object() is not None
# ----------------------------
# QwenReRanker Test Case