feat: support to pass in list size (#487)

src/core/algorithm/diskann/diskann_context.cc

@@ -14,6 +14,7 @@
 
 #include "diskann_context.h"
 #include <chrono>
+#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 &params) {
+  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,

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 &params) override;
 
   inline DistCalculator &dist_calculator() {
     return dc_;

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

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<const DiskAnnQueryParams *>(
+                  query_params.query_params.get());
+          diskann_query_param->list_size =
+              static_cast<uint32_t>(db_diskann_query_params->list_size());
+        }
+        return std::move(diskann_query_param);
       }
 
       case IndexType::VAMANA: {

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<DiskAnnQueryParam>;
 
+  // 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<DiskAnnQueryParam>(*this);
   }