Batch convert docs to achieve acceleration

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"}

python/tests/test_query_executor.py

@@ -286,20 +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 converted into a result list.
-        schema = MockCollectionSchema()
-        executor = QueryExecutor(schema)
+        # 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.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: ("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)

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]: ...

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,14 @@ 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
+        (the schema is resolved inside the binding from the collection),
+        avoiding per-doc Python/C++ crossings on the hot path.
+        """
+        tuples = collection.Query(query)
+        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 +164,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)
+        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]

src/binding/python/include/python_doc.h

@@ -14,6 +14,7 @@
 
 #include <pybind11/pybind11.h>
 #include <zvec/db/doc.h>
+#include <zvec/db/schema.h>
 
 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);

src/binding/python/model/python_collection.cc

@@ -15,9 +15,30 @@
 #include "python_collection.h"
 #include <pybind11/stl.h>
 #include <zvec/db/collection.h>
+#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(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[i] = py::none();
+    }
+  }
+  return out;
+}
+
+}  // namespace
+
 inline void throw_if_error(const Status &status) {
   switch (status.code()) {
     case StatusCode::OK:
@@ -255,29 +276,44 @@ void ZVecPyCollection::bind_dml_methods(
 
 void ZVecPyCollection::bind_dql_methods(
     py::class_<Collection, Collection::Ptr> &col) {
-  col.def("Query",
-          [](const Collection &self, const SearchQuery &query) {
-            Result<DocPtrList> 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).
+  // 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) {
+           Result<DocPtrList> result;
+           Result<CollectionSchema> schema_result;
+           {
+             py::gil_scoped_release release;
+             result = self.Query(query);
+             schema_result = self.Schema();
+           }
+           return docs_to_tuples(unwrap_expected(result),
+                                 unwrap_expected(schema_result));
+         },
+         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) {
             Result<DocPtrList> result;
+            Result<CollectionSchema> schema_result;
             {
               py::gil_scoped_release release;
               result = self.Query(query);
+              schema_result = self.Schema();
             }
-            // return DocPtrList
-            return unwrap_expected(result);
+            return docs_to_tuples(unwrap_expected(result),
+                                  unwrap_expected(schema_result));
           },
-          py::arg("query"), "Execute a multi query with re-ranking.")
+          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.")
       .def("GroupByQuery",
            [](const Collection &self, const GroupByVectorQuery &query) {
              Result<GroupResults> result;