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

feat(binding): use VectorViewClause zero-copy for Python dense vector query (#517)

This commit is contained in:
egolearner 2026-07-06 14:03:34 +08:00 committed by GitHub
parent 54340bb705
commit 998a711434
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 117 additions and 118 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
python/zvec/model/param/__init__.pyi
View File

@ -1147,7 +1147,13 @@ class _SearchQuery:
def __getstate__(self) -> tuple: ...
def __init__(self) -> None: ...
def __setstate__(self, arg0: tuple) -> None: ...
def set_vector(self, arg0: ..., arg1: typing.Any) -> None: ...
def set_vector(self, field_schema: typing.Any, obj: typing.Any) -> None:
"""
Set the query vector.
Dense vector source data must not be modified until the query finishes.
"""
@property
def output_fields(self) -> list[str] | None: ...
@output_fields.setter

4
src/binding/c/c_api.cc
View File

@ -5156,6 +5156,8 @@ zvec_error_code_t zvec_vector_query_set_query_vector(zvec_vector_query_t *query,
return ZVEC_ERROR_INVALID_ARGUMENT;
}
auto *ptr = reinterpret_cast<zvec::SearchQuery *>(query);
// Copies into VectorClause (not VectorViewClause) because the C API does
// not require `data` to stay alive after this call returns.
ptr->target_.set_vector(std::string(static_cast<const char *>(data), size));
return ZVEC_OK;
}
@ -5646,6 +5648,7 @@ zvec_error_code_t zvec_group_by_vector_query_set_query_vector(
return ZVEC_ERROR_INVALID_ARGUMENT;
}
auto *ptr = reinterpret_cast<zvec::GroupByVectorQuery *>(query);
// Copies into VectorClause — see comment on zvec_vector_query_set_query_vector.
ptr->target_.set_vector(std::string(static_cast<const char *>(data), size));
return ZVEC_OK;
}
@ -6067,6 +6070,7 @@ zvec_error_code_t zvec_sub_query_set_query_vector(
return ZVEC_ERROR_INVALID_ARGUMENT;
}
auto *ptr = reinterpret_cast<zvec::SubQuery *>(query);
// Copies into VectorClause — see comment on zvec_vector_query_set_query_vector.
auto &payload = std::get<zvec::VectorClause>(ptr->target_.clause_);
payload.query_vector_.assign(static_cast<const char *>(data), size);
return ZVEC_OK;

223
src/binding/python/model/param/python_param.cc
View File

@ -88,14 +88,6 @@ T checked_cast(const py::handle &h, const std::string &vector_field,
}
}
template <typename T>
std::string serialize_vector(const T *data, size_t n) {
std::string buf;
buf.resize(n * sizeof(T));
std::memcpy(buf.data(), data, n * sizeof(T));
return buf;
}
template <typename ValueType, typename ValueCastFn>
std::pair<std::string, std::string> serialize_sparse_vector(
const py::dict &sparse_dict, ValueCastFn &&value_caster) {
@ -1943,6 +1935,15 @@ void ZVecPyParams::bind_vector_query(py::module_ &m) {
SubQuery sub;
sub.num_candidates_ = sq.topk_;
sub.target_ = sq.target_;
// SubQuery is copied by value into MultiQuery. Materialize
// non-owning vector views so the copied SubQuery does not depend on
// the original _SearchQuery keep-alive relationship.
if (auto *vvc = sub.target_.get_vector_view_clause()) {
VectorClause vc{std::string(vvc->query_vector_),
std::string(vvc->sparse_indices_),
std::string(vvc->sparse_values_)};
sub.target_.clause_ = std::move(vc);
}
return sub;
},
py::arg("search_query"),
@ -1987,109 +1988,97 @@ void ZVecPyParams::bind_vector_query(py::module_ &m) {
}
})
// vector
.def("set_vector",
[](SearchQuery &self, const FieldSchema &field_schema,
const py::object &obj) {
const DataType data_type = field_schema.data_type();
.def(
"set_vector",
[](SearchQuery &self, const FieldSchema &field_schema,
const py::object &obj) {
const DataType data_type = field_schema.data_type();
// dense vector
if (FieldSchema::is_dense_vector_field(data_type)) {
if (!py::isinstance<py::array>(obj)) {
throw py::type_error("Dense vector[" + field_schema.name() +
"] expects a ndarray, got " +
std::string(py::str(py::type::of(obj))));
}
const auto arr = obj.cast<py::array>();
if (arr.ndim() != 1) {
throw py::type_error("Dense vector expects 1D array, got " +
std::to_string(arr.ndim()) + "D");
}
const auto buf = arr.request();
switch (data_type) {
case DataType::VECTOR_FP32: {
self.target_.set_vector(serialize_vector<float>(
static_cast<const float *>(buf.ptr), buf.size));
return;
}
case DataType::VECTOR_FP64: {
self.target_.set_vector(serialize_vector<double>(
static_cast<const double *>(buf.ptr), buf.size));
return;
}
case DataType::VECTOR_INT8: {
self.target_.set_vector(serialize_vector<int8_t>(
static_cast<const int8_t *>(buf.ptr), buf.size));
return;
}
case DataType::VECTOR_FP16: {
self.target_.set_vector(serialize_vector<uint16_t>(
static_cast<const uint16_t *>(buf.ptr), buf.size));
return;
}
default:
throw py::type_error(
"Unsupported dense vector type for ndarray input: " +
std::to_string(static_cast<int>(data_type)));
}
}
// sparse vector
if (FieldSchema::is_sparse_vector_field(data_type)) {
if (!py::isinstance<py::dict>(obj)) {
throw py::type_error("Sparse vector[" + field_schema.name() +
"] expects a Python dict, got " +
std::string(py::str(py::type::of(obj))));
}
const auto sparse = obj.cast<py::dict>();
// Dense vector data is referenced by the query object. Callers
// must not modify the source data until the query returns.
if (FieldSchema::is_dense_vector_field(data_type)) {
if (!py::isinstance<py::array>(obj)) {
throw py::type_error("Dense vector[" + field_schema.name() +
"] expects a ndarray, got " +
std::string(py::str(py::type::of(obj))));
}
const auto arr = obj.cast<py::array>();
if (arr.ndim() != 1) {
throw py::type_error("Dense vector expects 1D array, got " +
std::to_string(arr.ndim()) + "D");
}
const auto buf = arr.request();
self.target_.clause_ = VectorViewClause{
std::string_view(
static_cast<const char *>(buf.ptr),
static_cast<size_t>(buf.size) * buf.itemsize),
{},
{}};
return;
}
// sparse vector
if (FieldSchema::is_sparse_vector_field(data_type)) {
if (!py::isinstance<py::dict>(obj)) {
throw py::type_error("Sparse vector[" + field_schema.name() +
"] expects a Python dict, got " +
std::string(py::str(py::type::of(obj))));
}
const auto sparse = obj.cast<py::dict>();
switch (data_type) {
case DataType::SPARSE_VECTOR_FP16: {
auto [indices, values] =
serialize_sparse_vector<ailego::Float16>(
sparse, [](const py::handle &h, size_t idx) {
float f = checked_cast<float>(
h, "Sparse value[" + std::to_string(idx) + "]",
"FLOAT");
return ailego::Float16(f);
});
self.target_.set_sparse_vector(std::move(indices),
std::move(values));
break;
}
case DataType::SPARSE_VECTOR_FP32: {
auto [indices, values] = serialize_sparse_vector<float>(
sparse, [](const py::handle &h, size_t idx) {
return checked_cast<float>(
h, "Sparse value[" + std::to_string(idx) + "]",
"FLOAT");
});
self.target_.set_sparse_vector(std::move(indices),
std::move(values));
break;
}
default:
throw py::type_error(
"Unsupported sparse vector type: " +
std::to_string(static_cast<int>(data_type)));
}
return;
}
switch (data_type) {
case DataType::SPARSE_VECTOR_FP16: {
auto [indices, values] =
serialize_sparse_vector<ailego::Float16>(
sparse, [](const py::handle &h, size_t idx) {
float f = checked_cast<float>(
h, "Sparse value[" + std::to_string(idx) + "]",
"FLOAT");
return ailego::Float16(f);
});
self.target_.set_sparse_vector(std::move(indices),
std::move(values));
break;
}
case DataType::SPARSE_VECTOR_FP32: {
auto [indices, values] = serialize_sparse_vector<float>(
sparse, [](const py::handle &h, size_t idx) {
return checked_cast<float>(
h, "Sparse value[" + std::to_string(idx) + "]",
"FLOAT");
});
self.target_.set_sparse_vector(std::move(indices),
std::move(values));
break;
}
default:
throw py::type_error(
"Unsupported sparse vector type: " +
std::to_string(static_cast<int>(data_type)));
}
return;
}
throw py::type_error("Unsupported vector field type for field: " +
field_schema.name());
})
throw py::type_error("Unsupported vector field type for field: " +
field_schema.name());
},
py::arg("field_schema"), py::arg("obj"), py::keep_alive<1, 3>(),
"Set query vector. Dense vector source data must not be modified "
"until the query finishes.")
.def(
"get_vector",
[](const SearchQuery &self,
const FieldSchema &field_schema) -> py::object {
DataType data_type = field_schema.data_type();
const VectorClause *vc = self.target_.get_vector_clause();
// get_vector_view() works for both VectorClause and
// VectorViewClause.
auto vv = self.target_.get_vector_view();
if (FieldSchema::is_dense_vector_field(data_type)) {
if (vc == nullptr || vc->query_vector_.empty()) {
if (!vv || vv->query_vector_.empty()) {
throw std::runtime_error("No dense vector has been set");
}
size_t byte_size = vc->query_vector_.size();
const void *data = vc->query_vector_.data();
size_t byte_size = vv->query_vector_.size();
const void *data = vv->query_vector_.data();
switch (data_type) {
case DataType::VECTOR_FP32: {
@ -2135,29 +2124,29 @@ void ZVecPyParams::bind_vector_query(py::module_ &m) {
}
}
if (FieldSchema::is_sparse_vector_field(data_type)) {
if (vc == nullptr || vc->sparse_indices_.empty()) {
if (!vv || vv->sparse_indices_.empty()) {
return py::dict();
}
// Deserialize indices: stored as uint32_t[]
size_t indices_byte_size = vc->sparse_indices_.size();
size_t indices_byte_size = vv->sparse_indices_.size();
if (indices_byte_size % sizeof(uint32_t) != 0) {
throw std::runtime_error(
"Sparse indices buffer size not aligned to uint32_t");
}
size_t n = indices_byte_size / sizeof(uint32_t);
const uint32_t *indices = reinterpret_cast<const uint32_t *>(
vc->sparse_indices_.data());
vv->sparse_indices_.data());
// Deserialize values
switch (data_type) {
case DataType::SPARSE_VECTOR_FP32: {
if (vc->sparse_values_.size() != n * sizeof(float)) {
if (vv->sparse_values_.size() != n * sizeof(float)) {
throw std::runtime_error(
"Sparse FP32 values buffer size mismatch");
}
const float *values = reinterpret_cast<const float *>(
vc->sparse_values_.data());
vv->sparse_values_.data());
py::dict result;
for (size_t i = 0; i < n; ++i) {
result[py::int_(indices[i])] = py::float_(values[i]);
@ -2165,12 +2154,12 @@ void ZVecPyParams::bind_vector_query(py::module_ &m) {
return result;
}
case DataType::SPARSE_VECTOR_FP16: {
if (vc->sparse_values_.size() != n * sizeof(uint16_t)) {
if (vv->sparse_values_.size() != n * sizeof(uint16_t)) {
throw std::runtime_error(
"Sparse FP16 values buffer size mismatch");
}
const uint16_t *raw_bits = reinterpret_cast<const uint16_t *>(
vc->sparse_values_.data());
vv->sparse_values_.data());
py::dict result;
for (size_t i = 0; i < n; ++i) {
float f = ailego::FloatHelper::ToFP32(raw_bits[i]);
@ -2190,18 +2179,18 @@ void ZVecPyParams::bind_vector_query(py::module_ &m) {
py::arg("field_schema"))
.def(py::pickle(
[](const SearchQuery &self) {
const VectorClause *vc = self.target_.get_vector_clause();
auto vv = self.target_.get_vector_view();
const auto *fc = self.target_.get_fts_clause();
return py::make_tuple(self.topk_, self.target_.field_name_,
vc ? vc->query_vector_ : std::string(),
vc ? vc->sparse_indices_ : std::string(),
vc ? vc->sparse_values_ : std::string(),
self.filter_, self.include_vector_,
self.output_fields_,
self.target_.query_params_
? py::cast(self.target_.query_params_)
: py::none(),
fc ? py::cast(*fc) : py::none());
return py::make_tuple(
self.topk_, self.target_.field_name_,
vv ? std::string(vv->query_vector_) : std::string(),
vv ? std::string(vv->sparse_indices_) : std::string(),
vv ? std::string(vv->sparse_values_) : std::string(),
self.filter_, self.include_vector_, self.output_fields_,
self.target_.query_params_
? py::cast(self.target_.query_params_)
: py::none(),
fc ? py::cast(*fc) : py::none());
},
[](py::tuple t) {
if (t.size() != 10)