From 55a3c0e85314482dd4ebb88329d69fc10a80bdf4 Mon Sep 17 00:00:00 2001 From: "Dr.A" Date: Tue, 14 Apr 2026 14:34:49 +0200 Subject: [PATCH] feat: make max query topk configurable The native engine enforces a hard topk ceiling of 1024 on collection.query(). This adds a `max_query_topk` parameter to zvec.init() so users can raise the limit when they accept the extra memory/latency tradeoff. Changes: - Add `max_query_topk` to GlobalConfig::ConfigData with default 1024 - VectorQuery::validate() reads the limit from GlobalConfig instead of the hardcoded kMaxQueryTopk constant - Expose the parameter in Python zvec.init() with type/range validation - Parse and validate in the pybind11 config bridge - Add C API getter/setter: zvec_config_data_{set,get}_max_query_topk() Co-Authored-By: Claude Opus 4.6 (1M context) --- python/zvec/zvec.py | 12 +++++++++++ src/binding/c/c_api.cc | 21 +++++++++++++++++++ .../python/model/common/python_config.cc | 7 +++++++ src/db/common/config.cc | 6 ++++++ src/db/index/common/doc.cc | 6 ++++-- src/include/zvec/c_api.h | 18 ++++++++++++++++ src/include/zvec/db/config.h | 6 ++++++ 7 files changed, 74 insertions(+), 2 deletions(-) diff --git a/python/zvec/zvec.py b/python/zvec/zvec.py index 114fb49c9..45bc60d61 100644 --- a/python/zvec/zvec.py +++ b/python/zvec/zvec.py @@ -39,6 +39,7 @@ def init( invert_to_forward_scan_ratio: Optional[float] = None, brute_force_by_keys_ratio: Optional[float] = None, memory_limit_mb: Optional[int] = None, + max_query_topk: Optional[int] = None, ) -> None: """Initialize Zvec with configuration options. @@ -93,6 +94,11 @@ def init( approaching this limit. If ``None``, inferred from cgroup memory limit * 0.8 (e.g., in Docker). Must be > 0 if provided. + max_query_topk (Optional[int], optional): + Maximum allowed ``topk`` value in ``collection.query()``. + Default: ``1024``. Raise this when you need to retrieve more + results per query and accept the extra memory/latency cost. + Must be ≥ 1 if provided. Raises: RuntimeError: If Zvec is already initialized. @@ -159,6 +165,12 @@ def init( config_dict["brute_force_by_keys_ratio"] = brute_force_by_keys_ratio if memory_limit_mb is not None: config_dict["memory_limit_mb"] = memory_limit_mb + if max_query_topk is not None: + if not isinstance(max_query_topk, int) or isinstance(max_query_topk, bool): + raise TypeError("max_query_topk must be an integer") + if max_query_topk < 1: + raise ValueError("max_query_topk must be >= 1") + config_dict["max_query_topk"] = max_query_topk Initialize(config_dict) diff --git a/src/binding/c/c_api.cc b/src/binding/c/c_api.cc index 9144836b3..bb6ae3d08 100644 --- a/src/binding/c/c_api.cc +++ b/src/binding/c/c_api.cc @@ -648,6 +648,27 @@ uint32_t zvec_config_data_get_optimize_thread_count( return cpp_config->optimize_thread_count; } +zvec_error_code_t zvec_config_data_set_max_query_topk( + zvec_config_data_t *config, uint32_t max_topk) { + if (!config) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Config pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_config = reinterpret_cast(config); + cpp_config->max_query_topk = max_topk; + return ZVEC_OK; +} + +uint32_t zvec_config_data_get_max_query_topk( + const zvec_config_data_t *config) { + if (!config) { + return 1024; + } + auto *cpp_config = + reinterpret_cast(config); + return cpp_config->max_query_topk; +} + // ============================================================================= // Initialization and cleanup interface implementation diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc index bbcbb5bdb..9faedc644 100644 --- a/src/binding/python/model/common/python_config.cc +++ b/src/binding/python/model/common/python_config.cc @@ -177,6 +177,13 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { data.brute_force_by_keys_ratio = static_cast(v); } + // set max_query_topk + if (has_key(config_dict, "max_query_topk")) { + auto v = get_if(config_dict, "max_query_topk").value(); + if (v <= 0) throw py::value_error("max_query_topk must be positive"); + data.max_query_topk = static_cast(v); + } + // initialize (contains validate) Status status = GlobalConfig::Instance().Initialize(data); if (!status.ok()) { diff --git a/src/db/common/config.cc b/src/db/common/config.cc index 5938f5375..61bd419ca 100644 --- a/src/db/common/config.cc +++ b/src/db/common/config.cc @@ -37,6 +37,7 @@ GlobalConfig::ConfigData::ConfigData() query_thread_count(CgroupUtil::getCpuLimit()), invert_to_forward_scan_ratio(0.9), brute_force_by_keys_ratio(0.1), + max_query_topk(kMaxQueryTopk), optimize_thread_count(CgroupUtil::getCpuLimit()) {} Status GlobalConfig::Validate(const ConfigData &config) const { @@ -69,6 +70,11 @@ Status GlobalConfig::Validate(const ConfigData &config) const { "brute_force_by_keys_ratio must be between 0 and 1"); } + // Validate max_query_topk + if (config.max_query_topk == 0) { + return Status::InvalidArgument("max_query_topk must be greater than 0"); + } + // Validate optimize thread count if (config.optimize_thread_count == 0) { return Status::InvalidArgument( diff --git a/src/db/index/common/doc.cc b/src/db/index/common/doc.cc index 35271f36d..1c672ffad 100644 --- a/src/db/index/common/doc.cc +++ b/src/db/index/common/doc.cc @@ -18,6 +18,7 @@ #include #include #include +#include #include #include "db/common/constants.h" #include "db/index/common/type_helper.h" @@ -1202,9 +1203,10 @@ bool Doc::operator==(const Doc &other) const { } Status VectorQuery::validate(const FieldSchema *schema) const { - if ((uint32_t)topk_ > kMaxQueryTopk) { + uint32_t max_topk = GlobalConfig::Instance().max_query_topk(); + if ((uint32_t)topk_ > max_topk) { return Status::InvalidArgument("query validate failed: topk[", topk_, - "] is too large, max is ", kMaxQueryTopk); + "] is too large, max is ", max_topk); } if (output_fields_.has_value() && output_fields_->size() > kMaxOutputFieldSize) { diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index af21729ed..7c5331a39 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -697,6 +697,24 @@ zvec_config_data_set_optimize_thread_count(zvec_config_data_t *config, ZVEC_EXPORT uint32_t ZVEC_CALL zvec_config_data_get_optimize_thread_count(const zvec_config_data_t *config); +/** + * @brief Set max query topk in configuration data + * @param config Configuration data pointer + * @param max_topk Maximum allowed topk value for queries + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL +zvec_config_data_set_max_query_topk(zvec_config_data_t *config, + uint32_t max_topk); + +/** + * @brief Get max query topk from configuration data + * @param config Configuration data pointer + * @return uint32_t Max query topk + */ +ZVEC_EXPORT uint32_t ZVEC_CALL +zvec_config_data_get_max_query_topk(const zvec_config_data_t *config); + // ============================================================================= // Initialization and Cleanup Interface // ============================================================================= diff --git a/src/include/zvec/db/config.h b/src/include/zvec/db/config.h index 29fe19674..533b35215 100644 --- a/src/include/zvec/db/config.h +++ b/src/include/zvec/db/config.h @@ -92,6 +92,7 @@ class GlobalConfig : public ailego::Singleton { uint32_t query_thread_count; float invert_to_forward_scan_ratio; float brute_force_by_keys_ratio; + uint32_t max_query_topk; // optimize uint32_t optimize_thread_count; @@ -161,6 +162,11 @@ class GlobalConfig : public ailego::Singleton { return config_.brute_force_by_keys_ratio; } + //! Max query topk + uint32_t max_query_topk() const noexcept { + return config_.max_query_topk; + } + //! Optimize thread count uint32_t optimize_thread_count() const noexcept { return config_.optimize_thread_count;