From 439dd10f5ec080922a370c707129ebe7c695171a Mon Sep 17 00:00:00 2001 From: egolearner Date: Wed, 3 Jun 2026 21:17:49 +0800 Subject: [PATCH] feat(query): support FTS + vector hybrid retrieval in MultiQuery (#459) Previously MultiQuery only accepted vector sub-queries by using get_vector_field() to look up each sub-query's field. FTS sub-queries (whose field_name points to an FTS-indexed string column) would fail with "Vector field not found". Changes: - collection.cc: use get_field() uniformly in MultiQuery path; let validate_and_sanitize() check type compatibility internally, which is consistent with the single-query path. - query_executor.py: allow SingleVectorQueryExecutor to accept multi-query when it contains an FTS query (with reranker), and route to C++ MultiQuery fast path. - Add test_collection_fts_vector_hybrid.py covering hybrid retrieval ranking, scoring, filter, validation, and edge cases. --- .../test_collection_fts_vector_hybrid.py | 399 ++++++++++++++++++ python/tests/test_query_executor.py | 4 +- python/zvec/executor/query_executor.py | 62 +-- src/db/collection.cc | 9 +- src/db/reranker/reranker.cc | 9 +- 5 files changed, 445 insertions(+), 38 deletions(-) create mode 100644 python/tests/test_collection_fts_vector_hybrid.py diff --git a/python/tests/test_collection_fts_vector_hybrid.py b/python/tests/test_collection_fts_vector_hybrid.py new file mode 100644 index 0000000..f57a727 --- /dev/null +++ b/python/tests/test_collection_fts_vector_hybrid.py @@ -0,0 +1,399 @@ +# 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 +from zvec.typing import MetricType + + +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(topn=10, 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(topn=10, 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(topn=10, 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(topn=10, 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(topn=10, 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(topn=10, 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_rejected( + self, hybrid_collection_with_docs: Collection + ): + """Multi-query with duplicate field names should raise.""" + reranker = RrfReRanker(topn=10, rank_constant=60) + with pytest.raises(ValueError, match="appears more than once"): + hybrid_collection_with_docs.query( + queries=[ + Query(field_name="content", fts=Fts(match_string="hello")), + Query(field_name="content", fts=Fts(match_string="world")), + ], + topk=5, + reranker=reranker, + ) + + def test_multiple_vectors_without_fts_rejected( + self, hybrid_collection_with_docs: Collection + ): + """Two vector queries on a single-vector-field collection should raise.""" + reranker = RrfReRanker(topn=10, rank_constant=60) + with pytest.raises(ValueError, match="cannot query with multiple vectors"): + 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, + ) + + +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.""" + metrics = {"embedding": MetricType.IP} + weights = {"content": 0.5, "embedding": 0.5} + reranker = WeightedReRanker(topn=10, metrics=metrics, 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.""" + metrics = {"embedding": MetricType.IP} + weights = {"content": 0.4, "embedding": 0.6} + reranker = WeightedReRanker(topn=10, metrics=metrics, 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 + metrics = {"embedding": MetricType.IP} + weights_fts_heavy = {"content": 0.9, "embedding": 0.1} + reranker_fts = WeightedReRanker( + topn=10, metrics=metrics, 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 = {"content": 0.1, "embedding": 0.9} + reranker_vec = WeightedReRanker( + topn=10, metrics=metrics, 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 diff --git a/python/tests/test_query_executor.py b/python/tests/test_query_executor.py index 857bcc8..6b2266e 100644 --- a/python/tests/test_query_executor.py +++ b/python/tests/test_query_executor.py @@ -326,9 +326,7 @@ class TestMultiVectorQueryExecutor: 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-vector query" - ): + with pytest.raises(ValueError, match="Reranker is required for multi-query"): executor._do_validate(ctx) def test_do_validate_multiple_queries_with_reranker(self): diff --git a/python/zvec/executor/query_executor.py b/python/zvec/executor/query_executor.py index 9c39d23..663ebc5 100644 --- a/python/zvec/executor/query_executor.py +++ b/python/zvec/executor/query_executor.py @@ -281,11 +281,28 @@ class SingleVectorQueryExecutor(NoVectorQueryExecutor): def __init__(self, schema: CollectionSchema) -> None: super().__init__(schema) + def _validate_multi_query(self, ctx: QueryContext) -> None: + """Shared validation for multi-query: reranker required + no duplicate fields.""" + if ctx.reranker is None: + raise ValueError("Reranker is required for multi-query") + seen_fields = set() + for query in ctx.queries: + query._validate() + if query.field_name in seen_fields: + raise ValueError( + f"Query field name '{query.field_name}' appears more than once" + ) + seen_fields.add(query.field_name) + def _do_validate(self, ctx: QueryContext) -> None: if len(ctx.queries) > 1: - raise ValueError( - "Collection has only one vector field, cannot query with multiple vectors" - ) + # Allow FTS + vector hybrid multi-query (requires reranker) + if not any(q.has_fts() for q in ctx.queries): + raise ValueError( + "Collection has only one vector field, cannot query with multiple vectors" + ) + self._validate_multi_query(ctx) + return for query in ctx.queries: query._validate() @@ -299,22 +316,6 @@ class SingleVectorQueryExecutor(NoVectorQueryExecutor): vectors.append(self._do_build_query_with_vector(ctx, query, collection)) return vectors - -class MultiVectorQueryExecutor(SingleVectorQueryExecutor): - def __init__(self, schema: CollectionSchema) -> None: - super().__init__(schema) - - def _do_validate(self, ctx: QueryContext) -> None: - if len(ctx.queries) > 1 and ctx.reranker is None: - raise ValueError("Reranker is required for multi-vector query") - seen_fields = set() - for query in ctx.queries: - query._validate() - field = query.field_name - if field in seen_fields: - raise ValueError(f"Query field name '{field}' appears more than once") - seen_fields.add(field) - def execute(self, ctx: QueryContext, collection: _Collection) -> list[Doc]: # 1. validate query self._do_validate(ctx) @@ -323,12 +324,12 @@ class MultiVectorQueryExecutor(SingleVectorQueryExecutor): if not query_vectors: raise ValueError("No query to execute") - # Fast path: use C++ MultiQuery for multi-vector with C++ reranker + # Multi-query fast path: route FTS + vector hybrid to C++ MultiQuery if len(query_vectors) > 1 and ctx.reranker is not None: cpp_reranker = ctx.reranker._get_object() if cpp_reranker is not None: mvq = _MultiQuery() - mvq.queries = [self._to_sub_query(vq) for vq in query_vectors] + mvq.queries = [_SubQuery.from_vector_query(vq) for vq in query_vectors] mvq.topk = ctx.topk if ctx.filter: mvq.filter = ctx.filter @@ -339,19 +340,22 @@ class MultiVectorQueryExecutor(SingleVectorQueryExecutor): docs = collection.Query(mvq) return [convert_to_py_doc(doc, self._schema) for doc in docs] - # 3. execute query (fallback to Python path) + # 3. execute query docs = self._do_execute(query_vectors, collection) # 4. merge and rerank result return self._do_merge_rerank_results(ctx, docs) - def _do_execute( - self, vectors: list[_VectorQuery], collection: _Collection - ) -> dict[str, list[Doc]]: - return super()._do_execute(vectors, collection) - @staticmethod - def _to_sub_query(vq: _VectorQuery) -> _SubQuery: - return _SubQuery.from_vector_query(vq) +class MultiVectorQueryExecutor(SingleVectorQueryExecutor): + def __init__(self, schema: CollectionSchema) -> None: + super().__init__(schema) + + def _do_validate(self, ctx: QueryContext) -> None: + if len(ctx.queries) > 1: + self._validate_multi_query(ctx) + return + for query in ctx.queries: + query._validate() class QueryExecutorFactory: diff --git a/src/db/collection.cc b/src/db/collection.cc index b25d18b..1b9dd69 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -1721,12 +1721,13 @@ Result CollectionImpl::Query(const MultiQuery &query) const { auto [_, inserted] = seen_fields.insert(target.field_name_); if (!inserted) { return tl::make_unexpected(Status::InvalidArgument( - "Duplicate field name in multi-vector query: ", target.field_name_)); + "Duplicate field name in multi-query: ", target.field_name_)); } - auto *field_schema = schema_->get_vector_field(target.field_name_); + // Use get_field uniformly; validate_and_sanitize checks type compatibility. + auto *field_schema = schema_->get_field(target.field_name_); if (!field_schema) { - return tl::make_unexpected(Status::InvalidArgument( - "Vector field not found: ", target.field_name_)); + return tl::make_unexpected( + Status::InvalidArgument("Field not found: ", target.field_name_)); } SearchQuery sq; diff --git a/src/db/reranker/reranker.cc b/src/db/reranker/reranker.cc index 252e88f..7f55ef2 100644 --- a/src/db/reranker/reranker.cc +++ b/src/db/reranker/reranker.cc @@ -95,6 +95,11 @@ void WeightedReranker::bind_schema(CollectionSchema::Ptr schema) { Result WeightedReranker::normalize_score(double score, const FieldSchema &field) { + // FTS field: BM25 scores are non-negative; normalize via arctan to [0, 1). + if (field.index_type() == IndexType::FTS) { + return 2.0 * std::atan(score) / M_PI; + } + auto *vip = dynamic_cast(field.index_params().get()); if (!vip) { @@ -118,10 +123,10 @@ Result WeightedReranker::normalize_score(double score, Result WeightedReranker::rescore(double score, int /*rank*/, const std::string &field_name) const { - const auto *field = schema_->get_vector_field(field_name); + const auto *field = schema_->get_field(field_name); if (!field) { return tl::make_unexpected(Status::InvalidArgument( - "WeightedReranker: vector field not found: '", field_name + "'")); + "WeightedReranker: field not found: '", field_name + "'")); } auto normalized = normalize_score(score, *field); if (!normalized.has_value()) {