feat: add diskann index by richyreachy · Pull Request #369 · alibaba/zvec

richyreachy · 2026-04-22T12:08:50Z

Add diskann index into Zvec to lower memory usage in vector search as per the description: #325

… feat/diskann_index

feihongxu0824 · 2026-04-28T01:30:48Z

+                -Wl,--whole-archive
+                $<TARGET_FILE:core_knn_flat_static>
+                $<TARGET_FILE:core_knn_flat_sparse_static>
+                $<TARGET_FILE:core_knn_hnsw_static>


代码有重复，可以调整一下

feihongxu0824 · 2026-04-28T01:41:58Z

+        run: |
+          sudo apt-get update
+          sudo apt-get install -y --no-install-recommends \
+            libaio-dev


如果用户的环境没有装libaio-dev，会发生什么？

现在默认使用是需要安装libaio，可以通过配置的方式进行区分，千问的建议是通过linux安装包的方式安装libaio库：

Installation

zvec requires the libaio system library on linux platform.

On Ubuntu/Debian:

sudo apt-get install libaio1 libaio-dev pip install zvec

如果没有安装，会发生什么？这里预期的行为应该是如果用户不安装aio，不影响除diskann的其他功能使用

在新的PR里做了调整：#378

feihongxu0824 · 2026-04-28T01:56:43Z

+  }
+
+  auto &pool = ctx->expanded_nodes();
+  for (uint32_t i = 0; i < pool.size(); i++) {


可以使用std::remove_if + erase，效率高一些

feihongxu0824 · 2026-04-28T02:07:48Z

+
+  virtual ~DiskAnnQueryParams() = default;
+
+  int list_size() const {


加上参数的注释吧

feihongxu0824 · 2026-04-28T02:15:35Z

+
+    sector_internal_id_++;
+    if (sector_internal_id_ >= sector_vec_num_) {
+      std::vector<uint8_t> padding_(padding_size_, 0);


没有必要allocate一个临时的std::vector?

std::memset(data_ptr + data_size_, 0, padding_size_);

… feat/diskann_index

egolearner · 2026-05-27T09:15:05Z

    "proxima.stratified.trainer.cluster_params_in_level_";

+static const std::string MULTI_CHUNK_CLUSTER_COUNT =
+    "proxima.cluster.multi_chunk_cluster.count";


新代码要切换到zvec.xxx吗？

等一起统一调整，这个其它索引里还没有改掉

egolearner · 2026-05-27T09:18:55Z

+
+  //! Init
+  int init() {
+    // file_.open(path, std::ios::in | std::ios::out);


清理下注释掉的代码吧

egolearner · 2026-05-27T11:11:47Z

+  typedef std::shared_ptr<MultiChunkClusterAlgorithm> Pointer;
+
+  //! Constructor
+  MultiChunkClusterAlgorithm(void) {}


使用default，{}会影响编译器优化，析构函数类似

egolearner · 2026-05-27T11:13:24Z

+                "ValueType must be arithmetic");
+
+  //! Constructor
+  MultiChunkNumericalAlgorithm(void) {}


egolearner · 2026-05-27T11:16:41Z

+
+      bool result = algorithm.cluster_once(*local_threads, &cost);
+      if (result != true) {
+        LOG_ERROR("(%u) Failed to cluster.", i + 1);


%u打印会有编译warning，建议%zu (size_t)xxx打印

egolearner · 2026-05-27T11:17:26Z

+    (*finished)++;
+  }
+
+  return;


函数尾部的return可以去掉

egolearner · 2026-05-27T11:18:39Z

+  int cleanup(void);
+
+  //! Reset Cluster
+  int reset(void);


void可以去掉，c++里面reset()即可

egolearner · 2026-05-27T11:19:58Z

+        }
+      }
+
+      (*out)[id * chunk_count_ + chunk] = static_cast<uint32_t>(sel_index);


不用cast，sel_index本身是uint32的

egolearner · 2026-05-27T12:36:46Z

+  entity_.set_neighbors(id, pruned_list);
+  lock_pool_[lock_idx].unlock();
+
+  ret = inter_insert(id, pruned_list, ctx);


return inter_insert...否则会丢失错误

egolearner · 2026-05-27T12:37:44Z

+    return core::IndexError_Unsupported;
+  }
+
+  param_ = dynamic_cast<const DiskAnnIndexParam &>(param);


egolearner · 2026-05-27T12:38:31Z

+    return pq_chunk_num_;
+  }
+
+  void pq_chunk_num(int pq_chunk_num) {


其他的都带set，这个也加上吧

egolearner · 2026-05-27T12:40:59Z

+    20
+)pbdoc");
+  diskann_params
+      .def(py::init<int>(), py::arg("list_size") = 10, R"pbdoc(


c++默认为300，注释也写默认300，这里是10

egolearner · 2026-05-27T12:44:24Z

+        Default is ``MetricType.IP`` (inner product).
+    max_degree (int):.
+    list_size (int): .
+    pq_chunk_num (bool): .


pq_chunk_num (bool): —— 类型是 int，不是 bool。
示例里 >>> print(params.n_list) / >>> print(params.nprobe) 引用的属性 DiskAnn 类上根本不存在

… feat/diskann_index

egolearner · 2026-05-28T03:06:10Z

+               uint64_t n_retries = 0) {
+#if (defined(__linux) || defined(__linux__))
+  uint64_t iters =
+      DiskAnnUtil::round_up(read_reqs.size(), MAX_EVENTS) / MAX_EVENTS;


用div_round_up，省一次除

egolearner · 2026-05-28T03:15:47Z

+          break;
+        }
+      }
+      n_tries++;


retry不起作用，错误路径都返回或者break了

egolearner · 2026-05-28T03:19:36Z

+
+      if (ret != (int)n_ops) {
+        LOG_WARN(
+            "io_submit failed; returned: %d, expected=%lu, errno=%d, %s, "


libaio不会设置errno，直接打印ret吧

egolearner · 2026-05-28T12:44:07Z

-    )
+  if(MSVC)
+    # MSVC: STATIC-only stub to avoid creating an empty DLL with no exports
+    cc_library(


这里是有什么问题？之前ci也能跑？

调整和主干一致了

egolearner · 2026-05-28T12:46:42Z

+// every other index type stays fully functional.
+int EnsureDiskAnnRuntimeReady() {
+  const int status = ::zvec::LoadDiskAnnPlugin();
+  if (status == ::zvec::kDiskAnnPluginOk) {


zvec前面的::可以去掉？应该不太会冲突

src/binding/python/binding.cc 35: [](const std::string &path) { return ::zvec::LoadDiskAnnPlugin(path); }, 40: m.def("is_diskann_plugin_loaded", &::zvec::IsDiskAnnPluginLoaded, 42: m.def("is_libaio_available", &::zvec::IsLibAioAvailable, 48: m.attr("DISKANN_PLUGIN_OK") = static_cast<int>(::zvec::kDiskAnnPluginOk); 50: static_cast<int>(::zvec::kDiskAnnPluginUnsupportedPlatform); 52: static_cast<int>(::zvec::kDiskAnnPluginLibAioMissing); 54: static_cast<int>(::zvec::kDiskAnnPluginDlopenFailed); src/core/plugin/diskann_plugin.cc 88: if (::dladdr(reinterpret_cast<void *>(&::zvec::LoadDiskAnnPlugin), &info) == 105: if (::dladdr(reinterpret_cast<void *>(&::zvec::LoadDiskAnnPlugin), &info) == src/core/interface/indexes/diskann_index.cc 34: const int status = ::zvec::LoadDiskAnnPlugin(); 35: if (status == ::zvec::kDiskAnnPluginOk) { 39: case ::zvec::kDiskAnnPluginLibAioMissing: 46: case ::zvec::kDiskAnnPluginUnsupportedPlatform: 51: case ::zvec::kDiskAnnPluginDlopenFailed:

egolearner · 2026-05-28T12:56:24Z

+      LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str());
+      return core::IndexError_Runtime;
+    }
+    is_trained_ = true;


设置is_trained_后面add直接报错，这里创建streamer的作用是什么？

这个streamer是提供readonly查询的，不提供add能力，和ivf保持一致

egolearner · 2026-05-28T12:59:08Z

+    auto holder =
+        std::make_shared<zvec::core::MultiPassIndexHolder<DataType::DT_FP16>>(
+            param_.dimension);
+    for (auto doc : doc_cache_) {


ivf_index.cc也有类似实现，抽个公共实现吧

egolearner · 2026-05-28T13:01:39Z

+      input_vector_meta_.dimension() * input_vector_meta_.unit_size());
+
+  std::lock_guard<std::mutex> lock(mutex_);
+  while (doc_cache_.size() <= doc_id) {


已经上锁了，可以直接

if (size() <= doc_id) { doc_cache_.resize(doc+1, {kInvalidKey, fake_data}); }

egolearner · 2026-05-28T13:05:01Z

+  auto dumper = core::IndexFactory::CreateDumper("FileDumper");
+
+  dumper->create(file_path_);
+  builder_->dump(dumper);


需要检查错误？

egolearner · 2026-05-28T13:08:38Z


  int flush(void) override {
-    return IndexError_NotImplemented;
+    return 0;


只读的场景什么时候需要flush?

这里因为index调用到了flush，所以还是做了0值返回，保证接口的一致。

feihongxu0824 · 2026-05-29T03:31:23Z

+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ZVEC_PLUGIN_DISKANN_PLUGIN_H


用#pragma once吧

feihongxu0824 · 2026-05-29T03:34:40Z

+# --------------------------------------------------------------------------- #
+# Plugin gating.
+# --------------------------------------------------------------------------- #
+_PLUGIN_LOAD_STATUS = zvec.load_diskann_plugin()


这里最好不要让用户手动调用一次吧，import zvec内可以自动调用load_diskann_plugin

这个load_diskann_plugin不是必须调用，可以注释掉，或者去掉，会做隐式加载。如果没有load成功，会报错。

... ...
def main() -> int:
# ---- 1. Load the DiskAnn plugin. -------------------------------------- #
#status = zvec.load_diskann_plugin()
#if status != 0:
# print(
# f"[error] DiskAnn plugin failed to load (status={status}). "
# "Make sure libzvec_diskann_plugin.so sits next to _zvec.so.",
# file=sys.stderr,
# )
# return 1
#print("[ok] DiskAnn plugin loaded.")

# ---- 2. Prepare a throwaway workspace and the collection. ------------- # ...

feat: add diskann index

e38efb9

richyreachy requested review from Cuiyus, JalinWang, chinaux, feihongxu0824, iaojnh and zhourrr as code owners April 22, 2026 12:08

github-actions Bot assigned richyreachy Apr 22, 2026

richyreachy added 12 commits April 23, 2026 13:58

fix: add libaio dependency

43c9db5

Merge branch 'main' into feat/diskann_index

043456d

fix: fix initialization

dae465c

Merge branch 'feat/diskann_index' of github.com:richyreachy/zvec into…

1202e42

… feat/diskann_index

Merge branch 'main' into feat/diskann_index

d9503d3

refactor: fix ut

c24cd82

Merge branch 'feat/diskann_index' of github.com:richyreachy/zvec into…

87a1b8c

… feat/diskann_index

fix: fix seed value

4e668da

fix: fix seed value

fe60376

fix: update ut

dc3d0a8

Merge branch 'main' into feat/diskann_index

071d275

fix: fix ut

78eea04

feihongxu0824 reviewed Apr 28, 2026

View reviewed changes

Comment thread src/core/algorithm/diskann/CMakeLists.txt Outdated

feihongxu0824 reviewed Apr 28, 2026

View reviewed changes

Comment thread src/include/zvec/db/query_params.h

richyreachy added 7 commits April 28, 2026 20:25

Merge branch 'main' into feat/diskann_index

41f893d

fix: fix according to pr comments

4cc1ecb

Merge branch 'feat/diskann_index' of github.com:richyreachy/zvec into…

35192e7

… feat/diskann_index

fix: fix on comments

53f9ef9

fix: fix yaml

f3ee4ca

feat: support dynamic load

fd71a06

feat: add dynamic load

17570d7

richyreachy added 5 commits May 27, 2026 15:40

fix: fix ut

764465e

Merge branch 'main' into feat/diskann_index

7dbc8b9

fix: add provider

b477a05

Merge branch 'feat/diskann_index' of github.com:richyreachy/zvec into…

9a59894

… feat/diskann_index

fix: fix clang tidy

f908213

egolearner reviewed May 27, 2026

View reviewed changes

richyreachy added 4 commits May 28, 2026 10:18

Merge branch 'main' into feat/diskann_index

e155770

fix: fix ut

20ec830

Merge branch 'feat/diskann_index' of github.com:richyreachy/zvec into…

1af68bc

… feat/diskann_index

fix: fix ut

8a513d1

iaojnh requested changes May 28, 2026

View reviewed changes

Comment thread src/core/algorithm/diskann/diskann_searcher.h

Comment thread src/core/algorithm/diskann/CMakeLists.txt

Comment thread src/core/algorithm/diskann/diskann_builder.cc Outdated

richyreachy added 5 commits May 28, 2026 19:27

fix: update code

aaa3b8a

fix: update code

c9238f9

fix: remove unused comments

15ea014

fix: update code

f42f4ec

fix: fix code according to comments

e64912a

egolearner reviewed May 28, 2026

View reviewed changes

feihongxu0824 reviewed May 29, 2026

View reviewed changes

richyreachy added 11 commits May 29, 2026 14:26

fix: update code

f8f950d

fix: update code

8044f8f

fix: update code

6513565

fix: fix cmake

4f1ffb4

fix: fix update

91fb9e3

fix: update code

60ba001

Merge branch 'main' into feat/diskann_index

34cadcc

fix: add file

29e5f48

Merge branch 'main' into feat/diskann_index

7dd5898

Merge branch 'main' into feat/diskann_index

812d9d3

fix: update code

9626d92


		virtual ~DiskAnnQueryParams() = default;

		int list_size() const {

Conversation

richyreachy commented Apr 22, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Installation

On Ubuntu/Debian:

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!