perf: batch convert docs to achieve acceleration - #647
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
This PR changes the Python query hot path to avoid per-document Python/C++ binding overhead by batch-materializing query results as lightweight tuples in C++, then constructing Doc objects from those tuples in Python.
Changes:
- Refactors
_Collection.Querybindings to return a batch of(id, score, fields, vectors)tuples, parameterized by schema. - Extracts per-doc materialization into a shared C++ helper (
ZVecPyDoc::doc_to_tuple) and reuses it for both per-doc and batch paths. - Updates the Python
QueryExecutorto consume tuples viaDoc._from_tuple, and adds correctness tests for the new batch-materialized behavior.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| src/binding/python/model/python_doc.cc | Extracts get_all logic into doc_to_tuple and binds get_all to it. |
| src/binding/python/model/python_collection.cc | Changes Query bindings to return batch-materialized tuples and adds docs_to_tuples helper. |
| src/binding/python/include/python_doc.h | Declares doc_to_tuple and adds schema dependency for the new signature. |
| python/zvec/executor/query_executor.py | Switches query execution to consume tuple batches and build Doc via Doc._from_tuple. |
| python/tests/test_query_executor.py | Adjusts unit test to patch Doc._from_tuple instead of convert_to_py_doc. |
| python/tests/test_batch_materialize.py | Adds integration-style correctness tests for batch-materialized query results. |
Suppressed comments (1)
src/binding/python/include/python_doc.h:1
- Including
<zvec/db/schema.h>in this public binding header increases compile-time coupling for every translation unit that includespython_doc.h. Sincedoc_to_tupleonly needsCollectionSchemaby reference in the declaration, you can forward-declareclass CollectionSchema;innamespace zvechere and move the schema include into the corresponding.ccfile.
// Copyright 2025-present the zvec project
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| Result<CollectionSchema> schema_result; | ||
| { | ||
| py::gil_scoped_release release; | ||
| result = self.Query(query); |
There was a problem hiding this comment.
self.Query(query) & self.Schema() 不是同一把锁吧？ 并发情况下，如果drop_column，会不会有解析丢字段的问题？
There was a problem hiding this comment.
确实有这个问题，Query() 和Schema()是两次独立的加锁，两者之间存在无锁窗口。新增 QueryWithSchema 原子接口，在同一把 schema 读锁内返回 docs 与查询所用的 schema 快照（QuerySnapshot），绑定层改为单次调用，消除了原 Query() + Schema() 两次加锁之间可被 DropColumn 插入的窗口
There was a problem hiding this comment.
我觉得当前的解决方案不太优雅。为了解决pybind 层带来的开销，我们要去维护一个query_with_schema的接口。（或者至少这个接口应该是internal的）。
能不能将这个batch的操作在python & pybind层实现。（在python 有一个比较恶心的问题，就是需要在python 层加锁，去保证 query & get_schema 的原子操作。） @egolearner @feihongxu0824 @zhourrr 看下有啥好的法子没
There was a problem hiding this comment.
不增加公共api的一个思路，供参考
建议核心模板叫 query_result_snapshot_impl，避免 DB 层出现 Python 语义；execute_for_python 留在 binding 层。实现分三层。
1. 内部结果类型
新增非公开头文件 src/db/collection_query_internal.h：
#pragma once
#include <memory>
#include <zvec/db/collection.h>
namespace zvec::internal {
struct QueryResultSnapshot {
DocPtrList docs;
std::shared_ptr<const CollectionSchema> schema;
};
Result<QueryResultSnapshot> query_result_snapshot(
const Collection &collection, const SearchQuery &query);
} // namespace zvec::internal如果当前确实需要支持 MultiQuery，再增加一个同名 overload；否则不要提前泛化。
2. CollectionImpl 内部模板
在 collection.cc 的 CollectionImpl 中声明：
template <typename Query>
Result<internal::QueryResultSnapshot> query_result_snapshot_impl(
const Query &query) const;模板只在 collection.cc 使用，因此定义也留在 .cc，不暴露 CollectionImpl：
template <typename Query>
Result<internal::QueryResultSnapshot>
CollectionImpl::query_result_snapshot_impl(const Query &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());
}
// schema_ 与 query_unsafe 使用的 schema 属于同一个读锁区间。
// shared_ptr 只增加引用计数，不复制 CollectionSchema。
std::shared_ptr<const CollectionSchema> schema_snapshot = schema_;
return internal::QueryResultSnapshot{
std::move(docs.value()),
std::move(schema_snapshot),
};
}这里的关键点：
query_unsafe()不再加锁，由调用者保证锁。- 查询完成和捕获
schema_都发生在同一个共享锁内。 - DDL 需要独占锁，因此不能在两者之间替换 schema。
- DDL 采用 clone-and-swap；解锁后即使发生
DropColumn，旧 schema 仍由shared_ptr保活。 - 不进行
CollectionSchema深拷贝。 - Python 对象转换不在锁内执行，避免长时间阻塞 DDL。
3. 内部 free function
binding 只有 Collection&，看不到 CollectionImpl，因此在 collection.cc 提供一个内部转发：
namespace zvec::internal {
Result<QueryResultSnapshot> query_result_snapshot(
const Collection &collection, const SearchQuery &query) {
const auto *impl = dynamic_cast<const CollectionImpl *>(&collection);
if (impl == nullptr) {
return tl::make_unexpected(
Status::NotSupported("Unsupported Collection implementation"));
}
return impl->query_result_snapshot_impl(query);
}
} // namespace zvec::internaldynamic_cast 的成本相对查询可以忽略，并且比假设具体实现后直接 static_cast 更安全。
该函数不进入公开 collection.h，也不增加 Collection virtual，所以不改变 C++ vtable。
4. Python binding
binding 层才叫 execute_for_python：
template <typename Query>
py::list execute_for_python(const Collection &collection,
const Query &query) {
Result<internal::QueryResultSnapshot> result;
{
py::gil_scoped_release release;
result = internal::query_result_snapshot(collection, query);
}
// 到这里 GIL 已恢复，schema 读锁已经释放。
auto snapshot = unwrap_expected(std::move(result));
return docs_to_tuples(snapshot.docs, *snapshot.schema);
}绑定保持简单：
col.def(
"Query",
[](const Collection &collection, const SearchQuery &query) {
return execute_for_python(collection, query);
},
py::arg("query"));完整时序是：
释放 GIL
→ 获取 schema 读锁
→ 执行一个 query
→ O(1) 保留同版本 schema 指针
→ 释放 schema 读锁
恢复 GIL
→ 将整批 Doc 转成 list[tuple]
→ 返回 Python
建议额外让 docs_to_tuples 每批只调用一次 schema.forward_fields() 和 schema.vector_fields()，不要在每个 doc_to_tuple 中重复构造字段列表。
这样既保持“一次 query、一批结果统一转换”，又把公开 Interface、锁、schema 生命周期和 Python 物化职责分得很清楚。
zzlin237
changed the title
zzlin237
changed the title
# Conflicts: # src/db/collection.cc
# Conflicts: # src/binding/python/model/python_collection.cc # src/db/collection.cc # src/include/zvec/db/collection.h
# Conflicts: # tests/db/collection_test.cc
| */ | ||
| struct ZVEC_API QuerySnapshot { | ||
| DocPtrList docs; | ||
| CollectionSchema schema; |
There was a problem hiding this comment.
schema可以用shared_ptr<CollectionSchema>，并且检查下schema变更必须先创建副本再修改，这样多数场景都只需要引用计数+1
struct名称改为QueryResultSnapshot? 或者更直白的DocsAndSchema
# Conflicts: # src/binding/python/model/python_collection.cc # src/db/collection.cc
upstream alibaba#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.
Change the query result materialization from converting each document individually across the Python/C++ boundary into a PyDoc, to batch-materializing results as lightweight tuples according to the schema in a single C++ call and then constructing Doc objects via Doc._from_tuple on the Python side.
This aims to eliminate the per-document language binding conversion overhead on the query hot path, thereby reducing query latency.
Dataset: cohere-1m, M: 15