[WIP] feat(diskann): add macOS support via kqueue-based async IO

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

@@ -180,7 +180,7 @@ jobs:
       # DiskAnn libaio round: install libaio and re-run tests 
       # ------------------------------------------------------------------ #
       - name: Install libaio runtime
-        if: matrix.platform == 'linux-x64'
+        if: matrix.platform == 'linux-x64' || matrix.platform == 'linux-arm64'
         run: |
           sudo apt-get update -y
           # libaio1t64 on Ubuntu 24.04+ (t64 transition), libaio1 on older
@@ -190,20 +190,35 @@ jobs:
           ldconfig -p | grep libaio || true
         shell: bash
 
+      - name: Verify DiskAnn uses libaio
+        if: matrix.platform == 'linux-x64' || matrix.platform == 'linux-arm64'
+        run: |
+          python - <<'PY'
+          import zvec
+          from zvec.typing import IOBackendType
+
+          backend = zvec.io_backend_type()
+          print(f"DiskAnn I/O backend: {backend.name}")
+          assert backend == IOBackendType.LIBAIO, (
+              f"expected LIBAIO after installing libaio, got {backend.name}"
+          )
+          PY
+        shell: bash
+
       - name: Run DiskAnn C++ Tests (w/ libaio)
-        if: matrix.platform == 'linux-x64'
+        if: matrix.platform == 'linux-x64' || matrix.platform == 'linux-arm64'
         run: |
           cd "$GITHUB_WORKSPACE/build"
-          ctest -R diskann --output-on-failure --parallel $NPROC
+          ctest -R diskann --no-tests=error --output-on-failure \
+            --parallel $NPROC
         shell: bash
 
       - name: Run DiskAnn Python Tests (w/ libaio)
-        if: matrix.platform == 'linux-x64'
+        if: matrix.platform == 'linux-x64' || matrix.platform == 'linux-arm64'
         run: |
           cd "$GITHUB_WORKSPACE"
           python -m pytest python/tests/test_collection_diskann.py -v
         shell: bash
-
       # Verify installing libaio does not affect existing non-DiskAnn tests.
       - name: Run HNSW Tests (w/ libaio)
         if: matrix.platform == 'linux-x64'

CMakeLists.txt

@@ -122,14 +122,17 @@ else()
 endif()
 message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}")
 
-# DiskAnn support (Linux x86_64 only; libaio loaded at runtime via dlopen)
-if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS)
+# DiskAnn support:
+#   - Linux (x86_64, i686, i386, aarch64, arm64) with libaio (loaded via dlopen)
+#   - macOS (x86_64, ARM64/Apple Silicon) with kqueue
+if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64" AND NOT ANDROID AND NOT IOS)
+    OR (CMAKE_SYSTEM_NAME STREQUAL "Darwin" 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")
+  message(STATUS "DiskAnn support disabled - supported on Linux (x86_64/ARM64 with libaio) and macOS (with kqueue)")
 endif()
 message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}")
 

examples/c++/CMakeLists.txt

@@ -15,6 +15,7 @@ endif()
 get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE)
 
 set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include)
+set(ZVEC_SRC_DIR ${ZVEC_ROOT_DIR}/src)
 set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib)
 
 include_directories(${ZVEC_INCLUDE_DIR})

python/tests/detail/fixture_helper.py

@@ -2,12 +2,10 @@
 import logging
 import platform
 
-DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in (
-    "x86_64",
-    "AMD64",
-    "i686",
-    "i386",
-)
+DISKANN_SUPPORTED = (
+    platform.system() == "Linux"
+    and platform.machine() in ("x86_64", "AMD64", "i686", "i386", "aarch64", "arm64")
+) or platform.system() == "Darwin"
 
 from typing import Any, Generator
 from zvec.typing import DataType, StatusCode, MetricType, QuantizeType
@@ -27,9 +25,8 @@ def _ensure_diskann_runtime_or_reason() -> str | None:
     _DISKANN_PRELOAD_DONE = True
 
     if not DISKANN_SUPPORTED:
-        _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64"
+        _DISKANN_PRELOAD_REASON = "DiskAnn is supported on Linux (x86_64/ARM64 with libaio) and macOS (kqueue)"
         return _DISKANN_PRELOAD_REASON
-
     _DISKANN_PRELOAD_REASON = None
     return None
 

python/tests/test_collection_diskann.py

@@ -17,8 +17,8 @@
 
 Two platform-level prerequisites are enforced at module import time:
 
-1. DiskAnn is currently built only for Linux x86_64 — other platforms are
-   skipped wholesale.
+1. DiskAnn is built for Linux (x86_64/ARM64 with libaio) and macOS (with kqueue) —
+   other platforms are skipped wholesale.
 2. libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() /
    DiskAnnStreamer::init(). If libaio is missing, DiskAnn falls back to
    synchronous pread() — the tests still run but with degraded performance.
@@ -39,8 +39,14 @@
 # Platform gating (must happen BEFORE we touch zvec).
 # --------------------------------------------------------------------------- #
 pytestmark = pytest.mark.skipif(
-    not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")),
-    reason="DiskAnn plugin is only supported on Linux x86_64",
+    not (
+        (
+            sys.platform == "linux"
+            and platform.machine() in ("x86_64", "AMD64", "aarch64", "arm64")
+        )
+        or sys.platform == "darwin"
+    ),
+    reason="DiskAnn is supported on Linux (x86_64/ARM64 with libaio) and macOS (kqueue)",
 )
 
 import zvec  # noqa: E402

python/tests/test_typing.py

@@ -49,6 +49,7 @@ def test_enum_names(member, name):
         (DataType.FLOAT, 8),
         (IndexType.HNSW, 1),
         (IOBackendType.PREAD, 0),
+        (IOBackendType.THREAD_POOL_PREAD, 2),
         (MetricType.COSINE, 3),
         (QuantizeType.INT8, 2),
         (StatusCode.OK, 0),
@@ -112,7 +113,9 @@ def test_index_type_has_member(member):
     assert member in IndexType.__members__
 
 
-@pytest.mark.parametrize("member", ["PREAD", "LIBAIO"])
+@pytest.mark.parametrize(
+    "member", ["PREAD", "LIBAIO", "THREAD_POOL_PREAD"]
+)
 def test_io_backend_type_has_member(member):
     assert member in IOBackendType.__members__
 

python/zvec/__init__.pyi

@@ -52,12 +52,13 @@ from .zvec import create_and_open, init, open
 def io_backend_type() -> IOBackendType:
     """Returns the current I/O backend type for DiskAnn async disk reads
     as an IOBackendType enum (zvec.typing.IOBackendType).
-    IOBackendType.LIBAIO if libaio is available, IOBackendType.PREAD otherwise."""
+    On Linux this is IOBackendType.LIBAIO when libaio is available; on macOS
+    this is IOBackendType.THREAD_POOL_PREAD."""
 
 def io_backend_description() -> str:
     """Returns a human-readable description of the current I/O backend.
-    When only pread is available, includes instructions for installing
-    libaio to enable async I/O."""
+    On Linux, when only pread is available, includes instructions for
+    installing libaio to enable async I/O."""
 
 def set_default_jieba_dict_dir(dir: str) -> None:
     """Register the process-wide default jieba dict directory."""

python/zvec/typing/__init__.pyi

@@ -130,6 +130,8 @@ class IOBackendType:
 
     - PREAD: Synchronous pread() — no async I/O.
     - LIBAIO: libaio loaded at runtime via dlopen().
+    - THREAD_POOL_PREAD: Blocking pread() executed by a shared worker pool,
+      with completion delivered through kqueue.
 
     Examples:
         >>> from zvec.typing import IOBackendType
@@ -142,13 +144,16 @@ class IOBackendType:
       PREAD
 
       LIBAIO
+
+      THREAD_POOL_PREAD
     """
 
     LIBAIO: typing.ClassVar[IOBackendType]  # value = <IOBackendType.LIBAIO: 1>
     PREAD: typing.ClassVar[IOBackendType]  # value = <IOBackendType.PREAD: 0>
+    THREAD_POOL_PREAD: typing.ClassVar[IOBackendType]  # value = <IOBackendType.THREAD_POOL_PREAD: 2>
     __members__: typing.ClassVar[
         dict[str, IOBackendType]
-    ]  # value = {'PREAD': <IOBackendType.PREAD: 0>, 'LIBAIO': <IOBackendType.LIBAIO: 1>}
+    ]  # value = {'PREAD': <IOBackendType.PREAD: 0>, 'LIBAIO': <IOBackendType.LIBAIO: 1>, 'THREAD_POOL_PREAD': <IOBackendType.THREAD_POOL_PREAD: 2>}
 
     def __eq__(self, other: typing.Any) -> bool: ...
     def __getstate__(self) -> int: ...

src/ailego/io/io_backend_def.h

@@ -19,7 +19,7 @@
 // operations are still performed by the underlying loaders; this class is
 // responsible only for backend initialization and reporting.
 //
-// When no async backend is available, the caller should fall back to
+// When no asynchronous backend is available, the caller should fall back to
 // synchronous pread().
 //
 // Usage:
@@ -40,6 +40,8 @@ inline const char *IOBackendTypeName(IOBackendType type) {
   switch (type) {
     case IOBackendType::kLibAio:
       return "libaio";
+    case IOBackendType::kThreadPoolPread:
+      return "thread_pool_pread";
     case IOBackendType::kPread:
       return "pread";
   }
@@ -52,6 +54,9 @@ inline const char *IOBackendDescription(IOBackendType type) {
   switch (type) {
     case IOBackendType::kLibAio:
       return "libaio async I/O backend loaded at runtime via dlopen().";
+    case IOBackendType::kThreadPoolPread:
+      return "Blocking pread() calls executed by a shared worker pool, with "
+             "completion delivered through kqueue.";
     case IOBackendType::kPread:
       return "No async I/O backend available. Install libaio (e.g. "
              "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) "
@@ -63,8 +68,8 @@ inline const char *IOBackendDescription(IOBackendType type) {
 
 // Singleton that loads and queries an I/O backend on demand.
 //
-// available() (no arg) tries the best backend with priority (libaio > pread)
-// and returns the loaded backend type.
+// available() (no arg) selects libaio on Linux when available, worker-pool
+// pread() on macOS, and synchronous pread() otherwise.
 // available(IOBackendType) tries a specific backend.
 // Use type() / name() to query the loaded backend without triggering a load.
 class IOBackend {
@@ -74,14 +79,18 @@ class IOBackend {
     return instance;
   }
 
-  // Try to load the best available backend (libaio > pread).
+  // Try to load the best available backend for the current platform.
   // Returns the loaded backend type.
   // Idempotent — if already loaded, returns immediately.
   IOBackendType available() {
     if (type_ != IOBackendType::kPread) {
       return type_;
     }
+#if defined(__APPLE__) && defined(__MACH__)
+    return available(IOBackendType::kThreadPoolPread);
+#else
     return available(IOBackendType::kLibAio);
+#endif
   }
 
   // Try to load the requested backend.  Returns the loaded backend type
@@ -91,6 +100,12 @@ class IOBackend {
     if (type_ == requested && type_ != IOBackendType::kPread) {
       return type_;
     }
+#if defined(__APPLE__) && defined(__MACH__)
+    if (requested == IOBackendType::kThreadPoolPread) {
+      type_ = IOBackendType::kThreadPoolPread;
+      return type_;
+    }
+#endif
 #if defined(__linux) || defined(__linux__)
     if (requested == IOBackendType::kLibAio) {
       if (LibAioLoader::Instance().load() &&
@@ -112,6 +127,10 @@ class IOBackend {
     return available() == IOBackendType::kLibAio;
   }
 
+  bool is_thread_pool_pread() {
+    return available() == IOBackendType::kThreadPoolPread;
+  }
+
   // Returns the loaded backend type.
   IOBackendType type() const {
     return type_;

src/binding/python/model/common/python_config.cc

@@ -228,18 +228,18 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) {
       },
       "Returns the current I/O backend type for DiskAnn async disk reads "
       "as an IOBackendType enum (zvec.typing.IOBackendType). "
-      "IOBackendType.LIBAIO if libaio is available, "
-      "IOBackendType.PREAD otherwise.");
+      "On Linux this is IOBackendType.LIBAIO when libaio is available; "
+      "on macOS this is IOBackendType.THREAD_POOL_PREAD.");
 
   // Returns a human-readable description of the I/O backend, including
-  // installation guidance for libaio when only pread is available.
+  // installation guidance for libaio on Linux when only pread is available.
   m.def(
       "io_backend_description",
       []() -> std::string { return ailego::current_io_backend_description(); },
       "Returns a human-readable description of the current I/O backend. "
-      "When only pread is available, includes instructions for installing "
-      "libaio to enable async I/O.");
+      "On Linux, when only pread is available, includes instructions for "
+      "installing libaio to enable async I/O.");
 }
 
 
-}  // namespace zvec
+}  // namespace zvec

src/binding/python/typing/python_type.cc

@@ -146,14 +146,17 @@ Enumeration of supported I/O backend types for DiskAnn async disk reads.
 
 - PREAD: Synchronous pread() \u2014 no async I/O.
 - LIBAIO: libaio loaded at runtime via dlopen().
+- THREAD_POOL_PREAD: Blocking pread() executed by a shared worker pool, with
+  completion delivered through kqueue.
 
 Examples:
     >>> from zvec.typing import IOBackendType
     >>> print(IOBackendType.LIBAIO)
     IOBackendType.LIBAIO
 )pbdoc")
       .value("PREAD", ailego::IOBackendType::kPread)
-      .value("LIBAIO", ailego::IOBackendType::kLibAio);
+      .value("LIBAIO", ailego::IOBackendType::kLibAio)
+      .value("THREAD_POOL_PREAD", ailego::IOBackendType::kThreadPoolPread);
 }
 
 void ZVecPyTyping::bind_status(py::module_ &m) {
@@ -245,4 +248,4 @@ Construct a status with the given code and optional message.
       });
 }
 
-}  // namespace zvec
+}  // namespace zvec

src/core/algorithm/CMakeLists.txt

@@ -17,7 +17,7 @@ else()
   # Empty stub library for unsupported platforms
   file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/diskann_stub.cc
     "// Stub implementation for unsupported platforms\n"
-    "// DiskAnn only supports Linux x86_64\n"
+    "// DiskAnn supports Linux (x86_64/ARM64 with libaio) and macOS (kqueue)\n"
     "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n"
   )
 

src/core/algorithm/diskann/CMakeLists.txt

@@ -14,7 +14,7 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c)
 # ${CMAKE_DL_LIBS} for dlopen/dlsym/dlclose.
 set(CORE_KNN_DISKANN_LIBS core_framework core_knn_cluster)
 
-if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386")
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64")
     list(APPEND CORE_KNN_DISKANN_LIBS ${CMAKE_DL_LIBS})
 endif()
 
@@ -31,4 +31,4 @@ cc_library(
     INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
     LDFLAGS "${CORE_KNN_DISKANN_LDFLAGS}"
     VERSION "${PROXIMA_ZVEC_VERSION}"
-)
+)

src/core/algorithm/diskann/diskann_context.h

@@ -349,7 +349,13 @@ class DiskAnnContext : public IndexContext,
   uint32_t group_num_{0};
   std::map<std::string, TopkHeap> group_topk_heaps_{};
 
-  IOContext io_ctx_{0};
+  IOContext io_ctx_{
+#if defined(__APPLE__) || defined(__MACH__)
+      -1
+#else
+      0
+#endif
+  };
   SearchStats query_stats_;
 
   float *pq_table_dist_buffer_{nullptr};