diff --git a/src/core/algorithm/diskann/diskann_context.cc b/src/core/algorithm/diskann/diskann_context.cc index 44cff75..f13affb 100644 --- a/src/core/algorithm/diskann/diskann_context.cc +++ b/src/core/algorithm/diskann/diskann_context.cc @@ -14,6 +14,7 @@ #include "diskann_context.h" #include +#include "diskann_params.h" #include "diskann_pq_table.h" #include "diskann_util.h" @@ -93,6 +94,13 @@ DiskAnnContext::~DiskAnnContext() { } } +int DiskAnnContext::update(const ailego::Params ¶ms) { + uint32_t list_size = list_size_; + params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size); + list_size_ = list_size; + return 0; +} + int DiskAnnContext::update_context(ContextType type, const IndexMeta &meta, const IndexMetric::Pointer &measure, const DiskAnnEntity::Pointer &entity, diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index ff3863b..f8a736c 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -118,9 +118,7 @@ class DiskAnnContext : public IndexContext, } //! Update the parameters of context - virtual int update(const ailego::Params & /*params*/) override { - return 0; - } + int update(const ailego::Params ¶ms) override; inline DistCalculator &dist_calculator() { return dc_; diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 3ef0968..7615ed3 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -271,6 +271,14 @@ int DiskAnnIndex::_prepare_for_search( context->set_topk(diskann_search_param->topk); + // Propagate the query-time beam-search list size into the context. Must be + // at least topk to keep enough candidates for a correct result. + ailego::Params params; + params.set( + core::PARAM_DISKANN_SEARCHER_LIST_SIZE, + std::max(diskann_search_param->topk, diskann_search_param->list_size)); + context->update(params); + return 0; } diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index 29569dc..27dd0f1 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -212,7 +212,15 @@ class ProximaEngineHelper { "failed to build query param: " + diskann_query_param_result.error().message())); } - return std::move(diskann_query_param_result.value()); + auto &diskann_query_param = diskann_query_param_result.value(); + if (query_params.query_params) { + auto db_diskann_query_params = + dynamic_cast( + query_params.query_params.get()); + diskann_query_param->list_size = + static_cast(db_diskann_query_params->list_size()); + } + return std::move(diskann_query_param); } case IndexType::VAMANA: { diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 491d716..f14ebef 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -219,6 +219,10 @@ struct IVFQueryParam : public BaseIndexQueryParam { struct DiskAnnQueryParam : public BaseIndexQueryParam { using Pointer = std::shared_ptr; + // Beam-search candidate list size used at query time. Larger values improve + // recall at the cost of latency. + uint32_t list_size = kDefaultDiskAnnListSize; + BaseIndexQueryParam::Pointer Clone() const override { return std::make_shared(*this); }