diff --git a/CMakeLists.txt b/CMakeLists.txt index 956ee5599..cbf9fdbb4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -122,14 +122,15 @@ 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 or Windows x86/x86_64). +if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS) + OR (WIN32 AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64|amd64|i686|i386")) 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 - only supported on Linux or Windows x86/x86_64") endif() message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}") diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index 4e5c703e1..2818259b7 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -17,7 +17,10 @@ get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) -include_directories(${ZVEC_INCLUDE_DIR}) +# Add include and library search paths +# ${ZVEC_ROOT_DIR}/src is needed because public headers under src/include/ +# reference ailego/ headers (e.g. ) that live in src/. +include_directories(${ZVEC_INCLUDE_DIR} ${ZVEC_ROOT_DIR}/src) set(ZVEC_LIB_SEARCH_DIRS ${ZVEC_LIB_DIR}) # Support multi-config builds (MSVC puts libs in Debug/Release subdirectories) diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index f5e6c08e4..da8872c12 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -2,7 +2,10 @@ import logging import platform -DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in ( +DISKANN_SUPPORTED = platform.system() in ( + "Linux", + "Windows", +) and platform.machine() in ( "x86_64", "AMD64", "i686", @@ -27,7 +30,9 @@ 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 only supported on Linux or Windows x86/x86_64" + ) return _DISKANN_PRELOAD_REASON _DISKANN_PRELOAD_REASON = None diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index 8daf8c770..7a9764827 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -15,11 +15,11 @@ Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn index. -Two platform-level prerequisites are enforced at module import time: +Platform prerequisites are enforced at module import time: -1. DiskAnn is currently built only for Linux x86_64 — other platforms are +1. DiskAnn is built for Linux and Windows x86/x86_64; other platforms are skipped wholesale. -2. libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() / +2. On Linux, 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,11 @@ # 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 in ("linux", "win32") + and platform.machine() in ("x86_64", "AMD64", "i686", "i386") + ), + reason="DiskAnn is only supported on Linux or Windows x86/x86_64", ) import zvec # noqa: E402 @@ -524,10 +527,10 @@ def test_optimize_close_reopen_and_query( collection_option, multiple_docs: list[Doc], ): - """Test inserting 100 docs, optimize, close, reopen and query.""" - # Create collection and insert 100 documents + """Test a UTF-8 path through optimize, close, reopen and query.""" + # A non-ASCII path exercises CreateFileW in the Windows DiskAnn reader. temp_dir = tmp_path_factory.mktemp("zvec_diskann_optimize") - collection_path = temp_dir / "test_optimize_collection" + collection_path = temp_dir / "磁盘索引_テスト" coll = zvec.create_and_open( path=str(collection_path), diff --git a/python/tests/test_typing.py b/python/tests/test_typing.py index d566d5efc..9b1bb2c71 100644 --- a/python/tests/test_typing.py +++ b/python/tests/test_typing.py @@ -13,7 +13,10 @@ # limitations under the License. from __future__ import annotations +import sys + import pytest +import zvec from zvec import ( DataType, IndexType, @@ -34,6 +37,8 @@ (DataType.FLOAT, "FLOAT"), (IndexType.HNSW, "HNSW"), (IOBackendType.PREAD, "PREAD"), + (IOBackendType.IO_URING, "IO_URING"), + (IOBackendType.WINDOWS_OVERLAPPED, "WINDOWS_OVERLAPPED"), (MetricType.COSINE, "COSINE"), (QuantizeType.INT8, "INT8"), (StatusCode.OK, "OK"), @@ -49,6 +54,8 @@ def test_enum_names(member, name): (DataType.FLOAT, 8), (IndexType.HNSW, 1), (IOBackendType.PREAD, 0), + (IOBackendType.IO_URING, 2), + (IOBackendType.WINDOWS_OVERLAPPED, 3), (MetricType.COSINE, 3), (QuantizeType.INT8, 2), (StatusCode.OK, 0), @@ -112,11 +119,19 @@ def test_index_type_has_member(member): assert member in IndexType.__members__ -@pytest.mark.parametrize("member", ["PREAD", "LIBAIO"]) +@pytest.mark.parametrize( + "member", ["PREAD", "LIBAIO", "IO_URING", "WINDOWS_OVERLAPPED"] +) def test_io_backend_type_has_member(member): assert member in IOBackendType.__members__ +@pytest.mark.skipif(sys.platform != "win32", reason="Windows-specific backend") +def test_windows_io_backend_type(): + assert zvec.io_backend_type() == IOBackendType.WINDOWS_OVERLAPPED + assert "overlapped" in zvec.io_backend_description().lower() + + @pytest.mark.parametrize("member", ["FP16", "INT8", "INT4", "UNDEFINED"]) def test_quantize_type_has_member(member): assert member in QuantizeType.__members__ diff --git a/python/zvec/__init__.pyi b/python/zvec/__init__.pyi index 09e828dae..4a87d22ff 100644 --- a/python/zvec/__init__.pyi +++ b/python/zvec/__init__.pyi @@ -52,12 +52,11 @@ 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.""" + Linux selects IO_URING, LIBAIO, then PREAD; Windows uses + WINDOWS_OVERLAPPED.""" 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.""" + """Returns a platform-specific description of the current I/O backend.""" def set_default_jieba_dict_dir(dir: str) -> None: """Register the process-wide default jieba dict directory.""" diff --git a/python/zvec/typing/__init__.pyi b/python/zvec/typing/__init__.pyi index fab03f4b9..c3a31d243 100644 --- a/python/zvec/typing/__init__.pyi +++ b/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(). + - IO_URING: io_uring accessed directly through the Linux kernel ABI. + - WINDOWS_OVERLAPPED: Windows overlapped I/O using I/O completion ports. Examples: >>> from zvec.typing import IOBackendType @@ -142,13 +144,21 @@ class IOBackendType: PREAD LIBAIO + + IO_URING + + WINDOWS_OVERLAPPED """ LIBAIO: typing.ClassVar[IOBackendType] # value = + IO_URING: typing.ClassVar[IOBackendType] # value = PREAD: typing.ClassVar[IOBackendType] # value = + WINDOWS_OVERLAPPED: typing.ClassVar[ + IOBackendType + ] # value = __members__: typing.ClassVar[ dict[str, IOBackendType] - ] # value = {'PREAD': , 'LIBAIO': } + ] # value includes PREAD, LIBAIO, IO_URING, and WINDOWS_OVERLAPPED def __eq__(self, other: typing.Any) -> bool: ... def __getstate__(self) -> int: ... diff --git a/src/ailego/io/io_backend.h b/src/ailego/io/io_backend.h new file mode 100644 index 000000000..b7169be52 --- /dev/null +++ b/src/ailego/io/io_backend.h @@ -0,0 +1,18 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +// Compatibility include for internal users of the selector. +#include diff --git a/src/ailego/io/io_backend_def.h b/src/ailego/io/io_backend_def.h index d795c3046..363379917 100644 --- a/src/ailego/io/io_backend_def.h +++ b/src/ailego/io/io_backend_def.h @@ -29,41 +29,61 @@ #pragma once +#include +#include #include #include +#if defined(__linux) || defined(__linux__) +#include +#endif + namespace zvec { namespace ailego { // Returns a human-readable name for the given backend type. inline const char *IOBackendTypeName(IOBackendType type) { switch (type) { + case IOBackendType::kIoUring: + return "io_uring"; case IOBackendType::kLibAio: return "libaio"; case IOBackendType::kPread: return "pread"; + case IOBackendType::kWindowsOverlapped: + return "windows_overlapped"; } return "unknown"; } -// Returns a human-readable description for the given backend type. -// When the backend is kPread, includes installation guidance for libaio. +// Returns a platform-specific human-readable backend description. inline const char *IOBackendDescription(IOBackendType type) { switch (type) { + case IOBackendType::kIoUring: + return "io_uring async I/O backend accessed through the Linux kernel " + "ABI."; case IOBackendType::kLibAio: return "libaio async I/O backend loaded at runtime via dlopen()."; case IOBackendType::kPread: +#if defined(__linux) || defined(__linux__) return "No async I/O backend available. Install libaio (e.g. " "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " "and retry. DiskAnn will fall back to synchronous pread() \u2014 " "performance will be degraded."; +#else + return "Synchronous pread I/O backend."; +#endif + case IOBackendType::kWindowsOverlapped: + return "Windows unbuffered overlapped I/O backend using I/O completion " + "ports."; } return "Unknown I/O backend."; } // Singleton that loads and queries an I/O backend on demand. // -// available() (no arg) tries the best backend with priority (libaio > pread) +// available() (no arg) tries the best backend with priority +// (io_uring > libaio > pread) // and returns the loaded backend type. // available(IOBackendType) tries a specific backend. // Use type() / name() to query the loaded backend without triggering a load. @@ -74,24 +94,49 @@ class IOBackend { return instance; } - // Try to load the best available backend (libaio > pread). + // Select Windows overlapped I/O, or the best available Linux backend + // (io_uring > libaio > pread). // Returns the loaded backend type. // Idempotent — if already loaded, returns immediately. IOBackendType available() { +#if defined(_WIN32) || defined(_WIN64) + type_ = IOBackendType::kWindowsOverlapped; + return type_; +#else if (type_ != IOBackendType::kPread) { return type_; } - return available(IOBackendType::kLibAio); + IOBackendType type = available(IOBackendType::kIoUring); + if (type == IOBackendType::kPread) { + type = available(IOBackendType::kLibAio); + } + return type; +#endif } - // Try to load the requested backend. Returns the loaded backend type - // (may differ from requested if the load failed — falls back to kPread). + // Try to load the requested backend. On Linux, failed requests fall back to + // kPread. Windows always reports its native overlapped-I/O backend. // Idempotent — if the same backend is already loaded, returns immediately. IOBackendType available(IOBackendType requested) { +#if defined(_WIN32) || defined(_WIN64) + (void)requested; + type_ = IOBackendType::kWindowsOverlapped; + return type_; +#else if (type_ == requested && type_ != IOBackendType::kPread) { return type_; } #if defined(__linux) || defined(__linux__) + if (requested == IOBackendType::kIoUring) { + struct io_uring_params params; + std::memset(¶ms, 0, sizeof(params)); + int fd = static_cast(::syscall(__NR_io_uring_setup, 1, ¶ms)); + if (fd >= 0) { + ::close(fd); + type_ = IOBackendType::kIoUring; + return type_; + } + } if (requested == IOBackendType::kLibAio) { if (LibAioLoader::Instance().load() && LibAioLoader::Instance().is_available()) { @@ -102,6 +147,7 @@ class IOBackend { #endif type_ = IOBackendType::kPread; return type_; +#endif } bool is_pread() { @@ -112,6 +158,14 @@ class IOBackend { return available() == IOBackendType::kLibAio; } + bool is_iouring() { + return available() == IOBackendType::kIoUring; + } + + bool is_windows_overlapped() { + return available() == IOBackendType::kWindowsOverlapped; + } + // Returns the loaded backend type. IOBackendType type() const { return type_; @@ -130,7 +184,11 @@ class IOBackend { private: IOBackend() = default; +#if defined(_WIN32) || defined(_WIN64) + IOBackendType type_{IOBackendType::kWindowsOverlapped}; +#else IOBackendType type_{IOBackendType::kPread}; +#endif }; } // namespace ailego diff --git a/src/ailego/io/iouring_def.h b/src/ailego/io/iouring_def.h new file mode 100644 index 000000000..2c0cd7bb4 --- /dev/null +++ b/src/ailego/io/iouring_def.h @@ -0,0 +1,199 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Private header defining the io_uring kernel ABI structures and constants. +// +// This header is the io_uring analogue of libaio_def.h: it declares *only* +// the types, constants, and inline helpers that zvec needs from the io_uring +// kernel interface. By defining these structures ourselves we avoid any +// build-time dependency on or liburing-dev, mirroring the +// project's zero-dependency philosophy established by the libaio dlopen +// approach. +// +// The struct layouts (io_uring_sqe, io_uring_cqe, io_uring_params, +// io_sqring_offsets, io_cqring_offsets) are part of the Linux kernel ABI +// and are copied verbatim from . + +#pragma once + +#include + +#if defined(__linux) || defined(__linux__) + +// --------------------------------------------------------------------------- +// Syscall numbers +// --------------------------------------------------------------------------- +// io_uring was introduced in Linux 5.1 (2019). The three syscalls share the +// same numbers across all supported architectures. We prefer the values +// from when available and fall back to hardcoded numbers. +#include + +#ifndef __NR_io_uring_setup +#define __NR_io_uring_setup 425 +#endif + +#ifndef __NR_io_uring_enter +#define __NR_io_uring_enter 426 +#endif + +#ifndef __NR_io_uring_register +#define __NR_io_uring_register 427 +#endif + +// --------------------------------------------------------------------------- +// Constants (from ) +// --------------------------------------------------------------------------- + +// mmap offsets for the three shared regions. +#define IORING_OFF_SQ_RING 0ULL +#define IORING_OFF_CQ_RING 0x8000000ULL +#define IORING_OFF_SQES 0x10000000ULL + +// io_uring_enter flags. +#define IORING_ENTER_GETEVENTS (1U << 0) + +// io_uring_setup flags (none used by default). +// IORING_SETUP_IOPOLL (1U << 0) +// IORING_SETUP_SQPOLL (1U << 1) +// IORING_SETUP_SQ_AFF (1U << 2) + +// SQE opcode values. +#define IORING_OP_NOP 0 +#define IORING_OP_READV 1 +#define IORING_OP_WRITEV 2 +#define IORING_OP_FSYNC 3 +#define IORING_OP_READ_FIXED 4 +#define IORING_OP_WRITE_FIXED 5 +#define IORING_OP_POLL_ADD 6 +#define IORING_OP_POLL_REMOVE 7 +#define IORING_OP_SYNC_FILE_RANGE 8 +#define IORING_OP_SENDMSG 9 +#define IORING_OP_RECVMSG 10 +#define IORING_OP_TIMEOUT 11 +#define IORING_OP_TIMEOUT_REMOVE 12 +#define IORING_OP_ACCEPT 13 +#define IORING_OP_ASYNC_CANCEL 14 +#define IORING_OP_LINK_TIMEOUT 15 +#define IORING_OP_CONNECT 16 +#define IORING_OP_FALLOCATE 17 +#define IORING_OP_OPENAT 18 +#define IORING_OP_CLOSE 19 +#define IORING_OP_FILES_UPDATE 20 +#define IORING_OP_STATX 21 +#define IORING_OP_READ 22 +#define IORING_OP_WRITE 23 + +// --------------------------------------------------------------------------- +// Struct definitions (copied verbatim from ) +// --------------------------------------------------------------------------- + +// Submission queue entry — 64 bytes. +struct io_uring_sqe { + uint8_t opcode; // type of operation for this sqe + uint8_t flags; // IOSQE_ flags + uint16_t ioprio; // ioprio for the request + int32_t fd; // file descriptor to do IO on + union { + uint64_t off; // offset into file + uint64_t addr2; + }; + union { + uint64_t addr; // buffer or iovecs + uint64_t splice_off_in; + }; + uint32_t len; // buffer size or number of iovecs + union { + uint32_t rw_flags; // read/write flags (union of all flag types) + }; + uint64_t user_data; // data to be passed back at completion time + union { + struct { + uint16_t buf_index; // index into fixed buffers, if used + uint16_t personality; + }; + uint64_t __pad2[3]; + }; +}; + +// Completion queue entry — 16 bytes. +struct io_uring_cqe { + uint64_t user_data; // sqe->user_data + int32_t res; // result code for this event + uint32_t flags; +}; + +// SQ ring offsets — returned by io_uring_setup in io_uring_params. +struct io_sqring_offsets { + uint32_t head; + uint32_t tail; + uint32_t ring_mask; + uint32_t ring_entries; + uint32_t flags; + uint32_t dropped; + uint32_t array; + uint32_t resv1; + uint64_t resv2; +}; + +// CQ ring offsets — returned by io_uring_setup in io_uring_params. +struct io_cqring_offsets { + uint32_t head; + uint32_t tail; + uint32_t ring_mask; + uint32_t ring_entries; + uint32_t overflow; + uint32_t cqes; + uint32_t flags; + uint32_t resv1; + uint64_t resv2; +}; + +// Parameters passed to io_uring_setup(). +struct io_uring_params { + uint32_t sq_entries; + uint32_t cq_entries; + uint32_t flags; + uint32_t sq_thread_cpu; + uint32_t sq_thread_idle; + uint32_t features; + uint32_t wq_fd; + uint32_t resv[3]; + struct io_sqring_offsets sq_off; + struct io_cqring_offsets cq_off; +}; + +// --------------------------------------------------------------------------- +// Inline helper — prepare an SQE for a read operation. +// --------------------------------------------------------------------------- +static inline void io_uring_prep_read(struct io_uring_sqe *sqe, int fd, + void *buf, uint32_t nbytes, + uint64_t offset) { + sqe->opcode = IORING_OP_READ; + sqe->flags = 0; + sqe->ioprio = 0; + sqe->fd = fd; + sqe->off = offset; + sqe->addr = reinterpret_cast(buf); + sqe->len = nbytes; + sqe->rw_flags = 0; + sqe->user_data = 0; + sqe->buf_index = 0; + sqe->personality = 0; +} + +// --------------------------------------------------------------------------- +// End: struct and constant definitions from +// --------------------------------------------------------------------------- + +#endif // __linux__ diff --git a/src/ailego/io/iouring_loader.h b/src/ailego/io/iouring_loader.h new file mode 100644 index 000000000..be997d807 --- /dev/null +++ b/src/ailego/io/iouring_loader.h @@ -0,0 +1,259 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Raw-syscall wrapper for Linux io_uring. +// +// This class implements the io_uring submission/completion queue lifecycle +// using *only* kernel syscalls (io_uring_setup, io_uring_enter) and mmap. +// There is zero dependency on liburing or liburing-dev: no build-time header, +// no runtime .so, and no dlopen. This mirrors the project's existing +// zero-dependency philosophy established by the libaio dlopen approach, but +// goes one step further — io_uring is a pure kernel ABI accessed via syscall. +// +// Runtime detection is automatic: if the kernel does not support io_uring +// (pre-5.1 or io_uring disabled), io_uring_setup() returns -ENOSYS and +// setup() returns false, allowing callers to fall back to libaio or pread. +// +// The ring is **not** thread-safe. Each thread that performs I/O must have +// its own IoUringRing instance (managed through IoBackend / IOContext). + +#pragma once + +#if defined(__linux) || defined(__linux__) + +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +// Forward declaration — AlignedRead is defined in diskann_file_reader.h +// after this header is included. The forward declaration is sufficient +// because IoUringRing::execute takes it by reference. +struct AlignedRead; + +// Maximum number of SQEs we submit in a single io_uring_enter() call. +static constexpr uint32_t kIoUringMaxBatch = 128; + +class IoUringRing { + public: + IoUringRing() = default; + ~IoUringRing() { + teardown(); + } + + IoUringRing(const IoUringRing &) = delete; + IoUringRing &operator=(const IoUringRing &) = delete; + + // Create an io_uring with the given number of entries. + // Returns true on success, false if the kernel does not support io_uring + // or setup failed for any reason. + bool setup(uint32_t entries) { + struct io_uring_params params; + std::memset(¶ms, 0, sizeof(params)); + + // io_uring_setup is a raw syscall — returns fd (>=0) or -1 with errno. + ring_fd_ = static_cast( + syscall(__NR_io_uring_setup, static_cast(entries), ¶ms)); + if (ring_fd_ < 0) { + // ENOSYS = kernel doesn't support io_uring. + // EPERM = io_uring disabled via sysctl. + // EINVAL = invalid parameters. + if (errno != ENOSYS) { + LOG_WARN("io_uring_setup failed; errno=%d, %s", errno, + ::strerror(errno)); + } + return false; + } + + sq_entries_ = params.sq_entries; + cq_entries_ = params.cq_entries; + + // --- mmap the three shared regions --- + + // 1. SQ ring (includes head, tail, mask, entries, flags, dropped, array). + size_t sq_ring_sz = static_cast(params.sq_off.array) + + sq_entries_ * sizeof(uint32_t); + sq_ring_ptr_ = ::mmap(nullptr, sq_ring_sz, PROT_READ | PROT_WRITE, + MAP_SHARED, ring_fd_, IORING_OFF_SQ_RING); + if (sq_ring_ptr_ == MAP_FAILED) { + LOG_ERROR("mmap SQ ring failed: %s", ::strerror(errno)); + sq_ring_ptr_ = nullptr; + teardown(); + return false; + } + + // 2. SQE array. + size_t sqes_sz = sq_entries_ * sizeof(struct io_uring_sqe); + sqes_ptr_ = reinterpret_cast( + ::mmap(nullptr, sqes_sz, PROT_READ | PROT_WRITE, MAP_SHARED, ring_fd_, + IORING_OFF_SQES)); + if (sqes_ptr_ == MAP_FAILED) { + LOG_ERROR("mmap SQEs failed: %s", ::strerror(errno)); + sqes_ptr_ = nullptr; + teardown(); + return false; + } + + // 3. CQ ring (includes head, tail, mask, entries, overflow, cqes[]). + size_t cq_ring_sz = static_cast(params.cq_off.cqes) + + cq_entries_ * sizeof(struct io_uring_cqe); + cq_ring_ptr_ = ::mmap(nullptr, cq_ring_sz, PROT_READ | PROT_WRITE, + MAP_SHARED, ring_fd_, IORING_OFF_CQ_RING); + if (cq_ring_ptr_ == MAP_FAILED) { + LOG_ERROR("mmap CQ ring failed: %s", ::strerror(errno)); + cq_ring_ptr_ = nullptr; + teardown(); + return false; + } + + // --- Set up typed pointers into the mmap'd regions --- + + // SQ ring fields. + sq_head_ = reinterpret_cast(static_cast(sq_ring_ptr_) + + params.sq_off.head); + sq_tail_ = reinterpret_cast(static_cast(sq_ring_ptr_) + + params.sq_off.tail); + sq_ring_mask_ = reinterpret_cast( + static_cast(sq_ring_ptr_) + params.sq_off.ring_mask); + sq_ring_entries_ = reinterpret_cast( + static_cast(sq_ring_ptr_) + params.sq_off.ring_entries); + sq_flags_ = reinterpret_cast(static_cast(sq_ring_ptr_) + + params.sq_off.flags); + sq_dropped_ = reinterpret_cast( + static_cast(sq_ring_ptr_) + params.sq_off.dropped); + sq_array_ = reinterpret_cast(static_cast(sq_ring_ptr_) + + params.sq_off.array); + + // CQ ring fields. + cq_head_ = reinterpret_cast(static_cast(cq_ring_ptr_) + + params.cq_off.head); + cq_tail_ = reinterpret_cast(static_cast(cq_ring_ptr_) + + params.cq_off.tail); + cq_ring_mask_ = reinterpret_cast( + static_cast(cq_ring_ptr_) + params.cq_off.ring_mask); + cq_ring_entries_ = reinterpret_cast( + static_cast(cq_ring_ptr_) + params.cq_off.ring_entries); + cq_overflow_ = reinterpret_cast( + static_cast(cq_ring_ptr_) + params.cq_off.overflow); + cqes_ = reinterpret_cast( + static_cast(cq_ring_ptr_) + params.cq_off.cqes); + + // SQE array. + sqes_ = sqes_ptr_; + + // Initialize the SQ array to identity mapping so that logical index == + // physical SQE index. This is the simplest and most common configuration. + for (uint32_t i = 0; i < sq_entries_; i++) { + sq_array_[i] = i; + } + + return true; + } + + // Tear down the ring: munmap all regions and close the ring fd. + void teardown() { + if (sq_ring_ptr_ && sq_ring_ptr_ != MAP_FAILED) { + // We don't track the exact mmap size; munmap with a large enough size + // is safe because the kernel only unmaps what was actually mapped. + // However, to be correct we use the page-aligned size. + size_t sz = static_cast(sq_entries_) * sizeof(uint32_t) + 4096; + ::munmap(sq_ring_ptr_, sz); + } + if (sqes_ptr_ && sqes_ptr_ != MAP_FAILED) { + size_t sz = + static_cast(sq_entries_) * sizeof(struct io_uring_sqe); + ::munmap(sqes_ptr_, sz); + } + if (cq_ring_ptr_ && cq_ring_ptr_ != MAP_FAILED) { + size_t sz = + static_cast(cq_entries_) * sizeof(struct io_uring_cqe) + 4096; + ::munmap(cq_ring_ptr_, sz); + } + + sq_ring_ptr_ = nullptr; + sqes_ptr_ = nullptr; + cq_ring_ptr_ = nullptr; + sqes_ = nullptr; + cqes_ = nullptr; + sq_head_ = sq_tail_ = sq_ring_mask_ = sq_ring_entries_ = nullptr; + sq_flags_ = sq_dropped_ = sq_array_ = nullptr; + cq_head_ = cq_tail_ = cq_ring_mask_ = cq_ring_entries_ = nullptr; + cq_overflow_ = nullptr; + + if (ring_fd_ >= 0) { + ::close(ring_fd_); + ring_fd_ = -1; + } + + sq_entries_ = 0; + cq_entries_ = 0; + } + + bool is_valid() const { + return ring_fd_ >= 0; + } + + // Execute a batch of aligned read requests via io_uring. + // + // Implemented in diskann_file_reader.cc to avoid a circular dependency: + // AlignedRead is defined in diskann_file_reader.h *after* this header is + // included, so the method body cannot be inline here. + // + // On success returns 0. On failure returns -1; the caller should fall + // back to pread. + int execute(int fd, std::vector &read_reqs); + + private: + int ring_fd_{-1}; + + // mmap'd region base pointers (needed for munmap). + void *sq_ring_ptr_{nullptr}; + struct io_uring_sqe *sqes_ptr_{nullptr}; + void *cq_ring_ptr_{nullptr}; + + // SQ ring field pointers (into sq_ring_ptr_). + unsigned *sq_head_{nullptr}; + unsigned *sq_tail_{nullptr}; + unsigned *sq_ring_mask_{nullptr}; + unsigned *sq_ring_entries_{nullptr}; + unsigned *sq_flags_{nullptr}; + unsigned *sq_dropped_{nullptr}; + unsigned *sq_array_{nullptr}; + + // CQ ring field pointers (into cq_ring_ptr_). + unsigned *cq_head_{nullptr}; + unsigned *cq_tail_{nullptr}; + unsigned *cq_ring_mask_{nullptr}; + unsigned *cq_ring_entries_{nullptr}; + unsigned *cq_overflow_{nullptr}; + struct io_uring_cqe *cqes_{nullptr}; + + // SQE array pointer. + struct io_uring_sqe *sqes_{nullptr}; + + // Ring capacities. + unsigned sq_entries_{0}; + unsigned cq_entries_{0}; +}; + +} // namespace core +} // namespace zvec + +#endif // __linux__ diff --git a/src/binding/c/c_api.cc b/src/binding/c/c_api.cc index 1dfce56f2..d00436bc4 100644 --- a/src/binding/c/c_api.cc +++ b/src/binding/c/c_api.cc @@ -765,6 +765,20 @@ const char *zvec_get_default_jieba_dict_dir(void) { // I/O Backend Introspection // ============================================================================= +static_assert( + static_cast(zvec::ailego::IOBackendType::kPread) == + ZVEC_IO_BACKEND_TYPE_PREAD); +static_assert( + static_cast(zvec::ailego::IOBackendType::kLibAio) == + ZVEC_IO_BACKEND_TYPE_LIBAIO); +static_assert( + static_cast(zvec::ailego::IOBackendType::kIoUring) == + ZVEC_IO_BACKEND_TYPE_IO_URING); +static_assert( + static_cast( + zvec::ailego::IOBackendType::kWindowsOverlapped) == + ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED); + zvec_io_backend_type_t zvec_get_io_backend_type(void) { auto type = zvec::ailego::current_io_backend_type(); return static_cast(static_cast(type)); diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc index d6ad1f48b..49e94219d 100644 --- a/src/binding/python/model/common/python_config.cc +++ b/src/binding/python/model/common/python_config.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "python_config.h" +#include #include #include @@ -228,18 +229,15 @@ 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."); + "Linux selects IO_URING, LIBAIO, then PREAD; Windows uses " + "WINDOWS_OVERLAPPED."); - // Returns a human-readable description of the I/O backend, including - // installation guidance for libaio when only pread is available. + // Returns a platform-specific description of the I/O backend. 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."); + "Returns a platform-specific description of the current I/O backend."); } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 4468a53ff..15f309c90 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -313,7 +313,16 @@ void ZVecPyCollection::bind_dql_methods( "given vector column. One of 'mmap', 'buffer_pool', 'contiguous'. " "Raises KeyError if no HNSW index exists on the column, or " "ValueError if the column's index is not an HNSW index. Intended " - "for introspection and testing only; not part of the stable API."); + "for introspection and testing only; not part of the stable API.") + .def( + "_debug_io_backend_type", + [](const Collection &self) { + const auto result = self.DebugGetIoBackendType(); + return unwrap_expected(result); + }, + "Debug-only: returns the I/O backend type used by DiskAnn. " + "One of 'io_uring', 'libaio', or 'pread'. Intended for introspection " + "and testing only; not part of the stable API."); } } // namespace zvec diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc index d20cac755..183a94402 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -146,6 +146,8 @@ 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(). +- IO_URING: io_uring accessed directly through the Linux kernel ABI. +- WINDOWS_OVERLAPPED: Windows overlapped I/O using I/O completion ports. Examples: >>> from zvec.typing import IOBackendType @@ -153,7 +155,9 @@ Enumeration of supported I/O backend types for DiskAnn async disk reads. IOBackendType.LIBAIO )pbdoc") .value("PREAD", ailego::IOBackendType::kPread) - .value("LIBAIO", ailego::IOBackendType::kLibAio); + .value("LIBAIO", ailego::IOBackendType::kLibAio) + .value("IO_URING", ailego::IOBackendType::kIoUring) + .value("WINDOWS_OVERLAPPED", ailego::IOBackendType::kWindowsOverlapped); } void ZVecPyTyping::bind_status(py::module_ &m) { @@ -245,4 +249,4 @@ Construct a status with the given code and optional message. }); } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/algorithm/diskann/CMakeLists.txt b/src/core/algorithm/diskann/CMakeLists.txt index 5f1b7bf7f..ca00b4523 100644 --- a/src/core/algorithm/diskann/CMakeLists.txt +++ b/src/core/algorithm/diskann/CMakeLists.txt @@ -9,16 +9,18 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c) # wheel, so no separate runtime .so is needed. The shared variant is used by # C++ tools and tests that link against it directly. # -# libaio is loaded at runtime via dlopen()/dlsym() (see libaio_loader.h), so -# we do NOT link against -laio. The only system-level link dependency is -# ${CMAKE_DL_LIBS} for dlopen/dlsym/dlclose. +# On Linux, libaio is loaded at runtime via dlopen()/dlsym() (see +# libaio_loader.h), so we do NOT link against -laio. The only system-level +# link dependency is ${CMAKE_DL_LIBS} for dlopen/dlsym/dlclose. +# On Windows, overlapped I/O is built in, so no extra +# system-level dependency is needed. 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") list(APPEND CORE_KNN_DISKANN_LIBS ${CMAKE_DL_LIBS}) endif() -if(NOT APPLE) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") set(CORE_KNN_DISKANN_LDFLAGS "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") endif() @@ -31,4 +33,26 @@ cc_library( INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm LDFLAGS "${CORE_KNN_DISKANN_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" -) \ No newline at end of file +) + +if(MSVC) + # NOMINMAX must reach the _objects target, since _add_library() compiles + # sources there; the SHARED/STATIC targets only bundle object files and a + # PRIVATE define on them never reaches the compiler. Without this, the + # windows.h min/max macros collide with std::min/std::max and + # numeric_limits::max()/min() (error C4003). + target_compile_definitions(core_knn_diskann_objects PRIVATE NOMINMAX) + target_compile_definitions(core_knn_diskann PRIVATE NOMINMAX) +endif() + +# Ensure the Windows shared-library target links all dependencies explicitly. +foreach(_dep zvec_ailego core_framework core_knn_cluster) + if(TARGET ${_dep}) + target_include_directories(core_knn_diskann PRIVATE + $) + add_dependencies(core_knn_diskann ${_dep}) + if(MSVC) + target_link_libraries(core_knn_diskann ${_dep}) + endif() + endif() +endforeach() diff --git a/src/core/algorithm/diskann/diskann_builder.cc b/src/core/algorithm/diskann/diskann_builder.cc index b3597f1c5..f932e01aa 100644 --- a/src/core/algorithm/diskann/diskann_builder.cc +++ b/src/core/algorithm/diskann/diskann_builder.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include diff --git a/src/core/algorithm/diskann/diskann_builder_entity.cc b/src/core/algorithm/diskann/diskann_builder_entity.cc index 4381293a5..be0e7c75b 100644 --- a/src/core/algorithm/diskann/diskann_builder_entity.cc +++ b/src/core/algorithm/diskann/diskann_builder_entity.cc @@ -14,6 +14,7 @@ #include "diskann_builder_entity.h" #include +#include #include "diskann_algorithm.h" #include "diskann_util.h" diff --git a/src/core/algorithm/diskann/diskann_context.cc b/src/core/algorithm/diskann/diskann_context.cc index f13affb74..b48785653 100644 --- a/src/core/algorithm/diskann/diskann_context.cc +++ b/src/core/algorithm/diskann/diskann_context.cc @@ -27,7 +27,8 @@ DiskAnnContext::DiskAnnContext(const IndexMeta &meta, : dc_(entity.get(), measure, meta.dimension()), entity_{entity} {} int DiskAnnContext::init(ContextType type, uint32_t graph_degree, - uint32_t pq_chunk_num, uint32_t element_size) { + uint32_t pq_chunk_num, uint32_t element_size, + uint32_t disk_element_size) { type_ = type; element_size_ = element_size; pq_chunk_num_ = pq_chunk_num; @@ -60,7 +61,9 @@ int DiskAnnContext::init(ContextType type, uint32_t graph_degree, DiskAnnUtil::alloc_aligned((void **)&pq_coord_buffer_, graph_degree * pq_chunk_num_ * sizeof(uint8_t), 256); - DiskAnnUtil::alloc_aligned((void **)&coord_buffer_, element_size_, 256); + DiskAnnUtil::alloc_aligned( + (void **)&coord_buffer_, + disk_element_size > 0 ? disk_element_size : element_size_, 256); DiskAnnUtil::alloc_aligned( (void **)§or_buffer_, DiskAnnUtil::kMaxSectorReadNum * DiskAnnUtil::kSectorSize, @@ -82,12 +85,12 @@ int DiskAnnContext::init(ContextType type, uint32_t graph_degree, } DiskAnnContext::~DiskAnnContext() { - free(query_); - free(query_rotated_); - free(pq_table_dist_buffer_); - free(pq_coord_buffer_); - free(coord_buffer_); - free(sector_buffer_); + DiskAnnUtil::free_aligned(query_); + DiskAnnUtil::free_aligned(query_rotated_); + DiskAnnUtil::free_aligned(pq_table_dist_buffer_); + DiskAnnUtil::free_aligned(pq_coord_buffer_); + DiskAnnUtil::free_aligned(coord_buffer_); + DiskAnnUtil::free_aligned(sector_buffer_); if (type_ == kSearcherContext) { destroy_io_ctx(io_ctx_); @@ -105,7 +108,7 @@ int DiskAnnContext::update_context(ContextType type, const IndexMeta &meta, const IndexMetric::Pointer &measure, const DiskAnnEntity::Pointer &entity, uint32_t magic_num) { - if (ailego_unlikely(type != type_)) { + if (ailego_unlikely(type != static_cast(type_))) { LOG_ERROR( "DiskAnnContext does not support shared by different type, " "src=%u dst=%u", diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index dd824ff23..32a82cef6 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -60,7 +60,7 @@ class DiskAnnContext : public IndexContext, public: //! Init int init(ContextType type, uint32_t graph_degree, uint32_t pq_chunk_num, - uint32_t element_size); + uint32_t element_size, uint32_t disk_element_size = 0); //! Update context, the context may be shared by different searcher/streamer int update_context(ContextType type, const IndexMeta &meta, @@ -349,7 +349,7 @@ class DiskAnnContext : public IndexContext, uint32_t group_num_{0}; std::map group_topk_heaps_{}; - IOContext io_ctx_{0}; + IOContext io_ctx_{}; SearchStats query_stats_; float *pq_table_dist_buffer_{nullptr}; diff --git a/src/core/algorithm/diskann/diskann_entity.h b/src/core/algorithm/diskann/diskann_entity.h index af302290d..02dc64e06 100644 --- a/src/core/algorithm/diskann/diskann_entity.h +++ b/src/core/algorithm/diskann/diskann_entity.h @@ -24,7 +24,7 @@ using diskann_key_t = uint64_t; using diskann_id_t = uint32_t; constexpr diskann_id_t kInvalidId = static_cast(-1); -constexpr diskann_key_t kInvalidKey = static_cast(-1); +constexpr diskann_key_t kInvalidKey = static_cast(-1); struct VectorInfo { float dist_; diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index cde0755e8..67a111cbd 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -19,8 +19,10 @@ #include #include #include +#include #include #include +#include #include #define MAX_EVENTS 1024 @@ -54,36 +56,80 @@ void log_diskann_io_backend() { #endif } -int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) +int setup_io_ctx(IOContext &ctx) { std::call_once(g_io_backend_log_once, log_diskann_io_backend); - if (ailego::IOBackend::Instance().is_pread()) { + auto &backend = ailego::IOBackend::Instance(); + ailego::IOBackendType selected = backend.available(); + + ctx = new IoBackend(); + + // Priority 1: io_uring (raw kernel syscalls — zero dependency). + if (selected == ailego::IOBackendType::kIoUring && + ctx->ring.setup(MAX_EVENTS)) { + ctx->backend = IoBackend::IO_URING; return 0; } - int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); + if (selected == ailego::IOBackendType::kIoUring) { + selected = backend.available(ailego::IOBackendType::kLibAio); + } - return ret; -#else + // Priority 2: libaio (dlopen — soft dependency). + if (selected == ailego::IOBackendType::kLibAio && + LibAioLoader::Instance().load()) { + int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx->aio_ctx); + if (ret == 0) { + ctx->backend = IoBackend::LIBAIO; + return 0; + } + LOG_WARN("io_setup failed; returned: %d, %s. falling back to pread", ret, + ::strerror(-ret)); + } + + // Priority 3: synchronous pread (always available). + ctx->backend = IoBackend::NONE; return 0; -#endif } int destroy_io_ctx(IOContext &ctx) { -#if (defined(__linux) || defined(__linux__)) - if (ailego::IOBackend::Instance().is_pread() || ctx == nullptr) { + if (ctx == nullptr || ctx == (IOContext)-1) { return 0; } - int ret = LibAioLoader::Instance().io_destroy(ctx); - if (ret == 0) { - ctx = nullptr; + + if (ctx->backend == IoBackend::IO_URING) { + ctx->ring.teardown(); + } else if (ctx->backend == IoBackend::LIBAIO && + LibAioLoader::Instance().is_available()) { + int ret; + do { + ret = LibAioLoader::Instance().io_destroy(ctx->aio_ctx); + } while (ret == -EINTR); + if (ret != 0) { + return ret; + } + ctx->aio_ctx = nullptr; } + // IoUringRing destructor also calls teardown() — idempotent and safe. - return ret; -#else + delete ctx; + ctx = nullptr; + return 0; +} +#elif defined(_WIN32) || defined(_WIN64) +int setup_io_ctx(IOContext &ctx) { + ctx.reqs.resize(MAX_IO_DEPTH); + memset(ctx.reqs.data(), 0, ctx.reqs.size() * sizeof(OVERLAPPED)); return 0; -#endif } +int destroy_io_ctx(IOContext &ctx) { + ctx.close_handles(); + ctx.reqs.clear(); + return 0; +} +#endif + +#if (defined(__linux) || defined(__linux__)) static int execute_io_pread(int fd, std::vector &read_reqs) { for (auto &req : read_reqs) { ssize_t bytes_read = ::pread(fd, req.buf, req.len, req.offset); @@ -102,7 +148,6 @@ static int execute_io_pread(int fd, std::vector &read_reqs) { return 0; } -#if (defined(__linux) || defined(__linux__)) // io_getevents() should only fail permanently for an invalid context or // invalid arguments. If that happens after submission, io_destroy() is the // only safe way to quiesce the context before synchronous I/O touches the same @@ -111,7 +156,7 @@ static bool reset_aio_context(IOContext &ctx) { auto &loader = LibAioLoader::Instance(); int ret; do { - ret = loader.io_destroy(ctx); + ret = loader.io_destroy(ctx->aio_ctx); } while (ret == -EINTR); if (ret != 0) { @@ -120,17 +165,18 @@ static bool reset_aio_context(IOContext &ctx) { return false; } - ctx = nullptr; - IOContext replacement = nullptr; + ctx->aio_ctx = nullptr; + io_context_t replacement = nullptr; ret = loader.io_setup(MAX_EVENTS, &replacement); if (ret != 0) { LOG_ERROR( "io_setup failed while recreating an AIO context; returned: %d, %s. " "this context will use pread", ret, ::strerror(-ret)); + ctx->backend = IoBackend::NONE; return true; } - ctx = replacement; + ctx->aio_ctx = replacement; return true; } @@ -164,8 +210,8 @@ int execute_io_libaio(IOContext &ctx, int fd, // again. while (submitted < n_ops) { size_t remaining = n_ops - submitted; - int ret = LibAioLoader::Instance().io_submit(ctx, (int64_t)remaining, - cbs.data() + submitted); + int ret = LibAioLoader::Instance().io_submit( + ctx->aio_ctx, (int64_t)remaining, cbs.data() + submitted); if (ret > 0 && static_cast(ret) <= remaining) { submitted += static_cast(ret); n_tries = 0; @@ -190,8 +236,8 @@ int execute_io_libaio(IOContext &ctx, int fd, while (completed < submitted) { size_t remaining = submitted - completed; int ret = LibAioLoader::Instance().io_getevents( - ctx, (int64_t)remaining, (int64_t)remaining, evts.data() + completed, - nullptr); + ctx->aio_ctx, (int64_t)remaining, (int64_t)remaining, + evts.data() + completed, nullptr); if (ret > 0 && static_cast(ret) <= remaining) { completed += static_cast(ret); continue; @@ -252,20 +298,136 @@ int execute_io_libaio(IOContext &ctx, int fd, return 0; } -#endif int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) { #if (defined(__linux) || defined(__linux__)) - if (ailego::IOBackend::Instance().is_pread() || ctx == nullptr) { + // Guard against null or sentinel contexts. + if (ctx == nullptr || ctx == (IOContext)-1) { + return execute_io_pread(fd, read_reqs); + } + + // Dispatch based on the active backend. + if (ctx->backend == IoBackend::IO_URING) { + int ret = ctx->ring.execute(fd, read_reqs); + if (ret == 0) { + return 0; + } + // io_uring failed — fall back to pread. + LOG_WARN("io_uring execute failed; falling back to pread"); return execute_io_pread(fd, read_reqs); } - return execute_io_libaio(ctx, fd, read_reqs, n_retries); + + if (ctx->backend == IoBackend::LIBAIO) { + return execute_io_libaio(ctx, fd, read_reqs, n_retries); + } + + // NONE backend — synchronous pread. + return execute_io_pread(fd, read_reqs); #else return execute_io_pread(fd, read_reqs); #endif } +// --------------------------------------------------------------------------- +// IoUringRing::execute — defined here (not in iouring_loader.h) because it +// accesses AlignedRead members, and AlignedRead is defined in +// diskann_file_reader.h after iouring_loader.h is included. +// --------------------------------------------------------------------------- +#if (defined(__linux) || defined(__linux__)) +int IoUringRing::execute(int fd, std::vector &read_reqs) { + if (!is_valid()) { + return -1; + } + if (read_reqs.empty()) { + return 0; + } + + // Process in batches limited by the SQ ring size. + uint32_t batch_size = + std::min(sq_entries_, static_cast(kIoUringMaxBatch)); + uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), batch_size); + + for (uint64_t iter = 0; iter < iters; iter++) { + uint64_t n_ops = + std::min(static_cast(read_reqs.size()) - iter * batch_size, + static_cast(batch_size)); + + // --- Phase 1: Fill SQEs --- + + unsigned tail = __atomic_load_n(sq_tail_, __ATOMIC_ACQUIRE); + unsigned mask = *sq_ring_mask_; + + for (uint64_t j = 0; j < n_ops; j++) { + unsigned idx = (tail + static_cast(j)) & mask; + unsigned sqe_idx = sq_array_[idx]; + struct io_uring_sqe *sqe = &sqes_[sqe_idx]; + + uint64_t req_idx = j + iter * batch_size; + io_uring_prep_read(sqe, fd, read_reqs[req_idx].buf, + static_cast(read_reqs[req_idx].len), + read_reqs[req_idx].offset); + // Store the request index so we can verify the completion. + sqe->user_data = req_idx; + } + + // Memory barrier: ensure SQE contents are visible before tail update. + __sync_synchronize(); + __atomic_store_n(sq_tail_, tail + static_cast(n_ops), + __ATOMIC_RELEASE); + + // --- Phase 2: Submit and wait for completions --- + + int ret = static_cast( + syscall(__NR_io_uring_enter, ring_fd_, static_cast(n_ops), + static_cast(n_ops), IORING_ENTER_GETEVENTS, + static_cast(nullptr), static_cast(0))); + if (ret < 0) { + LOG_WARN( + "io_uring_enter failed; errno=%d, %s, n_ops=%lu. " + "falling back to pread", + errno, ::strerror(errno), (unsigned long)n_ops); + return -1; + } + + // --- Phase 3: Process CQEs --- + + unsigned head = __atomic_load_n(cq_head_, __ATOMIC_ACQUIRE); + unsigned cq_mask = *cq_ring_mask_; + bool all_ok = true; + uint64_t processed = 0; + + for (unsigned i = head; processed < n_ops; i = (i + 1), processed++) { + struct io_uring_cqe *cqe = &cqes_[i & cq_mask]; + uint64_t req_idx = cqe->user_data; + + if (cqe->res < 0) { + LOG_WARN("io_uring read failed: req=%lu, res=%d, offset=%lu", + (unsigned long)req_idx, cqe->res, + (unsigned long)read_reqs[req_idx].offset); + all_ok = false; + } else if (static_cast(cqe->res) != read_reqs[req_idx].len) { + LOG_WARN("io_uring short read: req=%lu, got=%d, expected=%lu", + (unsigned long)req_idx, cqe->res, + (unsigned long)read_reqs[req_idx].len); + all_ok = false; + } + } + + // Advance the CQ head to consume the completions. + __sync_synchronize(); + __atomic_store_n(cq_head_, head + static_cast(n_ops), + __ATOMIC_RELEASE); + + if (!all_ok) { + return -1; + } + } + + return 0; +} +#endif // __linux__ + LinuxAlignedFileReader::LinuxAlignedFileReader(int file_desc) { this->file_desc = file_desc; } @@ -286,7 +448,7 @@ IOContext &LinuxAlignedFileReader::get_ctx() { std::unique_lock lk(ctx_mut); auto it = ctx_map.find(std::this_thread::get_id()); if (it == ctx_map.end()) { - LOG_ERROR("bad thread access; returning -1 as io_context_t"); + LOG_ERROR("bad thread access; returning invalid IOContext"); return this->bad_ctx; } else { return it->second; @@ -299,32 +461,14 @@ void LinuxAlignedFileReader::register_thread() { std::unique_lock lk(ctx_mut); if (ctx_map.find(thread_id) != ctx_map.end()) { LOG_ERROR("multiple calls to register_thread from the same thread"); - return; } IOContext ctx = nullptr; - - std::call_once(g_io_backend_log_once, log_diskann_io_backend); - if (ailego::IOBackend::Instance().is_pread()) { - lk.unlock(); - return; - } - int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); - if (ret != 0) { - if (ret == -EAGAIN) { - LOG_ERROR( - "io_setup failed with EAGAIN: Consider increasing " - "/proc/sys/fs/aio-max-nr"); - } else { - LOG_ERROR("io_setup failed; returned: %d, %s", ret, ::strerror(-ret)); - } - } else { - LOG_INFO("allocating ctx: %lu", (uint64_t)ctx); - + int ret = setup_io_ctx(ctx); + if (ret == 0 && ctx != nullptr) { ctx_map[thread_id] = ctx; } - lk.unlock(); #endif } @@ -345,11 +489,8 @@ void LinuxAlignedFileReader::deregister_thread() { ctx_map.erase(it); } - // io_destroy is a syscall; keep it outside the lock to avoid blocking others - if (ailego::IOBackend::Instance().available() != - ailego::IOBackendType::kPread) { - LibAioLoader::Instance().io_destroy(ctx); - } + // teardown is a syscall; keep it outside the lock to avoid blocking others + destroy_io_ctx(ctx); LOG_INFO("returned ctx from thread"); #endif } @@ -357,13 +498,9 @@ void LinuxAlignedFileReader::deregister_thread() { void LinuxAlignedFileReader::deregister_all_threads() { #if (defined(__linux) || defined(__linux__)) std::unique_lock lk(ctx_mut); - bool aio_available = ailego::IOBackend::Instance().available() != - ailego::IOBackendType::kPread; for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { IOContext ctx = x->second; - if (aio_available) { - LibAioLoader::Instance().io_destroy(ctx); - } + destroy_io_ctx(ctx); } ctx_map.clear(); #endif @@ -421,6 +558,290 @@ int LinuxAlignedFileReader::read(std::vector &read_reqs, return ret; } +#endif // (defined(__linux) || defined(__linux__)) + +#if defined(_WIN32) || defined(_WIN64) + +// ============================================================================ +// Windows implementation using unbuffered overlapped I/O. Each IOContext owns +// a separate file handle and completion port, so concurrent search contexts +// cannot consume one another's completion packets. +// ============================================================================ + +WindowsAlignedFileReader::WindowsAlignedFileReader() {} + +WindowsAlignedFileReader::~WindowsAlignedFileReader() { + deregister_all_threads(); + close(); +} + +IOContext &WindowsAlignedFileReader::get_ctx() { + std::unique_lock lk(ctx_mut); + auto thread_id = std::this_thread::get_id(); + auto it = ctx_map.find(thread_id); + if (it == ctx_map.end()) { + LOG_ERROR("unable to find IOContext for thread_id"); + static IOContext empty_ctx; + return empty_ctx; + } + return it->second; +} + +void WindowsAlignedFileReader::register_thread() { + auto thread_id = std::this_thread::get_id(); + std::unique_lock lk(ctx_mut); + if (ctx_map.find(thread_id) != ctx_map.end()) { + LOG_ERROR("multiple calls to register_thread from the same thread"); + return; + } + + IOContext ctx; + if (setup_io_ctx(ctx) != 0) { + return; + } + ctx_map.emplace(thread_id, std::move(ctx)); +} + +void WindowsAlignedFileReader::deregister_thread() { + auto thread_id = std::this_thread::get_id(); + std::unique_lock lk(ctx_mut); + auto it = ctx_map.find(thread_id); + if (it == ctx_map.end()) { + LOG_ERROR("deregister_thread: thread not registered"); + return; + } + destroy_io_ctx(it->second); + ctx_map.erase(it); +} + +void WindowsAlignedFileReader::deregister_all_threads() { + std::unique_lock lk(ctx_mut); + for (auto &kv : ctx_map) { + destroy_io_ctx(kv.second); + } + ctx_map.clear(); +} + +void WindowsAlignedFileReader::open(const std::string &fname) { + const std::wstring wide_fname = ailego::FileHelper::Utf8ToWide(fname); + if (wide_fname.empty()) { + LOG_ERROR("Failed to convert DiskAnn file path from UTF-8: %s", + fname.c_str()); + return; + } + + close(); + + HANDLE probe_handle = + CreateFileW(wide_fname.c_str(), GENERIC_READ, + FILE_SHARE_READ | FILE_SHARE_DELETE, NULL, OPEN_EXISTING, + FILE_ATTRIBUTE_READONLY | FILE_FLAG_NO_BUFFERING | + FILE_FLAG_OVERLAPPED | FILE_FLAG_RANDOM_ACCESS, + NULL); + if (probe_handle == INVALID_HANDLE_VALUE) { + DWORD error = GetLastError(); + LOG_ERROR("Failed to open file: %s (error=%lu)", fname.c_str(), error); + return; + } + CloseHandle(probe_handle); + + file_path_ = wide_fname; + LOG_INFO("Opened file: %s", fname.c_str()); +} + +void WindowsAlignedFileReader::close() { + file_path_.clear(); +} + +int WindowsAlignedFileReader::prepare_io_ctx(IOContext &ctx) { + if (ctx.file_handle != INVALID_HANDLE_VALUE && ctx.completion_port != NULL && + ctx.file_path == file_path_) { + return 0; + } + + ctx.close_handles(); + ctx.file_handle = + CreateFileW(file_path_.c_str(), GENERIC_READ, + FILE_SHARE_READ | FILE_SHARE_DELETE, NULL, OPEN_EXISTING, + FILE_ATTRIBUTE_READONLY | FILE_FLAG_NO_BUFFERING | + FILE_FLAG_OVERLAPPED | FILE_FLAG_RANDOM_ACCESS, + NULL); + if (ctx.file_handle == INVALID_HANDLE_VALUE) { + LOG_ERROR("Failed to open DiskAnn file for IOContext (error=%lu)", + GetLastError()); + return IndexError_Runtime; + } + + ctx.completion_port = CreateIoCompletionPort(ctx.file_handle, NULL, 0, 1); + if (ctx.completion_port == NULL) { + DWORD error = GetLastError(); + LOG_ERROR("CreateIoCompletionPort failed (error=%lu)", error); + ctx.close_handles(); + return IndexError_Runtime; + } + + ctx.file_path = file_path_; + return 0; +} + +int WindowsAlignedFileReader::read(std::vector &read_reqs, + IOContext &ctx, bool async) { + if (async == true) { + LOG_WARN("Async currently not supported"); + } + + if (file_path_.empty()) { + LOG_ERROR("Attempt to read from invalid file handle"); + return IndexError_Runtime; + } + + if (prepare_io_ctx(ctx) != 0) { + return IndexError_Runtime; + } + + // Ensure the context has enough stable OVERLAPPED slots. ReadFile stores the + // pointer until completion, so this vector must not reallocate mid-batch. + if (ctx.reqs.size() < MAX_IO_DEPTH) { + ctx.reqs.resize(MAX_IO_DEPTH); + } + + if (read_reqs.empty()) { + return 0; + } + + static const DWORD kSectorLen = 4096; + size_t n_reqs = read_reqs.size(); + uint64_t n_batches = DiskAnnUtil::div_round_up(n_reqs, MAX_IO_DEPTH); + + for (uint64_t batch = 0; batch < n_batches; batch++) { + uint64_t batch_start = MAX_IO_DEPTH * batch; + uint64_t batch_size = + std::min((uint64_t)(n_reqs - batch_start), (uint64_t)MAX_IO_DEPTH); + + // Only reset slots used by this batch. The common search path submits a + // small beam, so clearing all MAX_IO_DEPTH slots adds avoidable hot-path + // work. + memset(ctx.reqs.data(), 0, batch_size * sizeof(OVERLAPPED)); + + // Issue ReadFile calls for this batch. + bool batch_error = false; + uint64_t issued_count = 0; + for (uint64_t j = 0; j < batch_size; j++) { + AlignedRead &req = read_reqs[batch_start + j]; + OVERLAPPED &os = ctx.reqs[j]; + + uint64_t offset = req.offset; + uint64_t nbytes = req.len; + char *read_buf = (char *)req.buf; + + // Alignment assertions (must be sector-aligned for unbuffered I/O). + assert((size_t)read_buf % kSectorLen == 0); + assert(offset % kSectorLen == 0); + assert(nbytes % kSectorLen == 0); + + // Fill in the OVERLAPPED struct with the file offset. + os.Offset = (DWORD)(offset & 0xffffffff); + os.OffsetHigh = (DWORD)(offset >> 32); + + if (nbytes > (std::numeric_limits::max)()) { + LOG_ERROR("Read request is too large: %llu bytes", + (unsigned long long)nbytes); + batch_error = true; + break; + } + + BOOL ret = ReadFile(ctx.file_handle, read_buf, (DWORD)nbytes, NULL, &os); + if (ret == FALSE) { + DWORD error = GetLastError(); + if (error != ERROR_IO_PENDING) { + LOG_ERROR("Error queuing IO -- error=%lu", error); + batch_error = true; + break; + } + } + issued_count++; + } + + // Drain every successfully issued request, including after a submission + // error. The caller owns the buffers and may free them as soon as read() + // returns, so no operation may remain outstanding. + uint64_t completed_count = 0; + while (completed_count < issued_count) { + OVERLAPPED_ENTRY entries[MAX_IO_DEPTH]; + ULONG removed = 0; + const ULONG max_entries = static_cast( + std::min(issued_count - completed_count, MAX_IO_DEPTH)); + BOOL dequeued = GetQueuedCompletionStatusEx( + ctx.completion_port, entries, max_entries, &removed, INFINITE, FALSE); + if (dequeued == FALSE) { + LOG_ERROR("GetQueuedCompletionStatusEx failed (error=%lu)", + GetLastError()); + batch_error = true; + + // This should only happen if the completion port itself fails. Cancel + // and poll every request to completion so the caller's buffers remain + // safe. The handle is private to this IOContext. + CancelIoEx(ctx.file_handle, NULL); + for (uint64_t j = 0; j < issued_count; j++) { + while (true) { + DWORD ignored = 0; + if (GetOverlappedResult(ctx.file_handle, &ctx.reqs[j], &ignored, + FALSE)) { + break; + } + if (GetLastError() != ERROR_IO_INCOMPLETE) { + break; + } + SwitchToThread(); + } + } + ctx.close_handles(); + completed_count = issued_count; + break; + } + + completed_count += removed; + for (ULONG i = 0; i < removed; i++) { + OVERLAPPED *os = entries[i].lpOverlapped; + OVERLAPPED *begin = ctx.reqs.data(); + OVERLAPPED *end = begin + batch_size; + if (os < begin || os >= end) { + LOG_ERROR("Completion port returned an unknown OVERLAPPED request"); + batch_error = true; + continue; + } + + const uint64_t slot = static_cast(os - begin); + DWORD bytes_transferred = 0; + BOOL ok = + GetOverlappedResult(ctx.file_handle, os, &bytes_transferred, FALSE); + if (ok == FALSE) { + LOG_ERROR("Overlapped read %lu failed (error=%lu)", + (unsigned long)slot, GetLastError()); + batch_error = true; + continue; + } + + const uint64_t expected = read_reqs[batch_start + slot].len; + if (static_cast(bytes_transferred) != expected) { + LOG_ERROR( + "Overlapped read %lu completed with %lu bytes, expected %lu", + (unsigned long)slot, (unsigned long)bytes_transferred, + (unsigned long)expected); + batch_error = true; + } + } + } + + if (batch_error) { + return IndexError_Runtime; + } + } + + return 0; +} + +#endif // defined(_WIN32) || defined(_WIN64) } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index a1cb7c91a..250924e70 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -18,11 +18,21 @@ #include #if (defined(__linux) || defined(__linux__)) -#include // dlopen-based libaio wrapper +#include // raw-syscall io_uring wrapper +#include // dlopen-based libaio wrapper +#elif defined(_WIN32) || defined(_WIN64) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef _WIN32_WINNT +#define _WIN32_WINNT 0x0600 +#endif +#include #endif -#include #include +#include +#include #include #include #include "diskann_util.h" @@ -31,7 +41,80 @@ namespace zvec { namespace core { #if (defined(__linux) || defined(__linux__)) -typedef io_context_t IOContext; + +// IoBackend holds the per-thread I/O context for whichever async backend +// was successfully initialised at setup time. The priority is: +// 1. io_uring (raw kernel syscalls — zero dependency) +// 2. libaio (dlopen — soft dependency) +// 3. pread (always available — synchronous fallback) +// +// IOContext is a *pointer* to IoBackend, which preserves the existing +// sentinel conventions: nullptr means uninitialised and (IOContext)-1 is +// the invalid-handle sentinel returned by get_ctx() for unregistered +// threads. +struct IoBackend { + enum Backend : uint8_t { + NONE = 0, // synchronous pread + IO_URING = 1, // io_uring via raw syscalls + LIBAIO = 2, // libaio via dlopen + }; + + Backend backend{NONE}; + IoUringRing ring{}; + io_context_t aio_ctx{nullptr}; +}; + +typedef IoBackend *IOContext; + +#elif defined(_WIN32) || defined(_WIN64) +struct IOContext { + std::vector reqs; + HANDLE file_handle{INVALID_HANDLE_VALUE}; + HANDLE completion_port{NULL}; + std::wstring file_path; + + IOContext() = default; + IOContext(const IOContext &) = delete; + IOContext &operator=(const IOContext &) = delete; + + IOContext(IOContext &&other) noexcept + : reqs(std::move(other.reqs)), + file_handle(other.file_handle), + completion_port(other.completion_port), + file_path(std::move(other.file_path)) { + other.file_handle = INVALID_HANDLE_VALUE; + other.completion_port = NULL; + } + + IOContext &operator=(IOContext &&other) noexcept { + if (this != &other) { + close_handles(); + reqs = std::move(other.reqs); + file_handle = other.file_handle; + completion_port = other.completion_port; + file_path = std::move(other.file_path); + other.file_handle = INVALID_HANDLE_VALUE; + other.completion_port = NULL; + } + return *this; + } + + ~IOContext() { + close_handles(); + } + + void close_handles() { + if (file_handle != INVALID_HANDLE_VALUE) { + CloseHandle(file_handle); + file_handle = INVALID_HANDLE_VALUE; + } + if (completion_port != NULL) { + CloseHandle(completion_port); + completion_port = NULL; + } + file_path.clear(); + } +}; #else typedef uint32_t IOContext; #endif @@ -79,6 +162,7 @@ class AlignedFileReader { bool async = false) = 0; }; +#if (defined(__linux) || defined(__linux__)) class LinuxAlignedFileReader : public AlignedFileReader { private: int file_desc; @@ -102,6 +186,37 @@ class LinuxAlignedFileReader : public AlignedFileReader { int read(std::vector &read_reqs, IOContext &ctx, bool async = false); }; +#elif defined(_WIN32) || defined(_WIN64) +class WindowsAlignedFileReader : public AlignedFileReader { + private: + std::wstring file_path_; + + int prepare_io_ctx(IOContext &ctx); + + public: + WindowsAlignedFileReader(); + ~WindowsAlignedFileReader(); + + public: + IOContext &get_ctx(); + + void register_thread(); + void deregister_thread(); + void deregister_all_threads(); + void open(const std::string &fname); + void close(); + + int read(std::vector &read_reqs, IOContext &ctx, + bool async = false); +}; +#endif + +// Platform-agnostic reader typedef +#if (defined(__linux) || defined(__linux__)) +typedef LinuxAlignedFileReader PlatformAlignedFileReader; +#elif defined(_WIN32) || defined(_WIN64) +typedef WindowsAlignedFileReader PlatformAlignedFileReader; +#endif } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_holder.h b/src/core/algorithm/diskann/diskann_holder.h index beb422ba2..8c2333a3e 100644 --- a/src/core/algorithm/diskann/diskann_holder.h +++ b/src/core/algorithm/diskann/diskann_holder.h @@ -25,7 +25,7 @@ struct DiskAnnIndexHolderMeta { uint32_t key_size_; uint32_t sector_size_; uint32_t doc_cnt_; - uint8_t reserve_[]; + uint8_t reserve_[4080]; //!< pad to kMetaSectorSize (4096) }; class DiskAnnIndexHolder : public IndexHolder { @@ -182,7 +182,7 @@ class DiskAnnIndexHolder : public IndexHolder { return IndexError_OpenFile; } - DiskAnnIndexHolderMeta holder_meta; + DiskAnnIndexHolderMeta holder_meta{}; holder_meta.element_size_ = element_size_; holder_meta.key_size_ = sizeof(diskann_key_t); holder_meta.sector_size_ = data_sector_size_; diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 32e7ff67c..3a0d847ab 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -30,13 +30,14 @@ DiskAnnIndexer::DiskAnnIndexer(const IndexMeta &meta) { DiskAnnIndexer::~DiskAnnIndexer() { destroy_io_ctx(init_ctx_); if (centroid_data_) { - free(centroid_data_); + DiskAnnUtil::free_aligned(centroid_data_); } DiskAnnUtil::free_aligned(coord_cache_buf_); } int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity) { entity_ = &entity; + beam_width_ = entity.beam_size(); auto storage = entity.get_storage(); auto vector_segment = entity.get_vector_segment(); @@ -45,7 +46,7 @@ int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity) { index_segment_offset_ = vector_segment->data_offset(); - reader_.reset(new LinuxAlignedFileReader()); + reader_.reset(new PlatformAlignedFileReader()); auto file_path = storage->file_path(); reader_->open(file_path); @@ -80,11 +81,16 @@ int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity) { sector_num_per_node_ = DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); - if (beam_width_ > sector_num_per_node_ * DiskAnnUtil::kMaxSectorReadNum) { - LOG_ERROR("Beamwidth can not be higher than kMaxSectorReadNum"); + if (beam_width_ == 0 || sector_num_per_node_ == 0 || + beam_width_ > DiskAnnUtil::kMaxSectorReadNum / sector_num_per_node_) { + LOG_ERROR( + "Invalid beam size %u: beam_size * sectors_per_node must be in " + "[1, %lu]", + beam_width_, (unsigned long)DiskAnnUtil::kMaxSectorReadNum); return IndexError_InvalidArgument; } + LOG_INFO("DiskAnn search beam size: %u", beam_width_); DiskAnnUtil::alloc_aligned((void **)(¢roid_data_), entrypoints_.size() * aligned_dim_ * sizeof(float), diff --git a/src/core/algorithm/diskann/diskann_indexer.h b/src/core/algorithm/diskann/diskann_indexer.h index c372d288f..8eae84a49 100644 --- a/src/core/algorithm/diskann/diskann_indexer.h +++ b/src/core/algorithm/diskann/diskann_indexer.h @@ -87,11 +87,11 @@ class DiskAnnIndexer { diskann_id_t medoid_; std::vector entrypoints_; - std::shared_ptr reader_{nullptr}; + std::shared_ptr reader_{nullptr}; PQTable::Pointer pq_table_; - IOContext init_ctx_{0}; + IOContext init_ctx_{}; std::vector neighbor_cache_buffer_; void *coord_cache_buf_{nullptr}; diff --git a/src/core/algorithm/diskann/diskann_params.h b/src/core/algorithm/diskann/diskann_params.h index fac0dc60a..35c48df87 100644 --- a/src/core/algorithm/diskann/diskann_params.h +++ b/src/core/algorithm/diskann/diskann_params.h @@ -41,6 +41,8 @@ static const std::string PARAM_DISKANN_SEARCHER_LIST_SIZE( "zvec.diskann.searcher.list_size"); static const std::string PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM( "zvec.diskann.searcher.cache_node_num"); +static const std::string PARAM_DISKANN_SEARCHER_BEAM_SIZE( + "zvec.diskann.searcher.beam_size"); static const std::string PARAM_DISKANN_REDUCER_INDEX_NAME( "zvec.diskann.reducer.index_name"); diff --git a/src/core/algorithm/diskann/diskann_reducer_entity.cc b/src/core/algorithm/diskann/diskann_reducer_entity.cc index 4ccfb6d21..936e76f74 100644 --- a/src/core/algorithm/diskann/diskann_reducer_entity.cc +++ b/src/core/algorithm/diskann/diskann_reducer_entity.cc @@ -105,7 +105,7 @@ int DiskAnnReducerEntity::load_key_segment() { return IndexError_InvalidFormat; } - size_t key_data_len = doc_cnt() * sizeof(key_t); + size_t key_data_len = doc_cnt() * sizeof(diskann_key_t); // load key mapping key_mapping_segment_ = container_->get(kDiskAnnKeyMappingSegmentId); diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index a34c546e5..a3bf4b73c 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "diskann_searcher.h" +#include #include "diskann_context.h" #include "diskann_indexer.h" #include "diskann_params.h" @@ -29,6 +30,7 @@ int DiskAnnSearcher::init(const ailego::Params &search_params) { search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + search_params.get(PARAM_DISKANN_SEARCHER_BEAM_SIZE, &beam_size_); return 0; } @@ -43,7 +45,7 @@ int DiskAnnSearcher::cleanup() { } int DiskAnnSearcher::load(IndexStorage::Pointer storage, - IndexMetric::Pointer measure) { + IndexMetric::Pointer /*measure*/) { LOG_INFO("DiskAnnSearcher::load Begin"); auto start_time = ailego::Monotime::MilliSeconds(); @@ -62,6 +64,7 @@ int DiskAnnSearcher::load(IndexStorage::Pointer storage, diskann_indexer_ = std::make_shared(meta_); + entity_.set_beam_size(beam_size_); int res = diskann_indexer_->init(entity_); if (res != 0) { return res; @@ -79,24 +82,33 @@ int DiskAnnSearcher::load(IndexStorage::Pointer storage, node_list.shrink_to_fit(); } - if (measure) { - measure_ = measure; - } else { - measure_ = IndexFactory::CreateMetric(meta_.metric_name()); - if (!measure_) { - LOG_ERROR("CreateMetric failed, name: %s", meta_.metric_name().c_str()); - return IndexError_NoExist; - } - ret = measure_->init(meta_, meta_.metric_params()); - if (ret != 0) { - LOG_ERROR("IndexMetric init failed, ret=%d", ret); - return ret; - } - if (measure_->query_metric()) { - measure_ = measure_->query_metric(); + search_meta_ = meta_; + if (meta_.metric_name() == "InnerProduct") { + search_meta_.set_metric("SquaredEuclidean", 0, ailego::Params()); + } else if (meta_.metric_name() == "Cosine") { + search_meta_.set_metric("SquaredEuclidean", 0, ailego::Params()); + if (meta_.data_type() == IndexMeta::DataType::DT_FP32) { + search_meta_.set_dimension(meta_.dimension() - 1); + } else { + search_meta_.set_dimension(meta_.dimension() - 2); } } + measure_ = IndexFactory::CreateMetric(search_meta_.metric_name()); + if (!measure_) { + LOG_ERROR("CreateMetric failed, name: %s", + search_meta_.metric_name().c_str()); + return IndexError_NoExist; + } + ret = measure_->init(search_meta_, search_meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("IndexMetric init failed, ret=%d", ret); + return ret; + } + if (measure_->query_metric()) { + measure_ = measure_->query_metric(); + } + stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time); state_ = STATE_LOADED; @@ -122,8 +134,8 @@ int DiskAnnSearcher::update_context(DiskAnnContext *ctx) const { return IndexError_Runtime; } - return ctx->update_context(DiskAnnContext::kSearcherContext, meta_, measure_, - entity, magic_); + return ctx->update_context(DiskAnnContext::kSearcherContext, search_meta_, + measure_, entity, magic_); } int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, @@ -287,15 +299,16 @@ IndexSearcher::Context::Pointer DiskAnnSearcher::create_context() const { return Context::Pointer(); } - DiskAnnContext *ctx = - new (std::nothrow) DiskAnnContext(meta_, measure_, search_ctx_entity); + DiskAnnContext *ctx = new (std::nothrow) + DiskAnnContext(search_meta_, measure_, search_ctx_entity); if (ctx == nullptr) { LOG_ERROR("Failed to allocate DiskAnn Context"); return Context::Pointer(); } if (ailego_unlikely(ctx->init( DiskAnnContext::kSearcherContext, search_ctx_entity->max_degree(), - search_ctx_entity->pq_chunk_num(), meta_.element_size())) != 0) { + search_ctx_entity->pq_chunk_num(), search_meta_.element_size(), + meta_.element_size())) != 0) { LOG_ERROR("Init DiskAnn Context failed"); delete ctx; diff --git a/src/core/algorithm/diskann/diskann_searcher.h b/src/core/algorithm/diskann/diskann_searcher.h index 99584fa35..060fda752 100644 --- a/src/core/algorithm/diskann/diskann_searcher.h +++ b/src/core/algorithm/diskann/diskann_searcher.h @@ -17,8 +17,6 @@ #include "diskann_context.h" #include "diskann_indexer.h" -class LinuxAlignedFileReader; - namespace zvec { namespace core { @@ -41,7 +39,8 @@ class DiskAnnSearcher : public IndexSearcher { int cleanup(void) override; //! Load Index from storage - int load(IndexStorage::Pointer storage, IndexMetric::Pointer metric) override; + int load(IndexStorage::Pointer storage, + IndexMetric::Pointer /*metric*/) override; //! Unload index from storage int unload(void) override; @@ -145,6 +144,7 @@ class DiskAnnSearcher : public IndexSearcher { IndexMetric::Pointer measure_{}; IndexMeta meta_{}; + IndexMeta search_meta_{}; ailego::Params params_{}; uint32_t list_size_{200}; diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.h b/src/core/algorithm/diskann/diskann_searcher_entity.h index 68ab49891..d73bf9723 100644 --- a/src/core/algorithm/diskann/diskann_searcher_entity.h +++ b/src/core/algorithm/diskann/diskann_searcher_entity.h @@ -58,6 +58,14 @@ class DiskAnnSearcherEntity : public DiskAnnEntity { return entrypoints_; } + uint32_t beam_size() const { + return beam_size_; + } + + void set_beam_size(uint32_t beam_size) { + beam_size_ = beam_size; + } + std::pair get_neighbors( diskann_id_t id) const override; diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 82e97dcd6..7374248ed 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "diskann_streamer.h" +#include #include "diskann_context.h" #include "diskann_index_provider.h" #include "diskann_indexer.h" @@ -33,6 +34,7 @@ int DiskAnnStreamer::init(const IndexMeta &meta, search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + search_params.get(PARAM_DISKANN_SEARCHER_BEAM_SIZE, &beam_size_); return 0; } @@ -65,6 +67,7 @@ int DiskAnnStreamer::open(IndexStorage::Pointer storage) { diskann_indexer_ = std::make_shared(meta_); + entity_.set_beam_size(beam_size_); int res = diskann_indexer_->init(entity_); if (res != 0) { return res; @@ -82,12 +85,25 @@ int DiskAnnStreamer::open(IndexStorage::Pointer storage) { node_list.shrink_to_fit(); } - measure_ = IndexFactory::CreateMetric(meta_.metric_name()); + search_meta_ = meta_; + if (meta_.metric_name() == "InnerProduct") { + search_meta_.set_metric("SquaredEuclidean", 0, ailego::Params()); + } else if (meta_.metric_name() == "Cosine") { + search_meta_.set_metric("SquaredEuclidean", 0, ailego::Params()); + if (meta_.data_type() == IndexMeta::DataType::DT_FP32) { + search_meta_.set_dimension(meta_.dimension() - 1); + } else { + search_meta_.set_dimension(meta_.dimension() - 2); + } + } + + measure_ = IndexFactory::CreateMetric(search_meta_.metric_name()); if (!measure_) { - LOG_ERROR("CreateMetric failed, name: %s", meta_.metric_name().c_str()); + LOG_ERROR("CreateMetric failed, name: %s", + search_meta_.metric_name().c_str()); return IndexError_NoExist; } - ret = measure_->init(meta_, meta_.metric_params()); + ret = measure_->init(search_meta_, search_meta_.metric_params()); if (ret != 0) { LOG_ERROR("IndexMetric init failed, ret=%d", ret); return ret; @@ -121,8 +137,8 @@ int DiskAnnStreamer::update_context(DiskAnnContext *ctx) const { return IndexError_Runtime; } - return ctx->update_context(DiskAnnContext::kSearcherContext, meta_, measure_, - entity, magic_); + return ctx->update_context(DiskAnnContext::kSearcherContext, search_meta_, + measure_, entity, magic_); } int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, @@ -322,15 +338,16 @@ IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { return Context::Pointer(); } - DiskAnnContext *ctx = - new (std::nothrow) DiskAnnContext(meta_, measure_, search_ctx_entity); + DiskAnnContext *ctx = new (std::nothrow) + DiskAnnContext(search_meta_, measure_, search_ctx_entity); if (ctx == nullptr) { LOG_ERROR("Failed to allocate DiskAnn Context"); return Context::Pointer(); } if (ailego_unlikely(ctx->init( DiskAnnContext::kSearcherContext, search_ctx_entity->max_degree(), - search_ctx_entity->pq_chunk_num(), meta_.element_size())) != 0) { + search_ctx_entity->pq_chunk_num(), search_meta_.element_size(), + meta_.element_size())) != 0) { LOG_ERROR("Init DiskAnn Context failed"); delete ctx; diff --git a/src/core/algorithm/diskann/diskann_streamer.h b/src/core/algorithm/diskann/diskann_streamer.h index 8d731f399..a3ecf6582 100644 --- a/src/core/algorithm/diskann/diskann_streamer.h +++ b/src/core/algorithm/diskann/diskann_streamer.h @@ -17,8 +17,6 @@ #include "diskann_context.h" #include "diskann_indexer.h" -class LinuxAlignedFileReader; - namespace zvec { namespace core { @@ -155,6 +153,7 @@ class DiskAnnStreamer : public IndexStreamer { IndexMetric::Pointer measure_{}; IndexMeta meta_{}; + IndexMeta search_meta_{}; ailego::Params params_{}; uint32_t list_size_{200}; diff --git a/src/core/algorithm/diskann/diskann_util.h b/src/core/algorithm/diskann/diskann_util.h index a02130bf0..a54fb5d71 100644 --- a/src/core/algorithm/diskann/diskann_util.h +++ b/src/core/algorithm/diskann/diskann_util.h @@ -17,6 +17,10 @@ #include #include "diskann_entity.h" +#if defined(_WIN32) || defined(_WIN64) +#include +#endif + namespace zvec { namespace core { @@ -35,15 +39,22 @@ class DiskAnnUtil { } static inline void alloc_aligned(void **ptr, size_t size, size_t align) { +#if defined(_WIN32) || defined(_WIN64) + *ptr = ::_aligned_malloc(size, align); +#else *ptr = ::aligned_alloc(align, size); +#endif } static inline void free_aligned(void *ptr) { if (ptr == nullptr) { return; } - +#if defined(_WIN32) || defined(_WIN64) + ::_aligned_free(ptr); +#else free(ptr); +#endif } template diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index e377f5342..641d7e389 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -15,6 +15,7 @@ #include #include #include +#include #include #if DISKANN_SUPPORTED #include "algorithm/diskann/diskann_params.h" @@ -27,7 +28,7 @@ namespace zvec::core_interface { int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { (void)param; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -35,24 +36,24 @@ int DiskAnnIndex::Open(const std::string &file_path, StorageOptions storage_options) { (void)file_path; (void)storage_options; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } int DiskAnnIndex::GenerateHolder() { - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } int DiskAnnIndex::Add(const VectorData &vector, uint32_t doc_id) { (void)vector; (void)doc_id; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } int DiskAnnIndex::Train() { - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -60,7 +61,7 @@ int DiskAnnIndex::_dense_fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { (void)doc_id; (void)vector_data_buffer; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -70,7 +71,7 @@ int DiskAnnIndex::_prepare_for_search( (void)query; (void)search_param; (void)context; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -80,7 +81,7 @@ int DiskAnnIndex::Merge(const std::vector &indexes, (void)indexes; (void)filter; (void)options; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -338,4 +339,18 @@ int DiskAnnIndex::Merge(const std::vector &indexes, #endif // DISKANN_SUPPORTED +ailego::IOBackendType DiskAnnIndex::io_backend_type() const { + auto &backend = ailego::IOBackend::Instance(); + ailego::IOBackendType type = backend.available(); +#if defined(__linux) || defined(__linux__) + if (type == ailego::IOBackendType::kPread) { + LOG_WARN( + "Only synchronous pread() is available. Install libaio " + "(e.g. 'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " + "for async I/O support — performance will be degraded without it."); + } +#endif + return type; +} + } // namespace zvec::core_interface diff --git a/src/db/collection.cc b/src/db/collection.cc index 33750f3b1..737034b2e 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -134,6 +134,8 @@ class CollectionImpl : public Collection { Result DebugGetHnswStorageMode( const std::string &column_name) const override; + Result DebugGetIoBackendType() const override; + private: void prepare_schema(); @@ -1833,6 +1835,11 @@ Result CollectionImpl::DebugGetHnswStorageMode( Status::NotFound("No HNSW index found for column '", column_name, "'")); } +Result CollectionImpl::DebugGetIoBackendType() const { + auto type = ailego::IOBackend::Instance().available(); + return std::string(ailego::IOBackendTypeName(type)); +} + Status CollectionImpl::recovery() { if (!FileHelper::DirectoryExists(path_.c_str())) { return Status::InvalidArgument("collection path{", path_, "} not exist."); diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 532696803..528335e5b 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -198,19 +198,15 @@ Status FieldSchema::validate() const { } if (index_params_->type() == IndexType::DISKANN) { - // DiskAnn requires Linux x86_64/i686/i386. The CMake variable - // DISKANN_SUPPORTED (defined in the top-level CMakeLists.txt) is the - // single source of truth for platform eligibility — it is also used by - // index_factory.cc to conditionally compile the DiskAnn index - // registration. Using the same macro here ensures that schema - // validation and index registration agree on supported platforms. - // - // libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() - // and DiskAnnStreamer::init(); if libaio is missing, DiskAnn falls - // back to synchronous pread() with degraded performance. + // DiskAnn is supported wherever DISKANN_SUPPORTED is enabled (Linux + // x86_64 with libaio/io_uring, Windows x86_64 with IOCP). On Linux the + // libaio runtime is loaded eagerly via dlopen inside + // DiskAnnBuilder::init() and DiskAnnStreamer::init(); if libaio is + // missing, DiskAnn falls back to synchronous pread() with degraded + // performance. #if !DISKANN_SUPPORTED return Status::NotSupported( - "DiskAnn is not supported on this platform (Linux x86_64 only)"); + "DiskAnn is not supported on this platform"); #endif } diff --git a/src/include/zvec/ailego/io/io_backend.h b/src/include/zvec/ailego/io/io_backend.h index 008b01514..b35ddead1 100644 --- a/src/include/zvec/ailego/io/io_backend.h +++ b/src/include/zvec/ailego/io/io_backend.h @@ -30,16 +30,18 @@ namespace ailego { // Supported I/O backend types. enum class IOBackendType { - kPread, // Synchronous pread() — no async I/O - kLibAio, // libaio loaded at runtime via dlopen() + kPread = 0, // Synchronous pread() — no async I/O + kLibAio = 1, // libaio loaded at runtime via dlopen() + kIoUring = 2, // io_uring accessed through the Linux kernel ABI + kWindowsOverlapped = 3, // Windows overlapped I/O using completion ports }; // Returns the currently active I/O backend type. -// Triggers backend initialization on first call (libaio > pread). +// On Linux, triggers selection in priority order: io_uring > libaio > pread. +// On Windows, returns kWindowsOverlapped. IOBackendType current_io_backend_type(); // Returns a human-readable description of the currently active I/O backend. -// When only pread is available, includes installation guidance for libaio. std::string current_io_backend_description(); } // namespace ailego diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 2d048689b..4cc9acd0a 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -789,6 +789,10 @@ typedef uint32_t zvec_io_backend_type_t; 0 /**< Synchronous pread() \u2014 no async I/O */ #define ZVEC_IO_BACKEND_TYPE_LIBAIO \ 1 /**< libaio loaded at runtime via dlopen() */ +#define ZVEC_IO_BACKEND_TYPE_IO_URING \ + 2 /**< io_uring accessed through the Linux kernel ABI */ +#define ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED \ + 3 /**< Windows overlapped I/O using I/O completion ports */ /** * @brief Get the current I/O backend type for DiskAnn async disk reads. @@ -796,7 +800,9 @@ typedef uint32_t zvec_io_backend_type_t; * Pure introspection \u2014 no side effects, no install hints. * * @return zvec_io_backend_type_t The loaded backend type - * (ZVEC_IO_BACKEND_TYPE_LIBAIO or ZVEC_IO_BACKEND_TYPE_PREAD). + * ZVEC_IO_BACKEND_TYPE_IO_URING, ZVEC_IO_BACKEND_TYPE_LIBAIO, + * ZVEC_IO_BACKEND_TYPE_PREAD, or + * ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED. */ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); @@ -805,7 +811,8 @@ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); * * @param type The backend type code. * @return Thread-local string valid until the next call on this thread; - * "libaio", "pread", or "unknown". + * "io_uring", "libaio", "pread", "windows_overlapped", or + * "unknown". */ ZVEC_EXPORT const char *ZVEC_CALL zvec_get_io_backend_type_name(zvec_io_backend_type_t type); @@ -813,8 +820,6 @@ zvec_get_io_backend_type_name(zvec_io_backend_type_t type); /** * @brief Get a human-readable description of the current I/O backend. * - * When only pread is available, includes installation guidance for libaio. - * * @return Thread-local string valid until the next call on this thread. */ ZVEC_EXPORT const char *ZVEC_CALL zvec_get_io_backend_description(void); diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 3006718e6..91c1cdc07 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -354,6 +355,10 @@ class ZVEC_CORE_API DiskAnnIndex : public Index { public: DiskAnnIndex() = default; + // Returns the I/O backend type currently loaded for DiskAnn async disk reads. + // If only pread is available, logs a hint to install libaio. + ailego::IOBackendType io_backend_type() const; + protected: virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) override; diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index c4b77575f..2cd30e573 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -118,6 +118,11 @@ class ZVEC_API Collection { //! introspection and testing; not part of the stable public API. virtual Result DebugGetHnswStorageMode( const std::string &column_name) const = 0; + + //! Debug-only: retrieve the I/O backend type used by DiskAnn. Returns + //! "io_uring", "libaio", or "pread". Intended for introspection and + //! testing; not part of the stable public API. + virtual Result DebugGetIoBackendType() const = 0; }; } // namespace zvec diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index 7365dcf1a..67731d00a 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -152,6 +152,31 @@ void test_error_handling_functions(void) { TEST_END(); } +void test_io_backend_functions(void) { + TEST_START(); + + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name( + ZVEC_IO_BACKEND_TYPE_PREAD), + "pread") == 0); + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name( + ZVEC_IO_BACKEND_TYPE_LIBAIO), + "libaio") == 0); + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name( + ZVEC_IO_BACKEND_TYPE_IO_URING), + "io_uring") == 0); + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name( + ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED), + "windows_overlapped") == 0); + +#ifdef _WIN32 + TEST_ASSERT(zvec_get_io_backend_type() == + ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED); + TEST_ASSERT(strstr(zvec_get_io_backend_description(), "overlapped") != NULL); +#endif + + TEST_END(); +} + void test_zvec_config() { TEST_START(); @@ -6376,6 +6401,7 @@ int main(void) { test_version_functions(); test_error_handling_functions(); + test_io_backend_functions(); test_zvec_config(); test_zvec_initialize(); test_zvec_string_functions(); diff --git a/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc index fdc1a396d..049a38082 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc @@ -192,7 +192,10 @@ TEST(DiskAnnLinuxAioTest, AccumulatesPartialSubmissionsAndCompletions) { state.submit_results = {2, 2}; state.completion_results = {1, 2, 1}; FakeAioGuard guard(&state); - IOContext ctx = reinterpret_cast(static_cast(1)); + IoBackend backend; + backend.backend = IoBackend::LIBAIO; + backend.aio_ctx = reinterpret_cast(static_cast(1)); + IOContext ctx = &backend; // An invalid fd makes any accidental pread fallback fail the test. EXPECT_EQ(execute_io_libaio(ctx, -1, requests), 0); @@ -221,7 +224,10 @@ TEST(DiskAnnLinuxAioTest, DrainsPartialSubmissionBeforePreadFallback) { state.submit_results = {2, -EAGAIN}; state.completion_results = {1, 1}; FakeAioGuard guard(&state); - IOContext ctx = reinterpret_cast(static_cast(1)); + IoBackend backend; + backend.backend = IoBackend::LIBAIO; + backend.aio_ctx = reinterpret_cast(static_cast(1)); + IOContext ctx = &backend; EXPECT_EQ(execute_io_libaio(ctx, file.fd(), requests), 0); EXPECT_EQ(state.submit_sizes, (std::vector{4, 2})); @@ -248,7 +254,10 @@ TEST(DiskAnnLinuxAioTest, DrainsAllCompletionsBeforePreadFallback) { state.completion_results = {1, 1, 2}; state.short_completion = 0; FakeAioGuard guard(&state); - IOContext ctx = reinterpret_cast(static_cast(1)); + IoBackend backend; + backend.backend = IoBackend::LIBAIO; + backend.aio_ctx = reinterpret_cast(static_cast(1)); + IOContext ctx = &backend; EXPECT_EQ(execute_io_libaio(ctx, file.fd(), requests), 0); EXPECT_EQ(state.completion_sizes, (std::vector{4, 3, 2})); diff --git a/tests/core/algorithm/diskann/diskann_searcher_test.cc b/tests/core/algorithm/diskann/diskann_searcher_test.cc index ff432886f..017e04375 100644 --- a/tests/core/algorithm/diskann/diskann_searcher_test.cc +++ b/tests/core/algorithm/diskann/diskann_searcher_test.cc @@ -106,6 +106,7 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { Params search_params; search_params.set("zvec.diskann.searcher.list_size", 500); + search_params.set(PARAM_DISKANN_SEARCHER_BEAM_SIZE, 4); ASSERT_EQ(0, searcher->init(search_params)); diff --git a/tests/core/interface/index_group_by_test.cc b/tests/core/interface/index_group_by_test.cc index bb41938fb..d2fd72b14 100644 --- a/tests/core/interface/index_group_by_test.cc +++ b/tests/core/interface/index_group_by_test.cc @@ -44,6 +44,7 @@ struct GroupByCase { bool is_sparse = false; uint32_t dimension = kDimension; bool with_refiner = false; + bool optional = false; // skip when plugin unavailable (e.g. DiskAnn) }; std::shared_ptr> AllPks() { @@ -246,7 +247,16 @@ class GroupByInterfaceTest : public ::testing::Test { ASSERT_EQ(0, source->Train()) << test_case.name; auto index = IndexFactory::CreateAndInitIndex(*test_case.index_param); - ASSERT_NE(nullptr, index) << test_case.name; + if (index == nullptr) { + if (test_case.optional) { + // Optional plugin (e.g. DiskAnn shared module) unavailable at runtime + source->Close(); + zvec::test_util::RemoveTestFiles(index_name + "*"); + zvec::test_util::RemoveTestFiles(source_index_name + "*"); + return; + } + ASSERT_NE(nullptr, index) << test_case.name; + } ASSERT_EQ( 0, index->Open(index_name, {StorageOptions::StorageType::kMMAP, true})) << test_case.name; @@ -522,14 +532,22 @@ TEST_F(GroupByInterfaceTest, UnsupportedIndexTypes) { /*dimension=*/kDimension, /*with_refiner=*/true}, #if DISKANN_SUPPORTED - {"unsupported_diskann_graph", DenseDiskAnnParam(), DiskAnnQuery()}, + {"unsupported_diskann_graph", DenseDiskAnnParam(), DiskAnnQuery(), + /*is_sparse=*/false, /*dimension=*/kDimension, + /*with_refiner=*/false, /*optional=*/true}, {"unsupported_diskann_linear", DenseDiskAnnParam(), - DiskAnnQuery(/*fetch_vector=*/false, /*is_linear=*/true)}, + DiskAnnQuery(/*fetch_vector=*/false, /*is_linear=*/true), + /*is_sparse=*/false, /*dimension=*/kDimension, + /*with_refiner=*/false, /*optional=*/true}, {"unsupported_diskann_bf_pks", DenseDiskAnnParam(), DiskAnnQuery(/*fetch_vector=*/false, /*is_linear=*/false, - /*with_bf_pks=*/true)}, + /*with_bf_pks=*/true), + /*is_sparse=*/false, /*dimension=*/kDimension, + /*with_refiner=*/false, /*optional=*/true}, {"unsupported_diskann_fetch_vector", DenseDiskAnnParam(), - DiskAnnQuery(/*fetch_vector=*/true)}, + DiskAnnQuery(/*fetch_vector=*/true), + /*is_sparse=*/false, /*dimension=*/kDimension, + /*with_refiner=*/false, /*optional=*/true}, #endif }; diff --git a/tools/core/recall_original.cc b/tools/core/recall_original.cc index 997ef5ad7..379e7351d 100644 --- a/tools/core/recall_original.cc +++ b/tools/core/recall_original.cc @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -135,12 +136,8 @@ class Recall { // Prepare file handler vector> output_fs; if (!output_.empty()) { - string cmd = "mkdir -p " + output_; - int ret = system(cmd.c_str()); - if (ret != 0) { - std::cerr << "execute cmd " << cmd << " failed" << std::endl; - return; - } + std::error_code ec; + std::filesystem::create_directories(output_, ec); struct stat sb; if (stat(output_.c_str(), &sb) == 0 && S_ISDIR(sb.st_mode)) { cout << "logs output to : " << output_ << endl; @@ -993,12 +990,8 @@ class SparseRecall { // Prepare file handler vector> output_fs; if (!output_.empty()) { - string cmd = "mkdir -p " + output_; - int ret = system(cmd.c_str()); - if (ret != 0) { - std::cerr << "execute cmd " << cmd << " failed" << std::endl; - return; - } + std::error_code ec; + std::filesystem::create_directories(output_, ec); struct stat sb; if (stat(output_.c_str(), &sb) == 0 && S_ISDIR(sb.st_mode)) { cout << "logs output to : " << output_ << endl;