From 518b6948818276050232774f2293ea4f4dd68be7 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 3 Aug 2026 17:20:28 +0800 Subject: [PATCH 1/6] reduce-overhead --- python/tests/test_batch_materialize.py | 180 ++++++++++++ python/tests/test_query_executor.py | 11 +- python/zvec/executor/query_executor.py | 17 +- src/binding/python/include/python_doc.h | 6 + src/binding/python/model/python_collection.cc | 58 +++- src/binding/python/model/python_doc.cc | 273 +++++++++--------- 6 files changed, 380 insertions(+), 165 deletions(-) create mode 100644 python/tests/test_batch_materialize.py diff --git a/python/tests/test_batch_materialize.py b/python/tests/test_batch_materialize.py new file mode 100644 index 000000000..989cbe3bc --- /dev/null +++ b/python/tests/test_batch_materialize.py @@ -0,0 +1,180 @@ +# 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. +"""Correctness tests for the batch-materialized query path. + +`Collection.query` goes through `_Collection.Query`, which batch-materializes +all hits into tuples in a single C++ call. These tests validate the +materialized output against two independent references: + +- a numpy brute-force ground truth over the inserted vectors (ids / scores); +- the `fetch` path, which materializes docs through a separate binding. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import zvec +from zvec import ( + Collection, + CollectionOption, + DataType, + Doc, + FieldSchema, + HnswIndexParam, + HnswQueryParam, + Query, + RrfReRanker, + VectorSchema, +) +from zvec.typing import MetricType + +DIM = 16 +N_DOCS = 200 + + +def _make_vectors() -> np.ndarray: + """Same deterministic vectors as inserted by the fixture.""" + return np.random.default_rng(42).random((N_DOCS, DIM), dtype=np.float32) + + +def _brute_force_topk( + query: np.ndarray, topk: int, mask: np.ndarray | None = None +) -> tuple[list[str], np.ndarray]: + """Exact L2sq top-k ids and distances over the ground-truth vectors.""" + dists = ((_make_vectors() - query) ** 2).sum(axis=1) + if mask is not None: + dists = np.where(mask, dists, np.inf) + idx = np.argsort(dists, kind="stable")[:topk] + return [str(i) for i in idx], dists[idx] + + +@pytest.fixture(scope="module") +def bm_collection(tmp_path_factory) -> Collection: + schema = zvec.CollectionSchema( + name="batch_mat_test", + fields=[ + FieldSchema("num", DataType.INT64, nullable=False), + FieldSchema("title", DataType.STRING, nullable=True), + ], + vectors=[ + VectorSchema( + "vec", + DataType.VECTOR_FP32, + dimension=DIM, + # explicit L2: score is the raw squared L2 distance (no + # metric normalization), matching the brute-force ground truth + index_param=HnswIndexParam(metric_type=MetricType.L2), + ), + ], + ) + path = tmp_path_factory.mktemp("zvec_batch_mat") / "coll" + coll = zvec.create_and_open( + path=str(path), + schema=schema, + option=CollectionOption(read_only=False, enable_mmap=True), + ) + + vectors = _make_vectors() + docs = [ + Doc( + id=str(i), + fields={"num": i, "title": f"doc-{i}"}, + vectors={"vec": vectors[i]}, + ) + for i in range(N_DOCS) + ] + for r in coll.insert(docs): + assert r.ok() + + yield coll + + try: + coll.destroy() + except Exception: + pass + + +class TestBatchMaterialize: + def _query_vec(self) -> np.ndarray: + return np.array([0.5] * DIM, dtype=np.float32) + + def test_matches_brute_force_ground_truth(self, bm_collection: Collection): + q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) + docs = bm_collection.query(q, topk=20) + assert len(docs) == 20 + + exp_ids, exp_dists = _brute_force_topk(self._query_vec(), 20) + assert [d.id for d in docs] == exp_ids + scores = [d.score for d in docs] + assert scores == sorted(scores) + for d, dist in zip(docs, exp_dists): + assert d.score == pytest.approx(float(dist), rel=1e-4) + + # scalar fields fully materialized, vectors excluded by default + for d in docs: + assert isinstance(d, Doc) + assert set(d.fields.keys()) == {"num", "title"} + assert d.fields["num"] == int(d.id) + assert d.fields["title"] == f"doc-{d.id}" + assert d.vectors == {} + + def test_fields_match_fetch_path(self, bm_collection: Collection): + q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) + docs = bm_collection.query(q, topk=10) + fetched = bm_collection.fetch([d.id for d in docs], include_vector=False) + for d in docs: + assert d.fields == fetched[d.id].fields + + @pytest.mark.parametrize("include_vector", [False, True]) + def test_include_vector(self, bm_collection: Collection, include_vector: bool): + q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) + docs = bm_collection.query(q, topk=5, include_vector=include_vector) + ground_truth = _make_vectors() + for d in docs: + assert bool(d.vectors) is include_vector + if include_vector: + assert np.allclose(d.vectors["vec"], ground_truth[int(d.id)]) + + def test_output_fields_subset(self, bm_collection: Collection): + q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) + docs = bm_collection.query(q, topk=5, output_fields=["num"]) + exp_ids, _ = _brute_force_topk(self._query_vec(), 5) + assert [d.id for d in docs] == exp_ids + for d in docs: + assert set(d.fields.keys()) == {"num"} + + def test_filter(self, bm_collection: Collection): + q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) + docs = bm_collection.query(q, topk=10, filter="num < 50") + mask = np.arange(N_DOCS) < 50 + exp_ids, exp_dists = _brute_force_topk(self._query_vec(), 10, mask) + assert [d.id for d in docs] == exp_ids + for d, dist in zip(docs, exp_dists): + assert d.fields["num"] < 50 + assert d.score == pytest.approx(float(dist), rel=1e-4) + + def test_empty_result(self, bm_collection: Collection): + q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) + docs = bm_collection.query(q, topk=10, filter="num < 0") + assert docs == [] + + def test_multi_query_rrf(self, bm_collection: Collection): + q1 = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) + q2 = Query(field_name="vec", vector=[0.1] * DIM, param=HnswQueryParam()) + docs = bm_collection.query([q1, q2], topk=10, reranker=RrfReRanker()) + assert len(docs) == 10 + for d in docs: + assert isinstance(d, Doc) + assert set(d.fields.keys()) == {"num", "title"} diff --git a/python/tests/test_query_executor.py b/python/tests/test_query_executor.py index 1bf084168..13dd16ffc 100644 --- a/python/tests/test_query_executor.py +++ b/python/tests/test_query_executor.py @@ -286,17 +286,14 @@ def test_do_merge_rerank_results_with_reranker(self): 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) + # Each query is executed serially and batch-materialized into results. + executor = QueryExecutor(MagicMock()) 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, - ): + with patch("zvec.executor.query_executor.Doc") as mock_doc: + mock_doc._from_tuple.side_effect = lambda t: t results = executor._execute_python_pipeline(vectors, collection) assert results == [["raw1"], ["raw2"]] assert collection.Query.call_count == 2 diff --git a/python/zvec/executor/query_executor.py b/python/zvec/executor/query_executor.py index cadbb17e0..344221b88 100644 --- a/python/zvec/executor/query_executor.py +++ b/python/zvec/executor/query_executor.py @@ -21,8 +21,7 @@ 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.doc import Doc, DocList from ..model.param.query import Query from ..model.schema import CollectionSchema from ..typing import DataType @@ -135,9 +134,13 @@ def execute(self, ctx: QueryContext, collection: _Collection) -> DocList: 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] + """Single/vector-less query: send a ``_SearchQuery`` to C++. + + Results are batch-materialized into tuples in a single C++ call, + avoiding per-doc Python/C++ crossings on the hot path. + """ + tuples = collection.Query(query, self._schema._get_object()) + return [Doc._from_tuple(t) if t is not None else None for t in tuples] def _execute_multi_query( self, ctx: QueryContext, queries: list[_SearchQuery], collection: _Collection @@ -160,8 +163,8 @@ def _execute_multi_query( 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] + tuples = collection.Query(multi_query, self._schema._get_object()) + return [Doc._from_tuple(t) if t is not None else None for t in tuples] def _build_multi_query( self, ctx: QueryContext, queries: list[_SearchQuery] diff --git a/src/binding/python/include/python_doc.h b/src/binding/python/include/python_doc.h index c38671057..6fdee18f6 100644 --- a/src/binding/python/include/python_doc.h +++ b/src/binding/python/include/python_doc.h @@ -14,6 +14,7 @@ #include #include +#include namespace py = pybind11; @@ -26,6 +27,11 @@ class ZVecPyDoc { public: static void Initialize(py::module_ &m); + // Materialize a single doc into (id, score, fields, vectors) following the + // collection schema. Shared by the per-doc `get_all` binding and the batch + // materialization path in the collection DQL bindings. Requires the GIL. + static py::tuple doc_to_tuple(Doc &self, const CollectionSchema &schema); + private: static void bind_doc_operator(py::module_ &m); static void bind_doc(py::module_ &m); diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 4468a53ff..0331e4268 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -15,9 +15,30 @@ #include "python_collection.h" #include #include +#include "python_doc.h" namespace zvec { +namespace { + +// Batch-materialize a DocPtrList into a list of (id, score, fields, vectors) +// tuples in a single GIL-held section, avoiding per-doc _Doc wrappers and +// per-doc Python->C++ crossings on the hot query path. +py::list docs_to_tuples(const DocPtrList &docs, + const CollectionSchema &schema) { + py::list out; + for (const auto &doc : docs) { + if (doc) { + out.append(ZVecPyDoc::doc_to_tuple(*doc, schema)); + } else { + out.append(py::none()); + } + } + return out; +} + +} // namespace + inline void throw_if_error(const Status &status) { switch (status.code()) { case StatusCode::OK: @@ -255,29 +276,38 @@ void ZVecPyCollection::bind_dml_methods( void ZVecPyCollection::bind_dql_methods( py::class_ &col) { - col.def("Query", - [](const Collection &self, const SearchQuery &query) { - Result result; - { - py::gil_scoped_release release; - result = self.Query(query); - } - // return DocPtrList - return unwrap_expected(result); - }) + // Query with the GIL released, then materialize all hits into + // (id, score, fields, vectors) tuples in one crossing (see docs_to_tuples). + col.def( + "Query", + [](const Collection &self, const SearchQuery &query, + const CollectionSchema &schema) { + Result result; + { + py::gil_scoped_release release; + result = self.Query(query); + } + return docs_to_tuples(unwrap_expected(result), schema); + }, + py::arg("query"), py::arg("schema"), + "Execute a query and return results as a list of " + "(id, score, fields, vectors) tuples materialized in one batch.") // MultiQuery: multi query with reranker .def( "Query", - [](const Collection &self, const MultiQuery &query) { + [](const Collection &self, const MultiQuery &query, + const CollectionSchema &schema) { Result result; { py::gil_scoped_release release; result = self.Query(query); } - // return DocPtrList - return unwrap_expected(result); + return docs_to_tuples(unwrap_expected(result), schema); }, - py::arg("query"), "Execute a multi query with re-ranking.") + py::arg("query"), py::arg("schema"), + "Execute a multi query with re-ranking and return results as a " + "list of (id, score, fields, vectors) tuples materialized in one " + "batch.") .def("GroupByQuery", [](const Collection &self, const GroupByVectorQuery &query) { Result result; diff --git a/src/binding/python/model/python_doc.cc b/src/binding/python/model/python_doc.cc index 68031977d..57bd60fbd 100644 --- a/src/binding/python/model/python_doc.cc +++ b/src/binding/python/model/python_doc.cc @@ -303,150 +303,149 @@ void ZVecPyDoc::bind_doc(py::module_ &m) { throw py::type_error("Unsupported type for field: " + field); } }); - doc.def( - "get_all", - [](Doc &self, const CollectionSchema &schema) -> py::tuple { - py::tuple result(4); - // 1. set doc id and score - result[0] = py::str(self.pk()); - result[1] = py::float_(self.score()); + doc.def("get_all", &ZVecPyDoc::doc_to_tuple, py::arg("schema"), + "Get all fields and vectors as a tuple: (id, score, fields, " + "vectors). Vectors are zero-copy numpy arrays (dense: ndarray, " + "sparse: (indices, values) tuple)."); +} - if (self.is_empty()) { - result[2] = py::none(); - result[3] = py::none(); - return result; - } - // 2. set scalar fields - py::dict fields; - for (const auto &field_meta : schema.forward_fields()) { - const std::string &field = field_meta->name(); - if (!self.has_value(field)) { - continue; - } +py::tuple ZVecPyDoc::doc_to_tuple(Doc &self, const CollectionSchema &schema) { + py::tuple result(4); + // 1. set doc id and score + result[0] = py::str(self.pk()); + result[1] = py::float_(self.score()); - try { - auto val = [&]() -> py::object { - switch (field_meta->data_type()) { - // base datatypes - case DataType::STRING: - return py::str(self.get(field).value()); - case DataType::BOOL: - return py::cast(self.get(field)); - case DataType::INT32: - return py::cast(self.get(field)); - case DataType::UINT32: - return py::cast(self.get(field)); - case DataType::INT64: - return py::cast(self.get(field)); - case DataType::UINT64: - return py::cast(self.get(field)); - case DataType::FLOAT: - return py::cast(self.get(field)); - case DataType::DOUBLE: - return py::cast(self.get(field)); + if (self.is_empty()) { + result[2] = py::none(); + result[3] = py::none(); + return result; + } + // 2. set scalar fields + py::dict fields; + for (const auto &field_meta : schema.forward_fields()) { + const std::string &field = field_meta->name(); + if (!self.has_value(field)) { + continue; + } - // array datatypes - case DataType::ARRAY_STRING: - return py::cast(self.get>(field)); - case DataType::ARRAY_INT32: - return py::cast(self.get>(field)); - case DataType::ARRAY_INT64: - return py::cast(self.get>(field)); - case DataType::ARRAY_UINT32: - return py::cast(self.get>(field)); - case DataType::ARRAY_UINT64: - return py::cast(self.get>(field)); - case DataType::ARRAY_FLOAT: - return py::cast(self.get>(field)); - case DataType::ARRAY_DOUBLE: - return py::cast(self.get>(field)); - case DataType::ARRAY_BOOL: - return py::cast(self.get>(field)); - default: - throw py::type_error("Unsupported type for field: " + field); - } - }(); - fields[py::str(field)] = val; - } catch (const std::exception &e) { - fields[py::str(field)] = py::none(); - } - } - if (!fields.empty()) { - result[2] = fields; - } else { - result[2] = py::none(); + try { + auto val = [&]() -> py::object { + switch (field_meta->data_type()) { + // base datatypes + case DataType::STRING: + return py::str(self.get(field).value()); + case DataType::BOOL: + return py::cast(self.get(field)); + case DataType::INT32: + return py::cast(self.get(field)); + case DataType::UINT32: + return py::cast(self.get(field)); + case DataType::INT64: + return py::cast(self.get(field)); + case DataType::UINT64: + return py::cast(self.get(field)); + case DataType::FLOAT: + return py::cast(self.get(field)); + case DataType::DOUBLE: + return py::cast(self.get(field)); + + // array datatypes + case DataType::ARRAY_STRING: + return py::cast(self.get>(field)); + case DataType::ARRAY_INT32: + return py::cast(self.get>(field)); + case DataType::ARRAY_INT64: + return py::cast(self.get>(field)); + case DataType::ARRAY_UINT32: + return py::cast(self.get>(field)); + case DataType::ARRAY_UINT64: + return py::cast(self.get>(field)); + case DataType::ARRAY_FLOAT: + return py::cast(self.get>(field)); + case DataType::ARRAY_DOUBLE: + return py::cast(self.get>(field)); + case DataType::ARRAY_BOOL: + return py::cast(self.get>(field)); + default: + throw py::type_error("Unsupported type for field: " + field); } - // 3. set vector fields - py::dict vectors; - for (const auto &vec_meta : schema.vector_fields()) { - const std::string &vec = vec_meta->name(); - if (!self.has_value(vec)) continue; + }(); + fields[py::str(field)] = val; + } catch (const std::exception &e) { + fields[py::str(field)] = py::none(); + } + } + if (!fields.empty()) { + result[2] = fields; + } else { + result[2] = py::none(); + } + // 3. set vector fields + py::dict vectors; + for (const auto &vec_meta : schema.vector_fields()) { + const std::string &vec = vec_meta->name(); + if (!self.has_value(vec)) continue; - try { - auto array = [&]() -> py::object { - switch (vec_meta->data_type()) { - case DataType::VECTOR_INT8: - return py::cast(self.get>(vec)); - case DataType::VECTOR_FP16: { - auto value = self.get>(vec); - if (value.has_value()) { - std::vector new_value; - new_value.reserve(value.value().size()); - for (auto &item : value.value()) { - new_value.push_back(static_cast(item)); - } - return py::cast(new_value); - } - return py::none(); - } - case DataType::VECTOR_FP32: - return py::cast(self.get>(vec)); - case DataType::VECTOR_FP64: - return py::cast(self.get>(vec)); - case DataType::SPARSE_VECTOR_FP16: { - auto vector = - self.get, - std::vector>>(vec); - const auto &indices = vector->first; - const auto &values = vector->second; - py::dict d; - for (size_t i = 0; i < indices.size(); ++i) { - d[py::int_(indices[i])] = - py::float_(static_cast(values[i])); - } - return d; - } - case DataType::SPARSE_VECTOR_FP32: { - auto vector = self.get< - std::pair, std::vector>>( - vec); - const auto &indices = vector->first; - const auto &values = vector->second; - py::dict d; - for (size_t i = 0; i < indices.size(); ++i) { - d[py::int_(indices[i])] = py::float_(values[i]); - } - return d; - } - default: - throw py::type_error("Unsupported type for field: " + vec); + try { + auto array = [&]() -> py::object { + switch (vec_meta->data_type()) { + case DataType::VECTOR_INT8: + return py::cast(self.get>(vec)); + case DataType::VECTOR_FP16: { + auto value = self.get>(vec); + if (value.has_value()) { + std::vector new_value; + new_value.reserve(value.value().size()); + for (auto &item : value.value()) { + new_value.push_back(static_cast(item)); } - }(); - vectors[py::str(vec)] = array; - } catch (const std::exception &e) { - vectors[py::str(vec)] = py::none(); + return py::cast(new_value); + } + return py::none(); } + case DataType::VECTOR_FP32: + return py::cast(self.get>(vec)); + case DataType::VECTOR_FP64: + return py::cast(self.get>(vec)); + case DataType::SPARSE_VECTOR_FP16: { + auto vector = self.get< + std::pair, std::vector>>( + vec); + const auto &indices = vector->first; + const auto &values = vector->second; + py::dict d; + for (size_t i = 0; i < indices.size(); ++i) { + d[py::int_(indices[i])] = + py::float_(static_cast(values[i])); + } + return d; + } + case DataType::SPARSE_VECTOR_FP32: { + auto vector = + self.get, std::vector>>( + vec); + const auto &indices = vector->first; + const auto &values = vector->second; + py::dict d; + for (size_t i = 0; i < indices.size(); ++i) { + d[py::int_(indices[i])] = py::float_(values[i]); + } + return d; + } + default: + throw py::type_error("Unsupported type for field: " + vec); } - if (!vectors.empty()) { - result[3] = vectors; - } else { - result[3] = py::none(); - } - return result; - }, - py::arg("schema"), - "Get all fields and vectors as a tuple: (id, score, fields, vectors). " - "Vectors are zero-copy numpy arrays (dense: ndarray, sparse: (indices, " - "values) tuple)."); + }(); + vectors[py::str(vec)] = array; + } catch (const std::exception &e) { + vectors[py::str(vec)] = py::none(); + } + } + if (!vectors.empty()) { + result[3] = vectors; + } else { + result[3] = py::none(); + } + return result; } } // namespace zvec \ No newline at end of file From 6364069fd3f028747c43927db0bfd48430e0523f Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 4 Aug 2026 11:00:03 +0800 Subject: [PATCH 2/6] debug with copilot --- python/tests/test_query_executor.py | 40 ++++++++++++++++--- python/zvec/__init__.pyi | 4 +- python/zvec/executor/query_executor.py | 7 ++-- src/binding/python/model/python_collection.cc | 32 +++++++++------ 4 files changed, 61 insertions(+), 22 deletions(-) diff --git a/python/tests/test_query_executor.py b/python/tests/test_query_executor.py index 13dd16ffc..89fc48d10 100644 --- a/python/tests/test_query_executor.py +++ b/python/tests/test_query_executor.py @@ -286,17 +286,47 @@ def test_do_merge_rerank_results_with_reranker(self): reranker.rerank.assert_called_once_with(docs_list, ctx.topk) def test_execute_python_pipeline(self): - # Each query is executed serially and batch-materialized into results. + # Each query is executed serially and batch-materialized into results: + # Doc._from_tuple is invoked for every non-None tuple returned by + # collection.Query, while None entries are passed through untouched. executor = QueryExecutor(MagicMock()) collection = MagicMock() - collection.Query.side_effect = [["raw1"], ["raw2"]] + collection.Query.side_effect = [["raw1", None], ["raw2"]] vectors = [MagicMock(), MagicMock()] with patch("zvec.executor.query_executor.Doc") as mock_doc: - mock_doc._from_tuple.side_effect = lambda t: t + mock_doc._from_tuple.side_effect = lambda t: ("doc", t) results = executor._execute_python_pipeline(vectors, collection) - assert results == [["raw1"], ["raw2"]] - assert collection.Query.call_count == 2 + + assert collection.Query.call_args_list == [ + ((vectors[0],), {}), + ((vectors[1],), {}), + ] + assert mock_doc._from_tuple.call_args_list == [ + (("raw1",), {}), + (("raw2",), {}), + ] + assert results == [[("doc", "raw1"), None], [("doc", "raw2")]] + + def test_execute_single_query_batch_materializes(self): + # _execute_single_query sends the query as-is to collection.Query + # (the schema is resolved inside the C++ binding) and converts each + # non-None returned tuple via Doc._from_tuple, keeping order and None. + executor = QueryExecutor(MagicMock()) + collection = MagicMock() + collection.Query.return_value = ["raw1", None, "raw2"] + query = MagicMock() + + with patch("zvec.executor.query_executor.Doc") as mock_doc: + mock_doc._from_tuple.side_effect = lambda t: ("doc", t) + results = executor._execute_single_query(query, collection) + + collection.Query.assert_called_once_with(query) + assert mock_doc._from_tuple.call_args_list == [ + (("raw1",), {}), + (("raw2",), {}), + ] + assert results == [("doc", "raw1"), None, ("doc", "raw2")] def test_build_search_query_by_missing_id_raises_value_error(self): vector_schema = VectorSchema(name="test", data_type=DataType.VECTOR_FP32) diff --git a/python/zvec/__init__.pyi b/python/zvec/__init__.pyi index 09e828dae..c9529788e 100644 --- a/python/zvec/__init__.pyi +++ b/python/zvec/__init__.pyi @@ -160,7 +160,9 @@ class _Collection: 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 Query( + self, arg0: param._SearchQuery + ) -> list[tuple[str, float, dict | None, dict | None] | None]: ... def Schema(self) -> schema._CollectionSchema: ... def Stats(self) -> schema.CollectionStats: ... def Update(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ... diff --git a/python/zvec/executor/query_executor.py b/python/zvec/executor/query_executor.py index 344221b88..9837c5d72 100644 --- a/python/zvec/executor/query_executor.py +++ b/python/zvec/executor/query_executor.py @@ -136,10 +136,11 @@ def _execute_single_query( ) -> DocList: """Single/vector-less query: send a ``_SearchQuery`` to C++. - Results are batch-materialized into tuples in a single C++ call, + Results are batch-materialized into tuples in a single C++ call + (the schema is resolved inside the binding from the collection), avoiding per-doc Python/C++ crossings on the hot path. """ - tuples = collection.Query(query, self._schema._get_object()) + tuples = collection.Query(query) return [Doc._from_tuple(t) if t is not None else None for t in tuples] def _execute_multi_query( @@ -163,7 +164,7 @@ def _execute_multi_query( return self._merge_and_rerank(ctx, docs_list) multi_query = self._build_multi_query(ctx, queries) - tuples = collection.Query(multi_query, self._schema._get_object()) + tuples = collection.Query(multi_query) return [Doc._from_tuple(t) if t is not None else None for t in tuples] def _build_multi_query( diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 0331e4268..12a6c98c2 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -26,12 +26,12 @@ namespace { // per-doc Python->C++ crossings on the hot query path. py::list docs_to_tuples(const DocPtrList &docs, const CollectionSchema &schema) { - py::list out; - for (const auto &doc : docs) { - if (doc) { - out.append(ZVecPyDoc::doc_to_tuple(*doc, schema)); + py::list out(docs.size()); + for (size_t i = 0; i < docs.size(); ++i) { + if (docs[i]) { + out[i] = ZVecPyDoc::doc_to_tuple(*docs[i], schema); } else { - out.append(py::none()); + out[i] = py::none(); } } return out; @@ -278,33 +278,39 @@ void ZVecPyCollection::bind_dql_methods( py::class_ &col) { // Query with the GIL released, then materialize all hits into // (id, score, fields, vectors) tuples in one crossing (see docs_to_tuples). + // The schema is taken from the collection itself, keeping the signature + // unchanged from the legacy per-doc binding. col.def( "Query", - [](const Collection &self, const SearchQuery &query, - const CollectionSchema &schema) { + [](const Collection &self, const SearchQuery &query) { Result result; + Result schema_result; { py::gil_scoped_release release; result = self.Query(query); + schema_result = self.Schema(); } - return docs_to_tuples(unwrap_expected(result), schema); + return docs_to_tuples(unwrap_expected(result), + unwrap_expected(schema_result)); }, - py::arg("query"), py::arg("schema"), + py::arg("query"), "Execute a query and return results as a list of " "(id, score, fields, vectors) tuples materialized in one batch.") // MultiQuery: multi query with reranker .def( "Query", - [](const Collection &self, const MultiQuery &query, - const CollectionSchema &schema) { + [](const Collection &self, const MultiQuery &query) { Result result; + Result schema_result; { py::gil_scoped_release release; result = self.Query(query); + schema_result = self.Schema(); } - return docs_to_tuples(unwrap_expected(result), schema); + return docs_to_tuples(unwrap_expected(result), + unwrap_expected(schema_result)); }, - py::arg("query"), py::arg("schema"), + py::arg("query"), "Execute a multi query with re-ranking and return results as a " "list of (id, score, fields, vectors) tuples materialized in one " "batch.") From 790fc30a54ea429fc0cbf1ff2a03ea0ecef8440a Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 17 Aug 2026 15:29:23 +0800 Subject: [PATCH 3/6] add QueryWithSchema --- src/binding/python/include/python_doc.h | 6 + src/binding/python/model/python_collection.cc | 34 ++- src/binding/python/model/python_doc.cc | 279 ++++++------------ src/db/collection.cc | 60 +++- src/include/zvec/db/collection.h | 32 ++ tests/db/collection_test.cc | 120 ++++++++ 6 files changed, 331 insertions(+), 200 deletions(-) diff --git a/src/binding/python/include/python_doc.h b/src/binding/python/include/python_doc.h index 6fdee18f6..787f82eb9 100644 --- a/src/binding/python/include/python_doc.h +++ b/src/binding/python/include/python_doc.h @@ -32,6 +32,12 @@ class ZVecPyDoc { // materialization path in the collection DQL bindings. Requires the GIL. static py::tuple doc_to_tuple(Doc &self, const CollectionSchema &schema); + // Convert a single Doc field value into a Python object according to its + // DataType. Shared by the per-field `get_any` binding and `doc_to_tuple`. + // Requires the GIL. + static py::object doc_value_to_py(Doc &self, const std::string &field, + DataType type); + private: static void bind_doc_operator(py::module_ &m); static void bind_doc(py::module_ &m); diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 12a6c98c2..157ff162e 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -68,6 +68,15 @@ T unwrap_expected(const tl::expected &exp) { return T{}; } +template +T unwrap_expected(tl::expected &&exp) { + if (exp.has_value()) { + return std::move(exp).value(); + } + throw_if_error(exp.error()); + return T{}; +} + void ZVecPyCollection::Initialize(pybind11::module_ &m) { py::class_(m, "_GroupResult") .def_readonly("group_by_value", &GroupResult::group_by_value_) @@ -278,20 +287,19 @@ void ZVecPyCollection::bind_dql_methods( py::class_ &col) { // Query with the GIL released, then materialize all hits into // (id, score, fields, vectors) tuples in one crossing (see docs_to_tuples). - // The schema is taken from the collection itself, keeping the signature - // unchanged from the legacy per-doc binding. + // QueryWithSchema returns the docs and the schema snapshot atomically under + // one read lock, so concurrent DDL cannot desynchronize them, while the + // binding signature stays unchanged from the legacy per-doc binding. col.def( "Query", [](const Collection &self, const SearchQuery &query) { - Result result; - Result schema_result; + Result result; { py::gil_scoped_release release; - result = self.Query(query); - schema_result = self.Schema(); + result = self.QueryWithSchema(query); } - return docs_to_tuples(unwrap_expected(result), - unwrap_expected(schema_result)); + auto snapshot = unwrap_expected(std::move(result)); + return docs_to_tuples(snapshot.docs, snapshot.schema); }, py::arg("query"), "Execute a query and return results as a list of " @@ -300,15 +308,13 @@ void ZVecPyCollection::bind_dql_methods( .def( "Query", [](const Collection &self, const MultiQuery &query) { - Result result; - Result schema_result; + Result result; { py::gil_scoped_release release; - result = self.Query(query); - schema_result = self.Schema(); + result = self.QueryWithSchema(query); } - return docs_to_tuples(unwrap_expected(result), - unwrap_expected(schema_result)); + auto snapshot = unwrap_expected(std::move(result)); + return docs_to_tuples(snapshot.docs, snapshot.schema); }, py::arg("query"), "Execute a multi query with re-ranking and return results as a " diff --git a/src/binding/python/model/python_doc.cc b/src/binding/python/model/python_doc.cc index 57bd60fbd..1b180df02 100644 --- a/src/binding/python/model/python_doc.cc +++ b/src/binding/python/model/python_doc.cc @@ -214,101 +214,107 @@ void ZVecPyDoc::bind_doc(py::module_ &m) { }); // binding doc get field - doc.def( - "get_any", - [](Doc &self, const std::string &field, - const DataType &type) -> py::object { - switch (type) { - // base datatypes - case DataType::STRING: - return py::cast(self.get(field)); - case DataType::BOOL: - return py::cast(self.get(field)); - case DataType::INT32: - return py::cast(self.get(field)); - case DataType::UINT32: - return py::cast(self.get(field)); - case DataType::INT64: - return py::cast(self.get(field)); - case DataType::UINT64: - return py::cast(self.get(field)); - case DataType::FLOAT: - return py::cast(self.get(field)); - case DataType::DOUBLE: - return py::cast(self.get(field)); - - // array datatypes - case DataType::ARRAY_STRING: - return py::cast(self.get>(field)); - case DataType::ARRAY_INT32: - return py::cast(self.get>(field)); - case DataType::ARRAY_INT64: - return py::cast(self.get>(field)); - case DataType::ARRAY_UINT32: - return py::cast(self.get>(field)); - case DataType::ARRAY_UINT64: - return py::cast(self.get>(field)); - case DataType::ARRAY_FLOAT: - return py::cast(self.get>(field)); - case DataType::ARRAY_DOUBLE: - return py::cast(self.get>(field)); - case DataType::ARRAY_BOOL: - return py::cast(self.get>(field)); - - // vector datatypes - case DataType::VECTOR_INT8: - return py::cast(self.get>(field)); - case DataType::VECTOR_FP16: { - auto value = self.get>(field); - if (value.has_value()) { - std::vector new_value; - new_value.reserve(value.value().size()); - for (auto &item : value.value()) { - new_value.push_back(static_cast(item)); - } - return py::cast(new_value); - } - return py::none(); - } - case DataType::VECTOR_FP32: - return py::cast(self.get>(field)); - case DataType::VECTOR_FP64: - return py::cast(self.get>(field)); - case DataType::SPARSE_VECTOR_FP16: { - auto vector = self.get< - std::pair, std::vector>>( - field); - const auto &indices = vector->first; - const auto &values = vector->second; - py::dict d; - for (size_t i = 0; i < indices.size(); ++i) { - d[py::int_(indices[i])] = - py::float_(static_cast(values[i])); - } - return d; - } - case DataType::SPARSE_VECTOR_FP32: { - auto vector = - self.get, std::vector>>( - field); - const auto &indices = vector->first; - const auto &values = vector->second; - py::dict d; - for (size_t i = 0; i < indices.size(); ++i) { - d[py::int_(indices[i])] = py::float_(values[i]); - } - return d; - } - default: - throw py::type_error("Unsupported type for field: " + field); - } - }); + doc.def("get_any", + [](Doc &self, const std::string &field, const DataType &type) + -> py::object { return doc_value_to_py(self, field, type); }); doc.def("get_all", &ZVecPyDoc::doc_to_tuple, py::arg("schema"), "Get all fields and vectors as a tuple: (id, score, fields, " "vectors). Vectors are zero-copy numpy arrays (dense: ndarray, " "sparse: (indices, values) tuple)."); } +py::object ZVecPyDoc::doc_value_to_py(Doc &self, const std::string &field, + DataType type) { + switch (type) { + // base datatypes + case DataType::STRING: + return py::cast(self.get(field)); + case DataType::BOOL: + return py::cast(self.get(field)); + case DataType::INT32: + return py::cast(self.get(field)); + case DataType::UINT32: + return py::cast(self.get(field)); + case DataType::INT64: + return py::cast(self.get(field)); + case DataType::UINT64: + return py::cast(self.get(field)); + case DataType::FLOAT: + return py::cast(self.get(field)); + case DataType::DOUBLE: + return py::cast(self.get(field)); + + // array datatypes + case DataType::ARRAY_STRING: + return py::cast(self.get>(field)); + case DataType::ARRAY_INT32: + return py::cast(self.get>(field)); + case DataType::ARRAY_INT64: + return py::cast(self.get>(field)); + case DataType::ARRAY_UINT32: + return py::cast(self.get>(field)); + case DataType::ARRAY_UINT64: + return py::cast(self.get>(field)); + case DataType::ARRAY_FLOAT: + return py::cast(self.get>(field)); + case DataType::ARRAY_DOUBLE: + return py::cast(self.get>(field)); + case DataType::ARRAY_BOOL: + return py::cast(self.get>(field)); + + // vector datatypes + case DataType::VECTOR_INT8: + return py::cast(self.get>(field)); + case DataType::VECTOR_FP16: { + auto value = self.get>(field); + if (value.has_value()) { + std::vector new_value; + new_value.reserve(value.value().size()); + for (auto &item : value.value()) { + new_value.push_back(static_cast(item)); + } + return py::cast(new_value); + } + return py::none(); + } + case DataType::VECTOR_FP32: + return py::cast(self.get>(field)); + case DataType::VECTOR_FP64: + return py::cast(self.get>(field)); + case DataType::SPARSE_VECTOR_FP16: { + auto vector = self.get< + std::pair, std::vector>>( + field); + if (!vector.has_value()) { + return py::none(); + } + const auto &indices = vector->first; + const auto &values = vector->second; + py::dict d; + for (size_t i = 0; i < indices.size(); ++i) { + d[py::int_(indices[i])] = py::float_(static_cast(values[i])); + } + return d; + } + case DataType::SPARSE_VECTOR_FP32: { + auto vector = + self.get, std::vector>>(field); + if (!vector.has_value()) { + return py::none(); + } + const auto &indices = vector->first; + const auto &values = vector->second; + py::dict d; + for (size_t i = 0; i < indices.size(); ++i) { + d[py::int_(indices[i])] = py::float_(values[i]); + } + return d; + } + default: + throw py::type_error("Unsupported type for field: " + field); + } +} + py::tuple ZVecPyDoc::doc_to_tuple(Doc &self, const CollectionSchema &schema) { py::tuple result(4); // 1. set doc id and score @@ -329,48 +335,8 @@ py::tuple ZVecPyDoc::doc_to_tuple(Doc &self, const CollectionSchema &schema) { } try { - auto val = [&]() -> py::object { - switch (field_meta->data_type()) { - // base datatypes - case DataType::STRING: - return py::str(self.get(field).value()); - case DataType::BOOL: - return py::cast(self.get(field)); - case DataType::INT32: - return py::cast(self.get(field)); - case DataType::UINT32: - return py::cast(self.get(field)); - case DataType::INT64: - return py::cast(self.get(field)); - case DataType::UINT64: - return py::cast(self.get(field)); - case DataType::FLOAT: - return py::cast(self.get(field)); - case DataType::DOUBLE: - return py::cast(self.get(field)); - - // array datatypes - case DataType::ARRAY_STRING: - return py::cast(self.get>(field)); - case DataType::ARRAY_INT32: - return py::cast(self.get>(field)); - case DataType::ARRAY_INT64: - return py::cast(self.get>(field)); - case DataType::ARRAY_UINT32: - return py::cast(self.get>(field)); - case DataType::ARRAY_UINT64: - return py::cast(self.get>(field)); - case DataType::ARRAY_FLOAT: - return py::cast(self.get>(field)); - case DataType::ARRAY_DOUBLE: - return py::cast(self.get>(field)); - case DataType::ARRAY_BOOL: - return py::cast(self.get>(field)); - default: - throw py::type_error("Unsupported type for field: " + field); - } - }(); - fields[py::str(field)] = val; + fields[py::str(field)] = + doc_value_to_py(self, field, field_meta->data_type()); } catch (const std::exception &e) { fields[py::str(field)] = py::none(); } @@ -387,56 +353,7 @@ py::tuple ZVecPyDoc::doc_to_tuple(Doc &self, const CollectionSchema &schema) { if (!self.has_value(vec)) continue; try { - auto array = [&]() -> py::object { - switch (vec_meta->data_type()) { - case DataType::VECTOR_INT8: - return py::cast(self.get>(vec)); - case DataType::VECTOR_FP16: { - auto value = self.get>(vec); - if (value.has_value()) { - std::vector new_value; - new_value.reserve(value.value().size()); - for (auto &item : value.value()) { - new_value.push_back(static_cast(item)); - } - return py::cast(new_value); - } - return py::none(); - } - case DataType::VECTOR_FP32: - return py::cast(self.get>(vec)); - case DataType::VECTOR_FP64: - return py::cast(self.get>(vec)); - case DataType::SPARSE_VECTOR_FP16: { - auto vector = self.get< - std::pair, std::vector>>( - vec); - const auto &indices = vector->first; - const auto &values = vector->second; - py::dict d; - for (size_t i = 0; i < indices.size(); ++i) { - d[py::int_(indices[i])] = - py::float_(static_cast(values[i])); - } - return d; - } - case DataType::SPARSE_VECTOR_FP32: { - auto vector = - self.get, std::vector>>( - vec); - const auto &indices = vector->first; - const auto &values = vector->second; - py::dict d; - for (size_t i = 0; i < indices.size(); ++i) { - d[py::int_(indices[i])] = py::float_(values[i]); - } - return d; - } - default: - throw py::type_error("Unsupported type for field: " + vec); - } - }(); - vectors[py::str(vec)] = array; + vectors[py::str(vec)] = doc_value_to_py(self, vec, vec_meta->data_type()); } catch (const std::exception &e) { vectors[py::str(vec)] = py::none(); } diff --git a/src/db/collection.cc b/src/db/collection.cc index 165e8b717..b6ca5c7e5 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -123,6 +123,11 @@ class CollectionImpl : public Collection { Result Query(const MultiQuery &query) const override; + Result QueryWithSchema( + const SearchQuery &query) const override; + + Result QueryWithSchema(const MultiQuery &query) const override; + Result GroupByQuery( const GroupByVectorQuery &query) const override; @@ -135,6 +140,12 @@ class CollectionImpl : public Collection { const std::string &column_name) const override; private: + // Query bodies without locking; the caller must hold schema_handle_mtx_ + // (at least shared) for the whole duration. + Result query_unsafe(const SearchQuery &query) const; + + Result query_unsafe(const MultiQuery &query) const; + void prepare_schema(); Status close_unsafe(); @@ -1611,6 +1622,49 @@ Result CollectionImpl::Query(const SearchQuery &query) const { CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); + return query_unsafe(query); +} + +Result CollectionImpl::Query(const MultiQuery &query) const { + std::shared_lock lock(schema_handle_mtx_); + + CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); + + return query_unsafe(query); +} + +Result CollectionImpl::QueryWithSchema( + const SearchQuery &query) const { + std::shared_lock lock(schema_handle_mtx_); + + CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); + + auto docs = query_unsafe(query); + if (!docs) { + return tl::make_unexpected(docs.error()); + } + // Snapshot the schema within the same critical section so it always matches + // the one used by query_unsafe, even under concurrent DDL. + return QuerySnapshot{std::move(*docs), *schema_}; +} + +Result CollectionImpl::QueryWithSchema( + const MultiQuery &query) const { + std::shared_lock lock(schema_handle_mtx_); + + CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); + + auto docs = query_unsafe(query); + if (!docs) { + return tl::make_unexpected(docs.error()); + } + // Snapshot the schema within the same critical section so it always matches + // the one used by query_unsafe, even under concurrent DDL. + return QuerySnapshot{std::move(*docs), *schema_}; +} + +Result CollectionImpl::query_unsafe( + const SearchQuery &query) const { // When field_name_ is set, use get_field to retrieve the schema uniformly. // validate checks that the field type matches the query type // (FTS query requires an FTS field, vector query requires a vector field). @@ -1637,11 +1691,7 @@ Result CollectionImpl::Query(const SearchQuery &query) const { return sql_engine_->execute(schema_, std::move(sanitized_query), segments); } -Result CollectionImpl::Query(const MultiQuery &query) const { - std::shared_lock lock(schema_handle_mtx_); - - CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); - +Result CollectionImpl::query_unsafe(const MultiQuery &query) const { if (query.queries.size() < 2) { return tl::make_unexpected(Status::InvalidArgument( "Invalid query: MultiQuery requires at least 2 sub-queries, got ", diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index c4b77575f..844d5bf03 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -25,6 +25,17 @@ namespace zvec { +/** + * @brief Atomic snapshot of a query execution: the result docs together with + * the exact schema the query was executed against. Both are captured within a + * single schema read-lock section, so they stay consistent even under + * concurrent DDL (e.g. DropColumn). + */ +struct ZVEC_API QuerySnapshot { + DocPtrList docs; + CollectionSchema schema; +}; + class ZVEC_API Collection { public: using Ptr = std::shared_ptr; @@ -102,6 +113,27 @@ class ZVEC_API Collection { virtual Result Query(const MultiQuery &query) const = 0; + /** + * @brief Execute a query and return the result docs together with the + * schema snapshot used by this execution. Unlike calling Query() and + * Schema() separately, both are captured atomically under the same + * read lock, so concurrent DDL cannot desynchronize them. + * + * @param query The search query. + * @return The query snapshot OR an error. + */ + virtual Result QueryWithSchema( + const SearchQuery &query) const = 0; + + /** + * @brief Multi-query variant of QueryWithSchema. + * + * @param query The multi query with re-ranking. + * @return The query snapshot OR an error. + */ + virtual Result QueryWithSchema( + const MultiQuery &query) const = 0; + virtual Result GroupByQuery( const GroupByVectorQuery &query) const = 0; diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index f98ef44f7..02511d66a 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -14,6 +14,8 @@ #include "zvec/db/collection.h" #include +#include +#include #include #include #include @@ -23,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -3845,6 +3848,123 @@ TEST_F(CollectionTest, Feature_Query_General) { } } +TEST_F(CollectionTest, Feature_QueryWithSchema_ConsistentWithQuery) { + FileHelper::RemoveDirectory(col_path); + + int doc_count = 100; + auto schema = TestHelper::CreateNormalSchema(); + auto options = CollectionOptions{false, false, 100 * 1024 * 1024}; + auto collection = TestHelper::CreateCollectionWithDoc(col_path, *schema, + options, 0, doc_count); + ASSERT_NE(collection, nullptr); + + auto query_doc = TestHelper::CreateDoc(1, *schema); + SearchQuery query; + query.topk_ = 10; + query.target_.field_name_ = "dense_fp32"; + auto vector = query_doc.get>("dense_fp32"); + ASSERT_TRUE(vector.has_value()); + query.target_.set_vector(std::string((char *)vector.value().data(), + vector.value().size() * sizeof(float))); + + auto plain = collection->Query(query); + ASSERT_TRUE(plain.has_value()); + + auto snapshot = collection->QueryWithSchema(query); + ASSERT_TRUE(snapshot.has_value()); + + // docs must match the plain Query result + ASSERT_EQ(snapshot->docs.size(), plain->size()); + for (size_t i = 0; i < plain->size(); ++i) { + ASSERT_EQ(*snapshot->docs[i], *(*plain)[i]); + } + // without concurrent DDL, the schema snapshot matches Schema() + auto current_schema = collection->Schema(); + ASSERT_TRUE(current_schema.has_value()); + ASSERT_EQ(snapshot->schema, current_schema.value()); +} + +TEST_F(CollectionTest, Feature_QueryWithSchema_AtomicUnderDropColumn) { + FileHelper::RemoveDirectory(col_path); + + int doc_count = 200; + auto schema = TestHelper::CreateNormalSchema(); + auto options = CollectionOptions{false, false, 100 * 1024 * 1024}; + auto collection = TestHelper::CreateCollectionWithDoc(col_path, *schema, + options, 0, doc_count); + ASSERT_NE(collection, nullptr); + + const std::string dropped_field = "int32"; + + auto query_doc = TestHelper::CreateDoc(1, *schema); + SearchQuery query; + query.topk_ = 10; + query.target_.field_name_ = "dense_fp32"; + auto vector = query_doc.get>("dense_fp32"); + ASSERT_TRUE(vector.has_value()); + query.target_.set_vector(std::string((char *)vector.value().data(), + vector.value().size() * sizeof(float))); + + std::atomic stop{false}; + std::atomic dropped{false}; + std::atomic invariant_violated{false}; + + // Invariant for every snapshot: the dropped column is present in the + // returned schema if and only if it is present in the returned docs, i.e. + // docs and schema come from the same consistent execution. + auto check_snapshot = [&](const Result &result) { + if (!result.has_value() || result->docs.empty()) { + return; + } + const bool schema_has_field = + result->schema.get_field(dropped_field) != nullptr; + for (const auto &doc : result->docs) { + if (doc && doc->has(dropped_field) != schema_has_field) { + invariant_violated = true; + } + } + }; + + std::thread drop_thread([&]() { + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + auto s = collection->DropColumn(dropped_field); + ASSERT_TRUE(s.ok()); + dropped = true; + }); + + // Hammer QueryWithSchema from two threads while the DDL runs. A small + // pause between queries keeps the shared lock from starving the DDL. + std::thread reader([&]() { + while (!stop) { + check_snapshot(collection->QueryWithSchema(query)); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + }); + + while (!dropped) { + check_snapshot(collection->QueryWithSchema(query)); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + // A few extra rounds after the DDL finished. + for (int i = 0; i < 10; ++i) { + check_snapshot(collection->QueryWithSchema(query)); + } + + stop = true; + reader.join(); + drop_thread.join(); + + ASSERT_FALSE(invariant_violated); + + // After the DDL, the field is gone from both schema and results. + auto final_result = collection->QueryWithSchema(query); + ASSERT_TRUE(final_result.has_value()); + ASSERT_EQ(final_result->schema.get_field(dropped_field), nullptr); + for (const auto &doc : final_result->docs) { + ASSERT_FALSE(doc->has(dropped_field)); + } +} + TEST_F(CollectionTest, Feature_Query_Empty) { auto func = [&](int doc_count, int topk) { FileHelper::RemoveDirectory(col_path); From 20f4cb97021a26f7e65dbfb67f6b88a02bc9415b Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 21 Aug 2026 11:49:40 +0800 Subject: [PATCH 4/6] fix --- src/binding/python/include/python_doc.h | 7 ++ src/binding/python/model/python_collection.cc | 42 ++++++---- src/binding/python/model/python_doc.cc | 11 ++- src/db/collection.cc | 77 ++++++++++++------- src/db/collection_query_internal.h | 48 ++++++++++++ src/include/zvec/db/collection.h | 32 -------- tests/db/collection_test.cc | 46 +++++------ 7 files changed, 164 insertions(+), 99 deletions(-) create mode 100644 src/db/collection_query_internal.h diff --git a/src/binding/python/include/python_doc.h b/src/binding/python/include/python_doc.h index 787f82eb9..e380cebfd 100644 --- a/src/binding/python/include/python_doc.h +++ b/src/binding/python/include/python_doc.h @@ -32,6 +32,13 @@ class ZVecPyDoc { // materialization path in the collection DQL bindings. Requires the GIL. static py::tuple doc_to_tuple(Doc &self, const CollectionSchema &schema); + // Same as doc_to_tuple but takes the pre-resolved forward/vector field lists + // directly, so batch materialization can resolve them once per batch instead + // of once per doc. Requires the GIL. + static py::tuple doc_to_tuple_with_fields( + Doc &self, const FieldSchemaPtrList &forward_fields, + const FieldSchemaPtrList &vector_fields); + // Convert a single Doc field value into a Python object according to its // DataType. Shared by the per-field `get_any` binding and `doc_to_tuple`. // Requires the GIL. diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 6e989749a..d5b3941e7 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -15,6 +15,7 @@ #include "python_collection.h" #include #include +#include "db/collection_query_internal.h" #include "python_doc.h" namespace zvec { @@ -23,13 +24,17 @@ namespace { // Batch-materialize a DocPtrList into a list of (id, score, fields, vectors) // tuples in a single GIL-held section, avoiding per-doc _Doc wrappers and -// per-doc Python->C++ crossings on the hot query path. +// per-doc Python->C++ crossings on the hot query path. The forward/vector field +// lists are resolved once per batch rather than once per doc. py::list docs_to_tuples(const DocPtrList &docs, const CollectionSchema &schema) { + const auto forward_fields = schema.forward_fields(); + const auto vector_fields = schema.vector_fields(); py::list out(docs.size()); for (size_t i = 0; i < docs.size(); ++i) { if (docs[i]) { - out[i] = ZVecPyDoc::doc_to_tuple(*docs[i], schema); + out[i] = ZVecPyDoc::doc_to_tuple_with_fields(*docs[i], forward_fields, + vector_fields); } else { out[i] = py::none(); } @@ -77,6 +82,21 @@ T unwrap_expected(tl::expected &&exp) { return T{}; } +// Run a query with the GIL released, capturing docs + schema atomically under a +// single schema read lock (internal::query_result_snapshot), then materialize +// the batch into tuples after the GIL is reacquired and the read lock released. +template +py::list execute_for_python(const Collection &collection, const Query &query) { + Result result; + { + py::gil_scoped_release release; + result = internal::query_result_snapshot(collection, query); + } + // GIL restored, schema read lock already released. + auto snapshot = unwrap_expected(std::move(result)); + return docs_to_tuples(snapshot.docs, *snapshot.schema); +} + void ZVecPyCollection::Initialize(pybind11::module_ &m) { py::class_(m, "_GroupResult") .def_readonly("group_by_value", &GroupResult::group_by_value_) @@ -296,19 +316,13 @@ void ZVecPyCollection::bind_dql_methods( py::class_ &col) { // Query with the GIL released, then materialize all hits into // (id, score, fields, vectors) tuples in one crossing (see docs_to_tuples). - // query_with_schema returns the docs and the schema snapshot atomically + // execute_for_python captures the docs and the schema snapshot atomically // under one read lock, so concurrent DDL cannot desynchronize them, while // the binding signature stays unchanged from the legacy per-doc binding. col.def( "Query", [](const Collection &self, const SearchQuery &query) { - Result result; - { - py::gil_scoped_release release; - result = self.query_with_schema(query); - } - auto snapshot = unwrap_expected(std::move(result)); - return docs_to_tuples(snapshot.docs, snapshot.schema); + return execute_for_python(self, query); }, py::arg("query"), "Execute a query and return results as a list of " @@ -317,13 +331,7 @@ void ZVecPyCollection::bind_dql_methods( .def( "Query", [](const Collection &self, const MultiQuery &query) { - Result result; - { - py::gil_scoped_release release; - result = self.query_with_schema(query); - } - auto snapshot = unwrap_expected(std::move(result)); - return docs_to_tuples(snapshot.docs, snapshot.schema); + return execute_for_python(self, query); }, py::arg("query"), "Execute a multi query with re-ranking and return results as a " diff --git a/src/binding/python/model/python_doc.cc b/src/binding/python/model/python_doc.cc index 1b180df02..764d12423 100644 --- a/src/binding/python/model/python_doc.cc +++ b/src/binding/python/model/python_doc.cc @@ -316,6 +316,13 @@ py::object ZVecPyDoc::doc_value_to_py(Doc &self, const std::string &field, } py::tuple ZVecPyDoc::doc_to_tuple(Doc &self, const CollectionSchema &schema) { + return doc_to_tuple_with_fields(self, schema.forward_fields(), + schema.vector_fields()); +} + +py::tuple ZVecPyDoc::doc_to_tuple_with_fields( + Doc &self, const FieldSchemaPtrList &forward_fields, + const FieldSchemaPtrList &vector_fields) { py::tuple result(4); // 1. set doc id and score result[0] = py::str(self.pk()); @@ -328,7 +335,7 @@ py::tuple ZVecPyDoc::doc_to_tuple(Doc &self, const CollectionSchema &schema) { } // 2. set scalar fields py::dict fields; - for (const auto &field_meta : schema.forward_fields()) { + for (const auto &field_meta : forward_fields) { const std::string &field = field_meta->name(); if (!self.has_value(field)) { continue; @@ -348,7 +355,7 @@ py::tuple ZVecPyDoc::doc_to_tuple(Doc &self, const CollectionSchema &schema) { } // 3. set vector fields py::dict vectors; - for (const auto &vec_meta : schema.vector_fields()) { + for (const auto &vec_meta : vector_fields) { const std::string &vec = vec_meta->name(); if (!self.has_value(vec)) continue; diff --git a/src/db/collection.cc b/src/db/collection.cc index f92214b81..f2959d323 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -35,6 +35,7 @@ #include #include #include +#include "db/collection_query_internal.h" #include "db/common/constants.h" #include "db/common/file_helper.h" #include "db/common/global_resource.h" @@ -126,12 +127,6 @@ class CollectionImpl : public Collection { Result query(const MultiQuery &query) const override; - Result query_with_schema( - const SearchQuery &query) const override; - - Result query_with_schema( - const MultiQuery &query) const override; - Result group_by_query( const GroupByVectorQuery &query) const override; @@ -143,6 +138,15 @@ class CollectionImpl : public Collection { Result debug_get_hnsw_storage_mode( const std::string &column_name) const override; + // Execute a query and capture the docs together with a shared_ptr snapshot of + // schema_, all within a single shared lock on schema_handle_mtx_. Used by the + // internal query_result_snapshot() free functions below. Public because those + // free functions are not members/friends, but CollectionImpl itself is not + // exposed in any public header, so this stays internal to this .cc. + template + Result query_result_snapshot_impl( + const Query &query) const; + private: // Query bodies without locking; the caller must hold schema_handle_mtx_ // (at least shared) for the whole duration. @@ -1706,8 +1710,9 @@ Result CollectionImpl::query(const MultiQuery &query) const { return query_unsafe(query); } -Result CollectionImpl::query_with_schema( - const SearchQuery &query) const { +template +Result +CollectionImpl::query_result_snapshot_impl(const Query &query) const { std::shared_lock lock(schema_handle_mtx_); CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); @@ -1718,24 +1723,12 @@ Result CollectionImpl::query_with_schema( return tl::make_unexpected(docs.error()); } // Snapshot the schema within the same critical section so it always matches - // the one used by query_unsafe, even under concurrent DDL. - return QuerySnapshot{std::move(*docs), *schema_}; -} - -Result CollectionImpl::query_with_schema( - const MultiQuery &query) const { - std::shared_lock lock(schema_handle_mtx_); - - CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); - CHECK_CLOSED_RETURN_STATUS_EXPECTED(closed_, false); - - auto docs = query_unsafe(query); - if (!docs) { - return tl::make_unexpected(docs.error()); - } - // Snapshot the schema within the same critical section so it always matches - // the one used by query_unsafe, even under concurrent DDL. - return QuerySnapshot{std::move(*docs), *schema_}; + // the one used by query_unsafe, even under concurrent DDL. The shared_ptr + // only bumps the refcount; DDL uses clone-and-swap under the exclusive lock, + // so this captured schema stays valid after the lock is released. + std::shared_ptr schema_snapshot = schema_; + return internal::QueryResultSnapshot{std::move(docs.value()), + std::move(schema_snapshot)}; } Result CollectionImpl::query_unsafe( @@ -2313,4 +2306,36 @@ std::vector CollectionImpl::get_all_persist_segments() const { return segment_manager_->get_segments(); } +namespace internal { + +namespace { +// The binding layer only holds a Collection reference and cannot see +// CollectionImpl, so recover the concrete type here. CollectionImpl is the sole +// Collection implementation; the dynamic_cast is negligible next to a query and +// safer than assuming the concrete type. Should a decorator/proxy Collection +// ever appear, this returns NotSupported at runtime instead of misbehaving. +template +Result query_result_snapshot_dispatch( + const Collection &collection, const Query &query) { + const auto *impl = dynamic_cast(&collection); + if (impl == nullptr) { + return tl::make_unexpected( + Status::NotSupported("Unsupported Collection implementation")); + } + return impl->query_result_snapshot_impl(query); +} +} // namespace + +Result query_result_snapshot(const Collection &collection, + const SearchQuery &query) { + return query_result_snapshot_dispatch(collection, query); +} + +Result query_result_snapshot(const Collection &collection, + const MultiQuery &query) { + return query_result_snapshot_dispatch(collection, query); +} + +} // namespace internal + } // namespace zvec diff --git a/src/db/collection_query_internal.h b/src/db/collection_query_internal.h new file mode 100644 index 000000000..3aaee70a5 --- /dev/null +++ b/src/db/collection_query_internal.h @@ -0,0 +1,48 @@ +// 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. + +#pragma once + +#include +#include + +namespace zvec { +namespace internal { + +// Result of a query captured atomically with the schema it executed against. +// The schema is held by shared_ptr: DDL uses clone-and-swap under the exclusive +// schema lock, so the snapshot stays valid after the read lock is released. +// +// This is an internal (non-public) type used by language bindings to +// materialize query results without racing against concurrent DDL. It is not +// part of the stable Collection API. +struct QueryResultSnapshot { + DocPtrList docs; + std::shared_ptr schema; +}; + +// Execute a query and capture the docs together with the schema they were +// executed against, atomically within a single schema read-lock section. +// +// These free functions assume the concrete Collection is a CollectionImpl (the +// only implementation); an unexpected implementation yields a NotSupported +// error at runtime. +Result query_result_snapshot(const Collection &collection, + const SearchQuery &query); + +Result query_result_snapshot(const Collection &collection, + const MultiQuery &query); + +} // namespace internal +} // namespace zvec diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index 70d944a6a..f34f60645 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -25,17 +25,6 @@ namespace zvec { -/** - * @brief Atomic snapshot of a query execution: the result docs together with - * the exact schema the query was executed against. Both are captured within a - * single schema read-lock section, so they stay consistent even under - * concurrent DDL (e.g. DropColumn). - */ -struct ZVEC_API QuerySnapshot { - DocPtrList docs; - CollectionSchema schema; -}; - class ZVEC_API Collection { public: using Ptr = std::shared_ptr; @@ -115,27 +104,6 @@ class ZVEC_API Collection { virtual Result query(const MultiQuery &query) const = 0; - /** - * @brief Execute a query and return the result docs together with the - * schema snapshot used by this execution. Unlike calling query() and - * schema() separately, both are captured atomically under the same - * read lock, so concurrent DDL cannot desynchronize them. - * - * @param query The search query. - * @return The query snapshot OR an error. - */ - virtual Result query_with_schema( - const SearchQuery &query) const = 0; - - /** - * @brief Multi-query variant of query_with_schema. - * - * @param query The multi query with re-ranking. - * @return The query snapshot OR an error. - */ - virtual Result query_with_schema( - const MultiQuery &query) const = 0; - virtual Result group_by_query( const GroupByVectorQuery &query) const = 0; diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index 53f1d2279..ca7d58920 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -33,6 +33,7 @@ #include #include #include +#include "db/collection_query_internal.h" #include "db/common/file_helper.h" #include "db/index/common/type_helper.h" #include "db/index/common/version_manager.h" @@ -4197,7 +4198,7 @@ TEST_F(CollectionTest, Feature_Query_General) { } } -TEST_F(CollectionTest, Feature_QueryWithSchema_ConsistentWithQuery) { +TEST_F(CollectionTest, Feature_QueryResultSnapshot_ConsistentWithQuery) { FileHelper::RemoveDirectory(col_path); int doc_count = 100; @@ -4219,7 +4220,7 @@ TEST_F(CollectionTest, Feature_QueryWithSchema_ConsistentWithQuery) { auto plain = collection->query(query); ASSERT_TRUE(plain.has_value()); - auto snapshot = collection->query_with_schema(query); + auto snapshot = internal::query_result_snapshot(*collection, query); ASSERT_TRUE(snapshot.has_value()); // docs must match the plain Query result @@ -4230,10 +4231,10 @@ TEST_F(CollectionTest, Feature_QueryWithSchema_ConsistentWithQuery) { // without concurrent DDL, the schema snapshot matches schema() auto current_schema = collection->schema(); ASSERT_TRUE(current_schema.has_value()); - ASSERT_EQ(snapshot->schema, current_schema.value()); + ASSERT_EQ(*snapshot->schema, current_schema.value()); } -TEST_F(CollectionTest, Feature_QueryWithSchema_AtomicUnderDropColumn) { +TEST_F(CollectionTest, Feature_QueryResultSnapshot_AtomicUnderDropColumn) { FileHelper::RemoveDirectory(col_path); int doc_count = 200; @@ -4261,18 +4262,19 @@ TEST_F(CollectionTest, Feature_QueryWithSchema_AtomicUnderDropColumn) { // Invariant for every snapshot: the dropped column is present in the // returned schema if and only if it is present in the returned docs, i.e. // docs and schema come from the same consistent execution. - auto check_snapshot = [&](const Result &result) { - if (!result.has_value() || result->docs.empty()) { - return; - } - const bool schema_has_field = - result->schema.get_field(dropped_field) != nullptr; - for (const auto &doc : result->docs) { - if (doc && doc->has(dropped_field) != schema_has_field) { - invariant_violated = true; - } - } - }; + auto check_snapshot = + [&](const Result &result) { + if (!result.has_value() || result->docs.empty()) { + return; + } + const bool schema_has_field = + result->schema->get_field(dropped_field) != nullptr; + for (const auto &doc : result->docs) { + if (doc && doc->has(dropped_field) != schema_has_field) { + invariant_violated = true; + } + } + }; std::thread drop_thread([&]() { std::this_thread::sleep_for(std::chrono::milliseconds(50)); @@ -4281,22 +4283,22 @@ TEST_F(CollectionTest, Feature_QueryWithSchema_AtomicUnderDropColumn) { dropped = true; }); - // Hammer query_with_schema from two threads while the DDL runs. A small + // Hammer query_result_snapshot from two threads while the DDL runs. A small // pause between queries keeps the shared lock from starving the DDL. std::thread reader([&]() { while (!stop) { - check_snapshot(collection->query_with_schema(query)); + check_snapshot(internal::query_result_snapshot(*collection, query)); std::this_thread::sleep_for(std::chrono::milliseconds(1)); } }); while (!dropped) { - check_snapshot(collection->query_with_schema(query)); + check_snapshot(internal::query_result_snapshot(*collection, query)); std::this_thread::sleep_for(std::chrono::milliseconds(1)); } // A few extra rounds after the DDL finished. for (int i = 0; i < 10; ++i) { - check_snapshot(collection->query_with_schema(query)); + check_snapshot(internal::query_result_snapshot(*collection, query)); } stop = true; @@ -4306,9 +4308,9 @@ TEST_F(CollectionTest, Feature_QueryWithSchema_AtomicUnderDropColumn) { ASSERT_FALSE(invariant_violated); // After the DDL, the field is gone from both schema and results. - auto final_result = collection->query_with_schema(query); + auto final_result = internal::query_result_snapshot(*collection, query); ASSERT_TRUE(final_result.has_value()); - ASSERT_EQ(final_result->schema.get_field(dropped_field), nullptr); + ASSERT_EQ(final_result->schema->get_field(dropped_field), nullptr); for (const auto &doc : final_result->docs) { ASSERT_FALSE(doc->has(dropped_field)); } From 58d82a17ffbed5bdb92cc112a571592f01314525 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 21 Aug 2026 13:11:12 +0800 Subject: [PATCH 5/6] fix: drop redundant kTaskCount lambda capture in ThreadQueue tests upstream #688 made kTaskCount static constexpr, which already avoids the MSVC capture requirement. The explicit capture left over from the merge is invalid (static storage duration cannot be captured) and broke clang-tidy. Align the file with upstream. --- tests/ailego/parallel/thread_queue_test.cc | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/ailego/parallel/thread_queue_test.cc b/tests/ailego/parallel/thread_queue_test.cc index 2d8aa26d0..1a1fff39a 100644 --- a/tests/ailego/parallel/thread_queue_test.cc +++ b/tests/ailego/parallel/thread_queue_test.cc @@ -56,9 +56,7 @@ TEST(ThreadQueue, General) { { std::unique_lock lock(count_mutex); completed = count_cond.wait_for(lock, std::chrono::seconds(10), - [&count, kTaskCount]() { - return count == kTaskCount; - }); + [&count]() { return count == kTaskCount; }); completed_count = count; } @@ -100,9 +98,7 @@ TEST(ThreadQueue, MutliThread) { { std::unique_lock lock(count_mutex); completed = count_cond.wait_for(lock, std::chrono::seconds(10), - [&count, kTaskCount]() { - return count == kTaskCount; - }); + [&count]() { return count == kTaskCount; }); completed_count = count; } From ab400b47ad15f54036c8e38eb1ae669b9bf3378d Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Mon, 24 Aug 2026 11:44:42 +0800 Subject: [PATCH 6/6] =?UTF-8?q?remove=20test=5Fbatch=5Fmaterialize.py?= =?UTF-8?q?=E2=80=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- python/tests/test_batch_materialize.py | 180 ------------------------- python/tests/test_collection.py | 63 +++++++++ src/binding/python/model/python_doc.cc | 4 +- 3 files changed, 65 insertions(+), 182 deletions(-) delete mode 100644 python/tests/test_batch_materialize.py diff --git a/python/tests/test_batch_materialize.py b/python/tests/test_batch_materialize.py deleted file mode 100644 index 989cbe3bc..000000000 --- a/python/tests/test_batch_materialize.py +++ /dev/null @@ -1,180 +0,0 @@ -# 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. -"""Correctness tests for the batch-materialized query path. - -`Collection.query` goes through `_Collection.Query`, which batch-materializes -all hits into tuples in a single C++ call. These tests validate the -materialized output against two independent references: - -- a numpy brute-force ground truth over the inserted vectors (ids / scores); -- the `fetch` path, which materializes docs through a separate binding. -""" - -from __future__ import annotations - -import numpy as np -import pytest -import zvec -from zvec import ( - Collection, - CollectionOption, - DataType, - Doc, - FieldSchema, - HnswIndexParam, - HnswQueryParam, - Query, - RrfReRanker, - VectorSchema, -) -from zvec.typing import MetricType - -DIM = 16 -N_DOCS = 200 - - -def _make_vectors() -> np.ndarray: - """Same deterministic vectors as inserted by the fixture.""" - return np.random.default_rng(42).random((N_DOCS, DIM), dtype=np.float32) - - -def _brute_force_topk( - query: np.ndarray, topk: int, mask: np.ndarray | None = None -) -> tuple[list[str], np.ndarray]: - """Exact L2sq top-k ids and distances over the ground-truth vectors.""" - dists = ((_make_vectors() - query) ** 2).sum(axis=1) - if mask is not None: - dists = np.where(mask, dists, np.inf) - idx = np.argsort(dists, kind="stable")[:topk] - return [str(i) for i in idx], dists[idx] - - -@pytest.fixture(scope="module") -def bm_collection(tmp_path_factory) -> Collection: - schema = zvec.CollectionSchema( - name="batch_mat_test", - fields=[ - FieldSchema("num", DataType.INT64, nullable=False), - FieldSchema("title", DataType.STRING, nullable=True), - ], - vectors=[ - VectorSchema( - "vec", - DataType.VECTOR_FP32, - dimension=DIM, - # explicit L2: score is the raw squared L2 distance (no - # metric normalization), matching the brute-force ground truth - index_param=HnswIndexParam(metric_type=MetricType.L2), - ), - ], - ) - path = tmp_path_factory.mktemp("zvec_batch_mat") / "coll" - coll = zvec.create_and_open( - path=str(path), - schema=schema, - option=CollectionOption(read_only=False, enable_mmap=True), - ) - - vectors = _make_vectors() - docs = [ - Doc( - id=str(i), - fields={"num": i, "title": f"doc-{i}"}, - vectors={"vec": vectors[i]}, - ) - for i in range(N_DOCS) - ] - for r in coll.insert(docs): - assert r.ok() - - yield coll - - try: - coll.destroy() - except Exception: - pass - - -class TestBatchMaterialize: - def _query_vec(self) -> np.ndarray: - return np.array([0.5] * DIM, dtype=np.float32) - - def test_matches_brute_force_ground_truth(self, bm_collection: Collection): - q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) - docs = bm_collection.query(q, topk=20) - assert len(docs) == 20 - - exp_ids, exp_dists = _brute_force_topk(self._query_vec(), 20) - assert [d.id for d in docs] == exp_ids - scores = [d.score for d in docs] - assert scores == sorted(scores) - for d, dist in zip(docs, exp_dists): - assert d.score == pytest.approx(float(dist), rel=1e-4) - - # scalar fields fully materialized, vectors excluded by default - for d in docs: - assert isinstance(d, Doc) - assert set(d.fields.keys()) == {"num", "title"} - assert d.fields["num"] == int(d.id) - assert d.fields["title"] == f"doc-{d.id}" - assert d.vectors == {} - - def test_fields_match_fetch_path(self, bm_collection: Collection): - q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) - docs = bm_collection.query(q, topk=10) - fetched = bm_collection.fetch([d.id for d in docs], include_vector=False) - for d in docs: - assert d.fields == fetched[d.id].fields - - @pytest.mark.parametrize("include_vector", [False, True]) - def test_include_vector(self, bm_collection: Collection, include_vector: bool): - q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) - docs = bm_collection.query(q, topk=5, include_vector=include_vector) - ground_truth = _make_vectors() - for d in docs: - assert bool(d.vectors) is include_vector - if include_vector: - assert np.allclose(d.vectors["vec"], ground_truth[int(d.id)]) - - def test_output_fields_subset(self, bm_collection: Collection): - q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) - docs = bm_collection.query(q, topk=5, output_fields=["num"]) - exp_ids, _ = _brute_force_topk(self._query_vec(), 5) - assert [d.id for d in docs] == exp_ids - for d in docs: - assert set(d.fields.keys()) == {"num"} - - def test_filter(self, bm_collection: Collection): - q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) - docs = bm_collection.query(q, topk=10, filter="num < 50") - mask = np.arange(N_DOCS) < 50 - exp_ids, exp_dists = _brute_force_topk(self._query_vec(), 10, mask) - assert [d.id for d in docs] == exp_ids - for d, dist in zip(docs, exp_dists): - assert d.fields["num"] < 50 - assert d.score == pytest.approx(float(dist), rel=1e-4) - - def test_empty_result(self, bm_collection: Collection): - q = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) - docs = bm_collection.query(q, topk=10, filter="num < 0") - assert docs == [] - - def test_multi_query_rrf(self, bm_collection: Collection): - q1 = Query(field_name="vec", vector=self._query_vec(), param=HnswQueryParam()) - q2 = Query(field_name="vec", vector=[0.1] * DIM, param=HnswQueryParam()) - docs = bm_collection.query([q1, q2], topk=10, reranker=RrfReRanker()) - assert len(docs) == 10 - for d in docs: - assert isinstance(d, Doc) - assert set(d.fields.keys()) == {"num", "title"} diff --git a/python/tests/test_collection.py b/python/tests/test_collection.py index b03d1e3cf..7fabdcb2f 100644 --- a/python/tests/test_collection.py +++ b/python/tests/test_collection.py @@ -1229,3 +1229,66 @@ def my_rerank_callback(query_results, fields, topn): ) assert len(result) > 0 assert len(result) <= 5 + + def test_collection_query_materializes_every_hit( + self, collection_with_multiple_docs: Collection + ): + """Query results are materialized as a batch; every hit must keep its own + scalar values, so a misaligned batch cannot go unnoticed.""" + # "dense" is [id + 0.1] * 128 under the default IP metric, so the hits + # are the largest ids in descending score order. + result = collection_with_multiple_docs.query( + Query(field_name="dense", vector=[100.1] * 128), topk=5 + ) + assert [doc.id for doc in result] == ["100", "99", "98", "97", "96"] + + scores = [doc.score for doc in result] + assert scores == sorted(scores, reverse=True) + + fetched = collection_with_multiple_docs.fetch([doc.id for doc in result]) + for doc in result: + # Values belong to this doc, not to a neighbour in the batch. + assert doc.field("id") == int(doc.id) + assert doc.field("name") == "test" + assert doc.field("weight") == 80.0 + assert doc.field("height") == 210 + # Cross-check against fetch, which materializes docs separately. + assert set(doc.field_names()) == set(fetched[doc.id].field_names()) + for name in doc.field_names(): + assert doc.field(name) == fetched[doc.id].field(name) + assert doc.vectors == {} + + def test_collection_query_materializes_vectors_per_hit( + self, collection_with_multiple_docs: Collection + ): + """With include_vector, each hit must carry its own vector values.""" + result = collection_with_multiple_docs.query( + Query(field_name="dense", vector=[100.1] * 128), + topk=3, + include_vector=True, + ) + assert len(result) == 3 + for doc in result: + expected = pytest.approx(int(doc.id) + 0.1, rel=1e-5) + assert len(doc.vector("dense")) == 128 + assert all(value == expected for value in doc.vector("dense")) + assert doc.vector("sparse") == {1: 1.0, 2: 2.0, 3: 3.0} + + def test_collection_multi_query_materializes_every_hit( + self, collection_with_multiple_docs: Collection, multiple_docs + ): + """The multi-query binding materializes results through its own code + path, so it needs the same per-hit guarantees as a single query.""" + 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=RrfReRanker(), + ) + assert len(result) > 0 + for doc in result: + assert doc.field("id") == int(doc.id) + assert doc.field("name") == "test" + assert doc.vectors == {} diff --git a/src/binding/python/model/python_doc.cc b/src/binding/python/model/python_doc.cc index 764d12423..ae03adaaf 100644 --- a/src/binding/python/model/python_doc.cc +++ b/src/binding/python/model/python_doc.cc @@ -219,8 +219,8 @@ void ZVecPyDoc::bind_doc(py::module_ &m) { -> py::object { return doc_value_to_py(self, field, type); }); doc.def("get_all", &ZVecPyDoc::doc_to_tuple, py::arg("schema"), "Get all fields and vectors as a tuple: (id, score, fields, " - "vectors). Vectors are zero-copy numpy arrays (dense: ndarray, " - "sparse: (indices, values) tuple)."); + "vectors). Dense vectors are returned as lists, sparse vectors as " + "{index: value} dicts."); } py::object ZVecPyDoc::doc_value_to_py(Doc &self, const std::string &field,