feat: add diskann index

.github/workflows/03-macos-linux-build.yml

@@ -60,6 +60,14 @@ jobs:
           sudo apt-get install -y clang libomp-dev
         shell: bash
 
+      - name: Install AIO 
+        if: runner.os == 'Linux' && runner.arch == 'X64'
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y --no-install-recommends \
+            libaio-dev
+        shell: bash
+
       - name: Print CPU info
         if: runner.os == 'Linux'
         run: lscpu
@@ -96,7 +104,7 @@ jobs:
             pytest-xdist \
             scikit-build-core \
             setuptools_scm
-        shell: bash
+        shell: bash       
 
       - name: Build from source
         run: |

.github/workflows/clang_tidy.yml

@@ -47,7 +47,7 @@ jobs:
         if: steps.changed_files.outputs.any_changed == 'true'
         run: |
           sudo apt-get update
-          sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build libomp-dev
+          sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build libomp-dev libaio-dev
 
       - name: Setup ccache
         if: steps.changed_files.outputs.any_changed == 'true'

.gitmodules

@@ -40,3 +40,6 @@
 [submodule "thirdparty/RaBitQ-Library/RaBitQ-Library-0.1"]
 	path = thirdparty/RaBitQ-Library/RaBitQ-Library-0.1
 	url = https://github.com/VectorDB-NTU/RaBitQ-Library.git
+[submodule "thirdparty/aio/libaio-0.3"]
+	path = thirdparty/aio/libaio-0.3
+	url = https://github.com/yugabyte/libaio.git

CMakeLists.txt

@@ -109,6 +109,17 @@ else()
 endif()
 message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}")
 
+# DiskAnn support (Linux x86_64 only, requires libaio)
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS)
+  set(DISKANN_SUPPORTED ON)
+  add_definitions(-DDISKANN_SUPPORTED=1)
+else()
+  set(DISKANN_SUPPORTED OFF)
+  add_definitions(-DDISKANN_SUPPORTED=0)
+  message(STATUS "DiskAnn support disabled - only supported on Linux x86_64")
+endif()
+message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}")
+
 option(USE_OSS_MIRROR "Use OSS mirror for faster third-party downloads" OFF)
 if(DEFINED ENV{USE_OSS_MIRROR} AND NOT "$ENV{USE_OSS_MIRROR}" STREQUAL "")
   set(USE_OSS_MIRROR "$ENV{USE_OSS_MIRROR}" CACHE BOOL "Use OSS mirror for faster third-party downloads" FORCE)
@@ -145,4 +156,16 @@ if(BUILD_PYTHON_BINDINGS)
 
     message(STATUS "Zvec install path: ${ZVEC_PY_INSTALL_DIR}")
     install(TARGETS _zvec LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR})
+
+    # DiskAnn ships as a runtime-loaded shared module
+    # (libzvec_diskann_plugin.so) that is brought online implicitly the
+    # first time a DiskAnn index is created — users never call any load
+    # function. The Python extension resolves the module next to _zvec.so
+    # (see the $ORIGIN rpath in src/binding/python/CMakeLists.txt); the
+    # module must therefore be installed alongside _zvec.so in the wheel.
+    # The target exists only on platforms where DiskAnn is buildable
+    # (currently Linux x86_64 with libaio).
+    if(TARGET core_knn_diskann)
+        install(TARGETS core_knn_diskann LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR})
+    endif()
 endif()

python/tests/detail/fixture_helper.py

@@ -1,9 +1,60 @@
 import pytest
 import logging
+import platform
+
+DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in (
+    "x86_64",
+    "AMD64",
+    "i686",
+    "i386",
+)
 
 from typing import Any, Generator
 from zvec.typing import DataType, StatusCode, MetricType, QuantizeType
 import zvec
+
+
+# Cache the DiskAnn plugin preload status so we pay the load cost once per
+# test session. The plugin normally auto-loads on first DiskAnn use, but we
+# preload it explicitly here so a missing libaio / misplaced plugin .so
+# surfaces as a clear pytest skip instead of a confusing
+# "Create vector column indexer failed" deep inside the collection code path.
+_DISKANN_PRELOAD_REASON: str | None = None
+_DISKANN_PRELOAD_DONE: bool = False
+
+
+def _ensure_diskann_runtime_or_reason() -> str | None:
+    """Preload the DiskAnn plugin and return None on success or a human-readable
+    skip reason on failure. Idempotent across calls."""
+    global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON
+    if _DISKANN_PRELOAD_DONE:
+        return _DISKANN_PRELOAD_REASON
+    _DISKANN_PRELOAD_DONE = True
+
+    if not DISKANN_SUPPORTED:
+        _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64"
+        return _DISKANN_PRELOAD_REASON
+
+    if not zvec.is_libaio_available():
+        _DISKANN_PRELOAD_REASON = (
+            "libaio is not available on this host; DiskAnn cannot run. "
+            "Install libaio1 (or libaio1t64 on Ubuntu 24.04+) and retry."
+        )
+        return _DISKANN_PRELOAD_REASON
+
+    status = zvec.load_diskann_plugin()
+    if status != zvec.DISKANN_PLUGIN_OK:
+        _DISKANN_PRELOAD_REASON = (
+            f"Failed to load DiskAnn plugin (status={status}); "
+            "check that libzvec_diskann_plugin.so is installed alongside "
+            "_zvec.so in the Python site-packages directory."
+        )
+        return _DISKANN_PRELOAD_REASON
+
+    _DISKANN_PRELOAD_REASON = None
+    return None
+
+
 from zvec import (
     CollectionOption,
     InvertIndexParam,
@@ -97,6 +148,15 @@ def full_schema_new(request) -> CollectionSchema:
     else:
         nullable, has_index, vector_index = True, False, HnswIndexParam()
 
+    # Skip DiskAnn tests on unsupported platforms or when the runtime cannot
+    # be brought up (missing libaio, plugin .so not installed, etc.).
+    from zvec.model.param import DiskAnnIndexParam
+
+    if isinstance(vector_index, DiskAnnIndexParam):
+        skip_reason = _ensure_diskann_runtime_or_reason()
+        if skip_reason is not None:
+            pytest.skip(skip_reason)
+
     scalar_index_param = None
     vector_index_param = None
     if has_index:

python/tests/detail/test_collection_recall.py

@@ -22,8 +22,10 @@
     HnswIndexParam,
     FlatIndexParam,
     IVFIndexParam,
+    DiskAnnIndexParam,
     HnswQueryParam,
     IVFQueryParam,
+    DiskAnnQueryParam,
 )
 
 from zvec.model.schema import FieldSchema, VectorSchema
@@ -179,10 +181,24 @@ def get_ground_truth_map(collection, test_docs, query_vectors_map, metric_type,
     for field_name, query_vectors in query_vectors_map.items():
         ground_truth_map[field_name] = {}
 
+        # Support per-field metric type: metric_type can be a dict mapping
+        # field_name -> MetricType, or a single MetricType applied to all fields.
+        if isinstance(metric_type, dict):
+            field_metric = metric_type.get(field_name, MetricType.IP)
+        else:
+            field_metric = metric_type
+
         for i, query_vector in enumerate(query_vectors):
             # Get the ground truth for this query
             relevant_doc_ids_scores = get_ground_truth_for_vector_query(
-                collection, query_vector, field_name, test_docs, i, metric_type, k, True
+                collection,
+                query_vector,
+                field_name,
+                test_docs,
+                i,
+                field_metric,
+                k,
+                True,
             )
             ground_truth_map[field_name][i] = relevant_doc_ids_scores
 
@@ -292,6 +308,7 @@ class TestRecall:
         [
             (True, True, HnswIndexParam()),
             (False, True, IVFIndexParam()),
+            (False, True, DiskAnnIndexParam()),
             (False, True, FlatIndexParam()),  # ——ok
             (
                 True,
@@ -371,6 +388,19 @@ class TestRecall:
                     use_soar=False,
                 ),
             ),
+            (
+                True,
+                True,
+                DiskAnnIndexParam(
+                    metric_type=MetricType.IP,
+                    max_degree=32,
+                ),
+            ),
+            (
+                True,
+                True,
+                DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32),
+            ),
         ],
         indirect=True,
     )
@@ -388,10 +418,16 @@ def test_recall_with_single_vector_valid_500(
     ):
         full_schema_params = request.getfixturevalue("full_schema_new")
 
+        # Build per-field metric type map so ground truth uses each field's
+        # actual index metric (fields may fall back to HnswIndexParam/IP).
+        field_metric_map = {}
         for vector_para in full_schema_params.vectors:
-            if vector_para.name == "vector_fp32_field":
-                metric_type = vector_para.index_param.metric_type
-                break
+            if vector_para.index_param is not None:
+                field_metric_map[vector_para.name] = vector_para.index_param.metric_type
+            else:
+                field_metric_map[vector_para.name] = MetricType.IP
+
+        metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP)
 
         multiple_docs = [
             generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num)
@@ -438,9 +474,13 @@ def test_recall_with_single_vector_valid_500(
                 multiple_docs[i].vectors[field_name] for i in range(query_num)
             ]
 
-        # Get ground truth mapping
+        # Get ground truth mapping (pass per-field metric map)
         ground_truth_map = get_ground_truth_map(
-            full_collection_new, multiple_docs, query_vectors_map, metric_type, top_k
+            full_collection_new,
+            multiple_docs,
+            query_vectors_map,
+            field_metric_map,
+            top_k,
         )
 
         # Validate ground truth mapping structure
@@ -479,8 +519,8 @@ def test_recall_with_single_vector_valid_500(
 
         print("(recall_at_k_stats:\n")
         print(recall_at_k_stats)
-        print("metric_type:")
-        print(metric_type)
+        print("field_metric_map:")
+        print(field_metric_map)
         # Print Recall@K statistics
         print(f"Recall@{top_k} using Ground Truth:")
         for field_name, stats in recall_at_k_stats.items():
@@ -552,7 +592,21 @@ def test_recall_with_single_vector_valid_500(
                     use_soar=True,
                 ),
             ),
-            # (True, True, IVFIndexParam(metric_type=MetricType.COSINE, n_list=150, n_iters=15, use_soar=False, )),
+            (
+                True,
+                True,
+                DiskAnnIndexParam(metric_type=MetricType.IP, max_degree=32),
+            ),
+            (
+                True,
+                True,
+                DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32),
+            ),
+            (
+                True,
+                True,
+                DiskAnnIndexParam(metric_type=MetricType.COSINE, max_degree=32),
+            ),
         ],
         indirect=True,
     )
@@ -571,10 +625,16 @@ def test_recall_with_single_vector_valid_2000(
     ):
         full_schema_params = request.getfixturevalue("full_schema_new")
 
+        # Build per-field metric type map so ground truth uses each field's
+        # actual index metric (fields may fall back to HnswIndexParam/IP).
+        field_metric_map = {}
         for vector_para in full_schema_params.vectors:
-            if vector_para.name == "vector_fp32_field":
-                metric_type = vector_para.index_param.metric_type
-                break
+            if vector_para.index_param is not None:
+                field_metric_map[vector_para.name] = vector_para.index_param.metric_type
+            else:
+                field_metric_map[vector_para.name] = MetricType.IP
+
+        metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP)
 
         multiple_docs = [
             generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num)
@@ -621,9 +681,13 @@ def test_recall_with_single_vector_valid_2000(
                 multiple_docs[i].vectors[field_name] for i in range(query_num)
             ]
 
-        # Get ground truth mapping
+        # Get ground truth mapping (pass per-field metric map)
         ground_truth_map = get_ground_truth_map(
-            full_collection_new, multiple_docs, query_vectors_map, metric_type, top_k
+            full_collection_new,
+            multiple_docs,
+            query_vectors_map,
+            field_metric_map,
+            top_k,
         )
 
         # Validate ground truth mapping structure
@@ -662,8 +726,8 @@ def test_recall_with_single_vector_valid_2000(
 
         print("(recall_at_k_stats:\n")
         print(recall_at_k_stats)
-        print("metric_type:")
-        print(metric_type)
+        print("field_metric_map:")
+        print(field_metric_map)
         # Print Recall@K statistics
         print(f"Recall@{top_k} using Ground Truth:")
         for field_name, stats in recall_at_k_stats.items():