diff --git a/.github/workflows/03-macos-linux-build.yml b/.github/workflows/03-macos-linux-build.yml index d74257679..6e97916d6 100644 --- a/.github/workflows/03-macos-linux-build.yml +++ b/.github/workflows/03-macos-linux-build.yml @@ -203,7 +203,6 @@ jobs: cd "$GITHUB_WORKSPACE" python -m pytest python/tests/test_collection_diskann.py -v shell: bash - # Verify installing libaio does not affect existing non-DiskAnn tests. - name: Run HNSW Tests (w/ libaio) if: matrix.platform == 'linux-x64' diff --git a/CMakeLists.txt b/CMakeLists.txt index 956ee5599..1565ab79e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -122,14 +122,17 @@ else() endif() message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}") -# DiskAnn support (Linux x86_64 only; libaio loaded at runtime via dlopen) -if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS) +# DiskAnn support: +# - Linux (x86_64, i686, i386, aarch64, arm64) with libaio (loaded via dlopen) +# - macOS (x86_64, ARM64/Apple Silicon) with kqueue +if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64" AND NOT ANDROID AND NOT IOS) + OR (CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND NOT IOS)) set(DISKANN_SUPPORTED ON) add_definitions(-DDISKANN_SUPPORTED=1) else() set(DISKANN_SUPPORTED OFF) add_definitions(-DDISKANN_SUPPORTED=0) - message(STATUS "DiskAnn support disabled - only supported on Linux x86_64") + message(STATUS "DiskAnn support disabled - supported on Linux (x86_64/ARM64 with libaio) and macOS (with kqueue)") endif() message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}") diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index 4e5c703e1..e9531236f 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -15,6 +15,7 @@ endif() get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) +set(ZVEC_SRC_DIR ${ZVEC_ROOT_DIR}/src) set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) include_directories(${ZVEC_INCLUDE_DIR}) diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index f5e6c08e4..538c65df3 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -2,12 +2,10 @@ import logging import platform -DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in ( - "x86_64", - "AMD64", - "i686", - "i386", -) +DISKANN_SUPPORTED = ( + platform.system() == "Linux" + and platform.machine() in ("x86_64", "AMD64", "i686", "i386", "aarch64", "arm64") +) or platform.system() == "Darwin" from typing import Any, Generator from zvec.typing import DataType, StatusCode, MetricType, QuantizeType @@ -27,9 +25,8 @@ def _ensure_diskann_runtime_or_reason() -> str | None: _DISKANN_PRELOAD_DONE = True if not DISKANN_SUPPORTED: - _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64" + _DISKANN_PRELOAD_REASON = "DiskAnn is supported on Linux (x86_64/ARM64 with libaio) and macOS (kqueue)" return _DISKANN_PRELOAD_REASON - _DISKANN_PRELOAD_REASON = None return None diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index 8daf8c770..458c93cad 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -17,8 +17,8 @@ Two platform-level prerequisites are enforced at module import time: -1. DiskAnn is currently built only for Linux x86_64 — other platforms are - skipped wholesale. +1. DiskAnn is built for Linux (x86_64/ARM64 with libaio) and macOS (with kqueue) — + other platforms are skipped wholesale. 2. libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() / DiskAnnStreamer::init(). If libaio is missing, DiskAnn falls back to synchronous pread() — the tests still run but with degraded performance. @@ -39,8 +39,14 @@ # Platform gating (must happen BEFORE we touch zvec). # --------------------------------------------------------------------------- # pytestmark = pytest.mark.skipif( - not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")), - reason="DiskAnn plugin is only supported on Linux x86_64", + not ( + ( + sys.platform == "linux" + and platform.machine() in ("x86_64", "AMD64", "aarch64", "arm64") + ) + or sys.platform == "darwin" + ), + reason="DiskAnn is supported on Linux (x86_64/ARM64 with libaio) and macOS (kqueue)", ) import zvec # noqa: E402 diff --git a/src/ailego/io/io_backend_def.h b/src/ailego/io/io_backend_def.h index d795c3046..10f73decc 100644 --- a/src/ailego/io/io_backend_def.h +++ b/src/ailego/io/io_backend_def.h @@ -29,6 +29,7 @@ #pragma once +#include #include #include @@ -78,30 +79,16 @@ class IOBackend { // Returns the loaded backend type. // Idempotent — if already loaded, returns immediately. IOBackendType available() { - if (type_ != IOBackendType::kPread) { - return type_; - } - return available(IOBackendType::kLibAio); + std::lock_guard lock(mutex_); + return available_locked(IOBackendType::kLibAio); } // Try to load the requested backend. Returns the loaded backend type // (may differ from requested if the load failed — falls back to kPread). // Idempotent — if the same backend is already loaded, returns immediately. IOBackendType available(IOBackendType requested) { - if (type_ == requested && type_ != IOBackendType::kPread) { - return type_; - } -#if defined(__linux) || defined(__linux__) - if (requested == IOBackendType::kLibAio) { - if (LibAioLoader::Instance().load() && - LibAioLoader::Instance().is_available()) { - type_ = IOBackendType::kLibAio; - return type_; - } - } -#endif - type_ = IOBackendType::kPread; - return type_; + std::lock_guard lock(mutex_); + return available_locked(requested); } bool is_pread() { @@ -114,22 +101,41 @@ class IOBackend { // Returns the loaded backend type. IOBackendType type() const { + std::lock_guard lock(mutex_); return type_; } // Human-readable name for the selected backend. const char *name() const { - return IOBackendTypeName(type_); + return IOBackendTypeName(type()); } // Human-readable description for the selected backend. const char *description() const { - return IOBackendDescription(type_); + return IOBackendDescription(type()); } private: IOBackend() = default; + IOBackendType available_locked(IOBackendType requested) { + if (type_ == requested && type_ != IOBackendType::kPread) { + return type_; + } +#if defined(__linux) || defined(__linux__) + if (requested == IOBackendType::kLibAio) { + if (LibAioLoader::Instance().load() && + LibAioLoader::Instance().is_available()) { + type_ = IOBackendType::kLibAio; + return type_; + } + } +#endif + type_ = IOBackendType::kPread; + return type_; + } + + mutable std::mutex mutex_; IOBackendType type_{IOBackendType::kPread}; }; diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 4468a53ff..7645a3325 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 'libaio', 'pread'. Intended for introspection and " + "testing only; not part of the stable API."); } } // namespace zvec diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index f874eba62..23aea7773 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -17,7 +17,7 @@ else() # Empty stub library for unsupported platforms file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/diskann_stub.cc "// Stub implementation for unsupported platforms\n" - "// DiskAnn only supports Linux x86_64\n" + "// DiskAnn supports Linux (x86_64/ARM64 with libaio) and macOS (kqueue)\n" "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" ) diff --git a/src/core/algorithm/diskann/CMakeLists.txt b/src/core/algorithm/diskann/CMakeLists.txt index 5f1b7bf7f..a0723056e 100644 --- a/src/core/algorithm/diskann/CMakeLists.txt +++ b/src/core/algorithm/diskann/CMakeLists.txt @@ -14,7 +14,7 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c) # ${CMAKE_DL_LIBS} for dlopen/dlsym/dlclose. set(CORE_KNN_DISKANN_LIBS core_framework core_knn_cluster) -if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386") +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64") list(APPEND CORE_KNN_DISKANN_LIBS ${CMAKE_DL_LIBS}) endif() @@ -31,4 +31,4 @@ cc_library( INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm LDFLAGS "${CORE_KNN_DISKANN_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" -) \ No newline at end of file +) diff --git a/src/core/algorithm/diskann/diskann_builder.cc b/src/core/algorithm/diskann/diskann_builder.cc index b3597f1c5..da76aee09 100644 --- a/src/core/algorithm/diskann/diskann_builder.cc +++ b/src/core/algorithm/diskann/diskann_builder.cc @@ -13,8 +13,10 @@ // limitations under the License. #include "diskann_builder.h" +#include #include #include +#include #include #include #include @@ -32,6 +34,12 @@ namespace core { int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { LOG_INFO("Begin DiskAnnBuilder::init"); + if (state_ != BUILD_STATE_INIT) { + LOG_ERROR("Cleanup DiskAnnBuilder before reinitializing it"); + return IndexError_NoReady; + } + + cleanup(); log_diskann_io_backend(); params.get(PARAM_DISKANN_BUILDER_MAX_DEGREE, &max_degree_); @@ -39,7 +47,7 @@ int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { params.get(PARAM_DISKANN_BUILDER_THREAD_COUNT, &build_thread_count_); if (build_thread_count_ == 0) { - build_thread_count_ = std::thread::hardware_concurrency(); + build_thread_count_ = std::max(1U, std::thread::hardware_concurrency()); } if (build_thread_count_ > std::thread::hardware_concurrency()) { @@ -62,7 +70,11 @@ int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { if (params.has(PARAM_DISKANN_BUILDER_MEMORY_LIMIT)) { params.get(PARAM_DISKANN_BUILDER_MEMORY_LIMIT, &memory_limit_); - if (memory_limit_ <= 0) { + const double memory_limit_bytes = get_memory_in_bytes(memory_limit_); + if (!std::isfinite(memory_limit_) || memory_limit_ <= 0 || + !std::isfinite(memory_limit_bytes) || + memory_limit_bytes > + static_cast(std::numeric_limits::max())) { LOG_ERROR("Invalid memory limit: %lf", memory_limit_); return IndexError_InvalidArgument; } @@ -88,9 +100,20 @@ int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { build_meta_.set_metric("SquaredEuclidean", 0, ailego::Params()); if (meta.data_type() == IndexMeta::DataType::DT_FP32) { + if (meta.dimension() <= 1) { + LOG_ERROR("Invalid FP32 cosine dimension: %u", meta.dimension()); + return IndexError_InvalidArgument; + } build_meta_.set_dimension(meta.dimension() - 1); - } else { + } else if (meta.data_type() == IndexMeta::DataType::DT_FP16) { + if (meta.dimension() <= 2) { + LOG_ERROR("Invalid FP16 cosine dimension: %u", meta.dimension()); + return IndexError_InvalidArgument; + } build_meta_.set_dimension(meta.dimension() - 2); + } else { + LOG_ERROR("Unsupported cosine data type: %u", meta.data_type()); + return IndexError_Unsupported; } } @@ -127,6 +150,30 @@ int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { } int DiskAnnBuilder::cleanup(void) { + holder_.reset(); + algo_.reset(); + trainer_.reset(); + metric_.reset(); + entity_.clear(); + raw_meta_.clear(); + build_meta_.clear(); + stats_.clear(); + data_file_.clear(); + max_degree_ = kDefaultMaxDegree; + list_size_ = kDefaultListSize; + memory_limit_ = 0.0; + memory_limit_set_ = false; + max_pq_chunk_num_ = kDefaultPqChunkNum; + pq_chunk_num_ = kDefaultPqChunkNum; + build_thread_count_ = 0; + max_train_sample_count_ = PQTable::kMaxTrainSampleCount; + train_sample_ratio_ = PQTable::kTrainSampleRatio; + universal_label_.clear(); + codebook_prefix_.clear(); + index_path_prefix_ = "./diskann"; + errcode_ = 0; + error_ = false; + state_ = BUILD_STATE_INIT; return 0; } @@ -221,30 +268,31 @@ int DiskAnnBuilder::calculate_pq_chunk_num() { return IndexError_InvalidLength; } + uint32_t requested_chunk_num = max_pq_chunk_num_; + if (requested_chunk_num == 0 || requested_chunk_num == kDefaultPqChunkNum) { + requested_chunk_num = + std::max(1U, static_cast(build_meta_.dimension() / 2)); + LOG_INFO( + "No Chunk Num input. Quantizing %u dimension data into %u dimension.", + build_meta_.dimension(), requested_chunk_num); + } + + pq_chunk_num_ = requested_chunk_num; if (memory_limit_set_) { - size_t memory_limit_bytes = get_memory_in_bytes(memory_limit_); - size_t pq_chunk_num = std::floor(memory_limit_bytes / doc_cnt); - if (pq_chunk_num <= 0) { + size_t memory_limit_bytes = + static_cast(get_memory_in_bytes(memory_limit_)); + size_t budget_chunk_num = memory_limit_bytes / doc_cnt; + if (budget_chunk_num == 0) { LOG_ERROR("Insufficient memory limit for vec, memory: %zu, vec num: %zu", memory_limit_bytes, doc_cnt); return IndexError_InvalidArgument; } + pq_chunk_num_ = std::min(pq_chunk_num_, + static_cast(std::min( + budget_chunk_num, build_meta_.dimension()))); } - pq_chunk_num_ = - pq_chunk_num_ < max_pq_chunk_num_ ? pq_chunk_num_ : max_pq_chunk_num_; - - // A chunk num of 0 (public API default) or the internal sentinel means - // "auto": quantize into half the dimensions. Resolve it before the - // upper-bound check so the default never reaches the divide-by-chunk path. - if (pq_chunk_num_ == 0 || pq_chunk_num_ == kDefaultPqChunkNum) { - pq_chunk_num_ = build_meta_.dimension() / 2; - LOG_INFO( - "No Chunk Num input. Quantizing %u dimension data into %u dimension.", - build_meta_.dimension(), pq_chunk_num_); - } - - if (pq_chunk_num_ > build_meta_.dimension()) { + if (pq_chunk_num_ == 0 || pq_chunk_num_ > build_meta_.dimension()) { LOG_ERROR("PQ Chunk Num is more than dimension, chunk num: %u, dim: %u", pq_chunk_num_, build_meta_.dimension()); return IndexError_InvalidArgument; @@ -393,13 +441,20 @@ void DiskAnnBuilder::do_build(uint64_t idx, size_t step_size, } DiskAnnContext::Pointer auto_ptr(ctx); - ctx->init(DiskAnnContext::kBuilderContext, max_degree_, pq_chunk_num_, - build_meta_.element_size()); + int ret = ctx->init(DiskAnnContext::kBuilderContext, max_degree_, + pq_chunk_num_, build_meta_.element_size()); + if (ailego_unlikely(ret != 0)) { + if (!error_.exchange(true)) { + LOG_ERROR("Failed to initialize build context"); + errcode_ = ret; + } + return; + } ctx->set_list_size(list_size_); for (uint64_t id = idx; id < entity_.doc_cnt(); id += step_size) { ctx->reset_query(entity_.get_vector(id)); - int ret = algo_->add_node(id, ctx); + ret = algo_->add_node(id, ctx); if (ailego_unlikely(ret != 0)) { if (!error_.exchange(true)) { LOG_ERROR("DiskAnn graph add node failed"); @@ -432,13 +487,20 @@ void DiskAnnBuilder::do_prune(uint64_t idx, size_t step_size, } DiskAnnContext::Pointer auto_ptr(ctx); - ctx->init(DiskAnnContext::kBuilderContext, max_degree_, pq_chunk_num_, - build_meta_.element_size()); + int ret = ctx->init(DiskAnnContext::kBuilderContext, max_degree_, + pq_chunk_num_, build_meta_.element_size()); + if (ailego_unlikely(ret != 0)) { + if (!error_.exchange(true)) { + LOG_ERROR("Failed to initialize prune context"); + errcode_ = ret; + } + return; + } ctx->set_list_size(list_size_); for (uint64_t id = idx; id < entity_.doc_cnt(); id += step_size) { ctx->reset_query(entity_.get_vector(id)); - int ret = algo_->prune_node(id, ctx); + ret = algo_->prune_node(id, ctx); if (ailego_unlikely(ret != 0)) { if (!error_.exchange(true)) { LOG_ERROR("DiskAnn graph add node failed"); @@ -474,6 +536,10 @@ int DiskAnnBuilder::train(IndexThreads::Pointer threads, LOG_ERROR("Init the builder before DiskAnnBuilder::train"); return IndexError_NoReady; } + if (!holder) { + LOG_ERROR("Invalid holder for DiskAnnBuilder::train"); + return IndexError_InvalidArgument; + } LOG_INFO("Begin DiskAnnBuilder::train"); @@ -529,6 +595,15 @@ int DiskAnnBuilder::do_norm(const void *data_ptr, std::string *norm_data) { int DiskAnnBuilder::build(IndexThreads::Pointer threads, IndexHolder::Pointer holder) { + if (state_ != BUILD_STATE_TRAINED) { + LOG_ERROR("Train the builder before DiskAnnBuilder::build"); + return IndexError_NoReady; + } + if (!holder) { + LOG_ERROR("Invalid holder for DiskAnnBuilder::build"); + return IndexError_InvalidArgument; + } + LOG_INFO("Start DiskAnnBuilder::build"); auto start_time = ailego::Monotime::MilliSeconds(); diff --git a/src/core/algorithm/diskann/diskann_builder.h b/src/core/algorithm/diskann/diskann_builder.h index edc80c483..25e7310a0 100644 --- a/src/core/algorithm/diskann/diskann_builder.h +++ b/src/core/algorithm/diskann/diskann_builder.h @@ -102,7 +102,7 @@ class DiskAnnBuilder : public IndexBuilder { std::string codebook_prefix_{""}; std::string index_path_prefix_{"./diskann"}; - BUILD_STATE state_; + BUILD_STATE state_{BUILD_STATE_INIT}; Stats stats_; int errcode_{0}; diff --git a/src/core/algorithm/diskann/diskann_builder_entity.cc b/src/core/algorithm/diskann/diskann_builder_entity.cc index 4381293a5..d3afd062d 100644 --- a/src/core/algorithm/diskann/diskann_builder_entity.cc +++ b/src/core/algorithm/diskann/diskann_builder_entity.cc @@ -20,9 +20,33 @@ namespace zvec { namespace core { +void DiskAnnBuilderEntity::clear() { + max_degree_ = 0; + list_size_ = 0; + memory_limit_ = 0; + num_threads_ = 0; + max_build_degree_ = 0; + max_observed_degree_ = 0; + neighbor_size_ = 0; + mem_index_file_.clear(); + index_path_prefix_.clear(); + vectors_buffer_.clear(); + keys_buffer_.clear(); + neighbors_buffer_.clear(); + entrypoints_.clear(); + meta_.clear(); + pq_full_pivot_data_.clear(); + pq_centroid_.clear(); + pq_chunk_offsets_.clear(); + block_compressed_data_.clear(); + meta_header_.clear(); + pq_meta_.clear(); +} + int DiskAnnBuilderEntity::init(const IndexMeta &meta, uint32_t max_degree, uint32_t list_size, double memory_limit, uint32_t build_threads) { + clear(); meta_ = meta; max_degree_ = max_degree; diff --git a/src/core/algorithm/diskann/diskann_builder_entity.h b/src/core/algorithm/diskann/diskann_builder_entity.h index d2223ab95..653ea377e 100644 --- a/src/core/algorithm/diskann/diskann_builder_entity.h +++ b/src/core/algorithm/diskann/diskann_builder_entity.h @@ -29,6 +29,8 @@ class DiskAnnBuilderEntity : public DiskAnnEntity { virtual ~DiskAnnBuilderEntity() = default; public: + void clear(); + int add_vector(diskann_key_t key, const void *vec) override; std::pair get_neighbors( @@ -104,4 +106,4 @@ class DiskAnnBuilderEntity : public DiskAnnEntity { }; } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_context.cc b/src/core/algorithm/diskann/diskann_context.cc index f13affb74..220979420 100644 --- a/src/core/algorithm/diskann/diskann_context.cc +++ b/src/core/algorithm/diskann/diskann_context.cc @@ -24,16 +24,27 @@ namespace core { DiskAnnContext::DiskAnnContext(const IndexMeta &meta, const IndexMetric::Pointer &measure, const DiskAnnEntity::Pointer &entity) - : dc_(entity.get(), measure, meta.dimension()), entity_{entity} {} + : IndexContext(measure), + 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) { + if (!entity_ || element_size == 0) { + LOG_ERROR("Invalid DiskAnn context parameters"); + return IndexError_InvalidArgument; + } + type_ = type; element_size_ = element_size; pq_chunk_num_ = pq_chunk_num; DiskAnnUtil::alloc_aligned((void **)&query_, element_size_, 32); DiskAnnUtil::alloc_aligned((void **)&query_rotated_, element_size_, 32); + if (!query_ || !query_rotated_) { + LOG_ERROR("Failed to allocate DiskAnn query buffers"); + return IndexError_NoMemory; + } int ret; switch (type) { @@ -47,6 +58,11 @@ int DiskAnnContext::init(ContextType type, uint32_t graph_degree, break; case kSearcherContext: + if (graph_degree == 0 || pq_chunk_num_ == 0) { + LOG_ERROR("Invalid DiskAnn search context dimensions"); + return IndexError_InvalidArgument; + } + ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), entity_->doc_cnt(), negative_probility_); if (ret != 0) { @@ -54,17 +70,24 @@ int DiskAnnContext::init(ContextType type, uint32_t graph_degree, return ret; } - DiskAnnUtil::alloc_aligned( - (void **)&pq_table_dist_buffer_, - PQTable::kPQCentroidNum * pq_chunk_num_ * sizeof(float), 256); - DiskAnnUtil::alloc_aligned((void **)&pq_coord_buffer_, - graph_degree * pq_chunk_num_ * sizeof(uint8_t), + DiskAnnUtil::alloc_aligned((void **)&pq_table_dist_buffer_, + static_cast(PQTable::kPQCentroidNum) * + pq_chunk_num_ * sizeof(float), 256); + DiskAnnUtil::alloc_aligned( + (void **)&pq_coord_buffer_, + static_cast(graph_degree) * pq_chunk_num_ * sizeof(uint8_t), + 256); DiskAnnUtil::alloc_aligned((void **)&coord_buffer_, element_size_, 256); DiskAnnUtil::alloc_aligned( (void **)§or_buffer_, DiskAnnUtil::kMaxSectorReadNum * DiskAnnUtil::kSectorSize, DiskAnnUtil::kSectorSize); + if (!pq_table_dist_buffer_ || !pq_coord_buffer_ || !coord_buffer_ || + !sector_buffer_) { + LOG_ERROR("Failed to allocate DiskAnn search buffers"); + return IndexError_NoMemory; + } ret = setup_io_ctx(io_ctx_); if (ret != 0) { @@ -132,11 +155,12 @@ int DiskAnnContext::update_context(ContextType type, const IndexMeta &meta, } entity_ = entity; - dc_.update(measure, meta.dimension()); + update_index_metric(measure); + dc_.update(entity_.get(), measure, meta.dimension()); magic_ = magic_num; return 0; } } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index ce8c0ca0c..6f9eb6af6 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -214,13 +214,33 @@ class DiskAnnContext : public IndexContext, for (auto &it : results_) { it.clear(); } + for (auto &it : group_results_) { + it.clear(); + } best_list_nodes_.clear(); expanded_nodes_.clear(); visit_filter_.clear(); + group_topk_heaps_.clear(); has_error_ = false; } + //! Preserve query-time options when replacing an incompatible pooled + //! context. Search scratch space and previous results are intentionally not + //! copied. + void copy_query_options_from(const DiskAnnContext &rhs) { + copy_common_query_options_from(rhs); + topk_ = rhs.topk_; + topk_heap_.clear(); + topk_heap_.limit(topk_); + group_num_ = rhs.group_num_; + group_topk_ = rhs.group_topk_; + group_topk_heaps_.clear(); + list_size_ = rhs.list_size_; + fetch_vector_ = rhs.fetch_vector_; + debug_mode_ = rhs.debug_mode_; + } + SearchStats &query_stats() { return query_stats_; } @@ -269,7 +289,7 @@ class DiskAnnContext : public IndexContext, } //! Get if group by search - inline bool group_by_search() { + inline bool group_by_search() const { return group_num_ > 0; } @@ -348,7 +368,13 @@ class DiskAnnContext : public IndexContext, uint32_t group_num_{0}; std::map group_topk_heaps_{}; - IOContext io_ctx_{0}; + IOContext io_ctx_{ +#if defined(__APPLE__) || defined(__MACH__) + -1 +#else + 0 +#endif + }; SearchStats query_stats_; float *pq_table_dist_buffer_{nullptr}; diff --git a/src/core/algorithm/diskann/diskann_dist_calculator.h b/src/core/algorithm/diskann/diskann_dist_calculator.h index e64b6f35e..43bcd54b4 100644 --- a/src/core/algorithm/diskann/diskann_dist_calculator.h +++ b/src/core/algorithm/diskann/diskann_dist_calculator.h @@ -35,7 +35,9 @@ class DistCalculator { dim_(dim), compare_cnt_(0) {} - void update(const IndexMetric::Pointer &measure, uint32_t dim) { + void update(const DiskAnnEntity *entity, const IndexMetric::Pointer &measure, + uint32_t dim) { + entity_ = entity; distance_ = measure->distance(); dim_ = dim; } diff --git a/src/core/algorithm/diskann/diskann_entity.h b/src/core/algorithm/diskann/diskann_entity.h index af302290d..bd5ce5142 100644 --- a/src/core/algorithm/diskann/diskann_entity.h +++ b/src/core/algorithm/diskann/diskann_entity.h @@ -72,17 +72,6 @@ struct DiskAnnMetaHeader { clear(); } - DiskAnnMetaHeader(const DiskAnnMetaHeader &header) { - memcpy(this, &header, sizeof(header)); - } - - DiskAnnMetaHeader &operator=(const DiskAnnMetaHeader &header) { - if (this != &header) { - memcpy(this, &header, sizeof(header)); - } - return *this; - } - void reset() { doc_cnt = 0U; } @@ -104,15 +93,6 @@ struct DiskAnnPqMeta { clear(); } - DiskAnnPqMeta(const DiskAnnPqMeta &meta) { - memcpy(this, &meta, sizeof(meta)); - } - - DiskAnnPqMeta &operator=(const DiskAnnPqMeta &meta) { - memcpy(this, &meta, sizeof(meta)); - return *this; - } - void clear() { memset(this, 0, sizeof(DiskAnnPqMeta)); } @@ -236,8 +216,8 @@ class DiskAnnEntity { constexpr static uint32_t kRevision = 0U; protected: - DiskAnnMetaHeader meta_header_; - DiskAnnPqMeta pq_meta_; + DiskAnnMetaHeader meta_header_{}; + DiskAnnPqMeta pq_meta_{}; }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index cde0755e8..e65fd0445 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -22,6 +22,10 @@ #include #include #include +#if defined(__APPLE__) || defined(__MACH__) +#include +#include +#endif #define MAX_EVENTS 1024 @@ -61,8 +65,18 @@ int setup_io_ctx(IOContext &ctx) { return 0; } int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); - return ret; +#elif defined(__APPLE__) || defined(__MACH__) + // Create a kqueue for this context. On macOS the kqueue is used to + // monitor file descriptor readiness for async-style I/O. + int kq = ::kqueue(); + if (kq == -1) { + LOG_ERROR("kqueue() failed in setup_io_ctx; errno=%d, %s", errno, + ::strerror(errno)); + return IndexError_Runtime; + } + ctx = kq; + return 0; #else return 0; #endif @@ -79,28 +93,162 @@ int destroy_io_ctx(IOContext &ctx) { } return ret; +#elif defined(__APPLE__) || defined(__MACH__) + if (ctx >= 0) { + ::close(ctx); + ctx = -1; + } + return 0; #else return 0; #endif } -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); - if (bytes_read < 0) { - LOG_ERROR("pread failed; errno=%d, %s, offset=%lu, len=%lu", errno, - ::strerror(errno), (unsigned long)req.offset, - (unsigned long)req.len); - return IndexError_Runtime; +static int execute_one_pread(int fd, const AlignedRead &req) { + auto *buf = static_cast(req.buf); + uint64_t offset = req.offset; + uint64_t remaining = req.len; + + while (remaining > 0) { + ssize_t bytes_read = + ::pread(fd, buf, static_cast(remaining), offset); + if (bytes_read > 0) { + buf += bytes_read; + offset += static_cast(bytes_read); + remaining -= static_cast(bytes_read); + continue; } - if ((size_t)bytes_read != req.len) { - LOG_ERROR("pread short read; got=%zd, expected=%lu", bytes_read, - (unsigned long)req.len); + if (bytes_read == 0) { + LOG_ERROR("pread returned EOF; offset=%llu, remaining=%llu", + (unsigned long long)offset, (unsigned long long)remaining); return IndexError_Runtime; } + if (errno == EINTR) { + continue; + } + + LOG_ERROR("pread failed; errno=%d, %s, offset=%llu, len=%llu", errno, + ::strerror(errno), (unsigned long long)offset, + (unsigned long long)remaining); + return IndexError_Runtime; + } + + return 0; +} + +static int execute_io_pread(int fd, std::vector &read_reqs) { + for (const auto &req : read_reqs) { + int ret = execute_one_pread(fd, req); + if (ret != 0) { + return ret; + } + } + return 0; +} + +#if defined(__APPLE__) || defined(__MACH__) +// Execute batch I/O on macOS using kqueue to monitor file descriptor +// readiness and pread for actual data transfer. +// +// On macOS, regular file descriptors are almost always "readable", so +// kqueue's primary value here is providing the same async I/O interface +// as Linux's libaio. For each read request we: +// 1. Attempt a non-blocking pread (via O_NONBLOCK on the fd). +// 2. If EAGAIN, wait on kqueue for EVFILT_READ readiness, then retry. +// 3. Fall back to blocking pread if kqueue encounters an error. +// +// The kqueue fd is passed in as the IOContext. If no valid kqueue is +// available, we fall back to plain blocking pread. +static int execute_io_kqueue(int kq, int fd, + std::vector &read_reqs) { + // If no kqueue available, fall back to blocking pread. + if (kq < 0) { + return execute_io_pread(fd, read_reqs); + } + + // Register the file descriptor with the kqueue for read events. + // EV_CLEAR gives edge-triggered semantics so we only get notified + // when new data becomes available. + struct kevent ke; + EV_SET(&ke, fd, EVFILT_READ, EV_ADD | EV_CLEAR, 0, 0, nullptr); + + for (auto &req : read_reqs) { + while (true) { + ssize_t bytes_read = ::pread(fd, req.buf, req.len, req.offset); + + if (bytes_read > 0) { + // Successfully read data; verify full read. + if ((size_t)bytes_read != req.len) { + // Partial read — retry for the remaining bytes. + // Update offset and buffer to read the rest. + char *buf_ptr = static_cast(req.buf) + bytes_read; + uint64_t new_offset = req.offset + bytes_read; + size_t remaining = req.len - bytes_read; + while (remaining > 0) { + ssize_t n = ::pread(fd, buf_ptr, remaining, new_offset); + if (n < 0) { + if (errno == EINTR) continue; + LOG_ERROR("pread retry failed; errno=%d, %s", errno, + ::strerror(errno)); + return IndexError_Runtime; + } + if (n == 0) break; + buf_ptr += n; + new_offset += n; + remaining -= n; + } + if (remaining > 0) { + LOG_ERROR("pread short read after retry; remaining=%zu", remaining); + return IndexError_Runtime; + } + } + break; // Success, move to next request. + } + + if (bytes_read == 0) { + // EOF — should not happen for a valid index file. + LOG_ERROR("pread returned 0 (EOF); offset=%lu, len=%lu", + (unsigned long)req.offset, (unsigned long)req.len); + return IndexError_Runtime; + } + + // bytes_read == -1, error + if (errno == EINTR) { + continue; // Retry on signal. + } + if (errno == EAGAIN || errno == EWOULDBLOCK) { + // Data not ready — wait on kqueue for readability. + struct kevent events[1]; + struct timespec ts; + ts.tv_sec = 5; // 5 second timeout as a safety net. + ts.tv_nsec = 0; + int n_ev = ::kevent(kq, &ke, 1, events, 1, &ts); + if (n_ev < 0) { + if (errno == EINTR) continue; + // kqueue error — fall back to blocking pread for this request. + LOG_WARN("kevent failed; errno=%d, %s, falling back to pread", errno, + ::strerror(errno)); + return execute_io_pread(fd, read_reqs); + } + if (n_ev == 0) { + // Timeout — fall back to blocking pread. + LOG_WARN("kqueue timeout, falling back to pread"); + return execute_io_pread(fd, read_reqs); + } + // Event triggered — retry pread. + continue; + } + + // Other error — fall back to blocking pread. + LOG_ERROR("pread failed; errno=%d, %s, falling back to pread", errno, + ::strerror(errno)); + return execute_io_pread(fd, read_reqs); + } } + return 0; } +#endif // __APPLE__ #if (defined(__linux) || defined(__linux__)) // io_getevents() should only fail permanently for an invalid context or @@ -261,7 +409,13 @@ int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, return execute_io_pread(fd, read_reqs); } return execute_io_libaio(ctx, fd, read_reqs, n_retries); +#elif defined(__APPLE__) || defined(__MACH__) + // On macOS, use kqueue-based I/O. The IOContext (ctx) is a kqueue fd. + (void)n_retries; + return execute_io_kqueue(ctx, fd, read_reqs); #else + (void)ctx; + (void)n_retries; return execute_io_pread(fd, read_reqs); #endif } @@ -299,7 +453,6 @@ 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; } @@ -321,10 +474,27 @@ void LinuxAlignedFileReader::register_thread() { } } else { LOG_INFO("allocating ctx: %lu", (uint64_t)ctx); - ctx_map[thread_id] = ctx; } + lk.unlock(); +#elif defined(__APPLE__) || defined(__MACH__) + 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 = -1; + int kq = ::kqueue(); + if (kq == -1) { + LOG_ERROR("kqueue() failed in register_thread; errno=%d, %s", errno, + ::strerror(errno)); + } else { + LOG_INFO("allocating kqueue ctx: %d", kq); + ctx = kq; + ctx_map[thread_id] = ctx; + } lk.unlock(); #endif } @@ -346,19 +516,36 @@ void LinuxAlignedFileReader::deregister_thread() { } // io_destroy is a syscall; keep it outside the lock to avoid blocking others - if (ailego::IOBackend::Instance().available() != - ailego::IOBackendType::kPread) { + if (!ailego::IOBackend::Instance().is_pread()) { LibAioLoader::Instance().io_destroy(ctx); } LOG_INFO("returned ctx from thread"); +#elif defined(__APPLE__) || defined(__MACH__) + auto thread_id = std::this_thread::get_id(); + IOContext ctx; + + { + std::lock_guard lk(ctx_mut); + auto it = ctx_map.find(thread_id); + if (it == ctx_map.end()) { + LOG_ERROR("deregister_thread: thread not registered"); + return; + } + ctx = it->second; + ctx_map.erase(it); + } + + if (ctx >= 0) { + ::close(ctx); + } + LOG_INFO("returned kqueue ctx from thread"); #endif } 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; + bool aio_available = !ailego::IOBackend::Instance().is_pread(); for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { IOContext ctx = x->second; if (aio_available) { @@ -366,6 +553,15 @@ void LinuxAlignedFileReader::deregister_all_threads() { } } ctx_map.clear(); +#elif defined(__APPLE__) || defined(__MACH__) + std::unique_lock lk(ctx_mut); + for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { + IOContext ctx = x->second; + if (ctx >= 0) { + ::close(ctx); + } + } + ctx_map.clear(); #endif } @@ -395,6 +591,32 @@ void LinuxAlignedFileReader::open(const std::string &fname) { ::strerror(errno)); } +#if defined(__APPLE__) || defined(__MACH__) + // macOS has no O_DIRECT. F_NOCACHE is its closest per-file equivalent: it + // asks the kernel to minimize caching for I/O through this descriptor. This + // is advisory rather than a guarantee that every read reaches the device. + // Disable read-ahead as well because DiskAnn performs random reads. + // + // Do not mmap the entire index and call msync(MS_INVALIDATE) here. That does + // not provide a reliable global cache eviction guarantee and makes open time + // and virtual-address usage scale with the size of the index. + if (this->file_desc != -1) { + if (::fcntl(this->file_desc, F_NOCACHE, 1) == -1) { + LOG_WARN( + "fcntl(F_NOCACHE) failed for %s (errno=%d: %s); reads will use " + "the page cache", + fname.c_str(), errno, ::strerror(errno)); + } else { + LOG_INFO("DiskAnn macOS: F_NOCACHE enabled for %s", fname.c_str()); + } + + if (::fcntl(this->file_desc, F_RDAHEAD, 0) == -1) { + LOG_WARN("fcntl(F_RDAHEAD, 0) failed for %s (errno=%d: %s)", + fname.c_str(), errno, ::strerror(errno)); + } + } +#endif + LOG_INFO("Opened file : %s", fname.c_str()); } @@ -421,6 +643,5 @@ int LinuxAlignedFileReader::read(std::vector &read_reqs, return ret; } - } // 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..9c8bd531c 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -21,8 +21,17 @@ #include // dlopen-based libaio wrapper #endif +#if defined(__APPLE__) || defined(__MACH__) +#include +#include +#include +#endif + #include #include +#include +#include +#include #include #include #include "diskann_util.h" @@ -30,8 +39,13 @@ namespace zvec { namespace core { +// On Linux, IOContext is the libaio io_context_t. +// On macOS, IOContext is an int holding a kqueue file descriptor. +// On other platforms, IOContext is a uint32_t placeholder. #if (defined(__linux) || defined(__linux__)) typedef io_context_t IOContext; +#elif defined(__APPLE__) || defined(__MACH__) +typedef int IOContext; #else typedef uint32_t IOContext; #endif @@ -52,9 +66,12 @@ struct AlignedRead { AlignedRead(uint64_t offset, uint64_t len, void *buf) : offset(offset), len(len), buf(buf) { +#if defined(__linux__) || defined(__linux) + // O_DIRECT requires 512-byte alignment on Linux. ailego_assert(static_cast(offset) % 512 == 0); ailego_assert(static_cast(len) % 512 == 0); ailego_assert(reinterpret_cast(buf) % 512 == 0); +#endif } }; @@ -79,6 +96,10 @@ class AlignedFileReader { bool async = false) = 0; }; +// Reader implementation used on all supported platforms. +// On Linux (x86_64 and ARM64) it uses libaio for asynchronous batch I/O. +// On macOS (including ARM/Apple Silicon) it uses kqueue to monitor file +// descriptor readiness and pread for actual data transfer. class LinuxAlignedFileReader : public AlignedFileReader { private: int file_desc; diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 32e7ff67c..bf5d3c105 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -86,11 +86,18 @@ int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity) { return IndexError_InvalidArgument; } - DiskAnnUtil::alloc_aligned((void **)(¢roid_data_), - entrypoints_.size() * aligned_dim_ * sizeof(float), - 32); + centroid_stride_ = DiskAnnUtil::round_up(meta_.element_size(), 32); + DiskAnnUtil::alloc_aligned(¢roid_data_, + entrypoints_.size() * centroid_stride_, 32); + if (centroid_data_ == nullptr) { + LOG_ERROR("Failed to allocate entrypoint vector buffer"); + return IndexError_NoMemory; + } - use_medroids_data_as_centroids(); + ret = use_medroids_data_as_centroids(); + if (ret != 0) { + return ret; + } return 0; } @@ -98,26 +105,20 @@ int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity) { int DiskAnnIndexer::use_medroids_data_as_centroids() { LOG_INFO("Loading centroid data from medoid vector data"); - std::vector nodes_to_read; - std::vector medoid_bufs; - std::vector> neighbor_bufs; - - std::vector centroid_buffer; - - size_t dim = meta_.dimension(); - centroid_buffer.resize(dim); - - nodes_to_read.push_back(medoid_); - medoid_bufs.push_back(&(centroid_buffer[0])); - neighbor_bufs.emplace_back(0, nullptr); - - auto read_status = read_nodes(nodes_to_read, medoid_bufs, neighbor_bufs); + std::vector entrypoint_buffers(entrypoints_.size()); + std::vector> neighbor_buffers( + entrypoints_.size(), std::make_pair(0U, nullptr)); + auto *entrypoint_data = static_cast(centroid_data_); + for (size_t i = 0; i < entrypoints_.size(); ++i) { + entrypoint_buffers[i] = entrypoint_data + i * centroid_stride_; + } - if (read_status[0] == true) { - for (uint32_t i = 0; i < dim; i++) centroid_data_[i] = centroid_buffer[i]; - } else { - LOG_ERROR("Failed to read medoid"); - return IndexError_Runtime; + auto read_status = + read_nodes(entrypoints_, entrypoint_buffers, neighbor_buffers); + if (std::find(read_status.begin(), read_status.end(), false) != + read_status.end()) { + LOG_ERROR("Failed to read one or more entrypoint vectors"); + return IndexError_ReadData; } return 0; @@ -137,6 +138,9 @@ std::vector DiskAnnIndexer::read_nodes( std::vector> &neighbor_buffers) { std::vector read_reqs; std::vector retval(node_ids.size(), true); + if (node_ids.empty()) { + return retval; + } uint8_t *buf = nullptr; auto sector_num = @@ -146,6 +150,11 @@ std::vector DiskAnnIndexer::read_nodes( DiskAnnUtil::alloc_aligned( (void **)&buf, node_ids.size() * sector_num * DiskAnnUtil::kSectorSize, DiskAnnUtil::kSectorSize); + if (buf == nullptr) { + LOG_ERROR("read_nodes: failed to allocate aligned read buffer"); + std::fill(retval.begin(), retval.end(), false); + return retval; + } for (size_t i = 0; i < node_ids.size(); ++i) { auto node_id = node_ids[i]; @@ -202,6 +211,9 @@ int DiskAnnIndexer::load_cache_list( LOG_INFO("Loading the cache list into memory"); size_t num_cached_nodes = node_list.size(); + if (num_cached_nodes == 0) { + return 0; + } neighbor_cache_buffer_.resize(num_cached_nodes * (max_degree_ + 1), 0); @@ -209,6 +221,11 @@ int DiskAnnIndexer::load_cache_list( DiskAnnUtil::alloc_aligned((void **)&coord_cache_buf_, coord_cache_buf_len * meta_.unit_size(), 8 * meta_.unit_size()); + if (coord_cache_buf_ == nullptr) { + LOG_ERROR("Failed to allocate coordinate cache buffer"); + neighbor_cache_buffer_.clear(); + return IndexError_NoMemory; + } memset(coord_cache_buf_, 0, coord_cache_buf_len * meta_.unit_size()); @@ -385,6 +402,21 @@ int DiskAnnIndexer::linear_search(DiskAnnContext *ctx) { auto &topk_heap = ctx->topk_heap(); topk_heap.clear(); + auto &group_topk_heaps = ctx->group_topk_heaps(); + group_topk_heaps.clear(); + auto emplace_candidate = [&](diskann_id_t id, VectorInfo info) { + if (ctx->group_by_search() && ctx->group_by().is_valid()) { + topk_heap.emplace(id, info); + std::string group_id = ctx->group_by()(get_key(id)); + auto &group_topk_heap = group_topk_heaps[group_id]; + if (group_topk_heap.empty()) { + group_topk_heap.limit(ctx->group_topk()); + } + group_topk_heap.emplace(id, std::move(info)); + } else { + topk_heap.emplace(id, std::move(info)); + } + }; IOContext &io_ctx = ctx->io_ctx(); void *aligned_query_raw = ctx->query(); @@ -479,7 +511,7 @@ int DiskAnnIndexer::linear_search(DiskAnnContext *ctx) { float cur_expanded_dist = dc.dist(aligned_query_raw, node_fp_coords_copy); - topk_heap.emplace( + emplace_candidate( std::get<0>(cached_neighbor), VectorInfo(cur_expanded_dist, make_vector_copy(node_fp_coords_copy))); } @@ -494,7 +526,7 @@ int DiskAnnIndexer::linear_search(DiskAnnContext *ctx) { float cur_expanded_dist = dc.dist(aligned_query_raw, data_buf); - topk_heap.emplace( + emplace_candidate( frontier_neighbor.first, VectorInfo(cur_expanded_dist, make_vector_copy(data_buf))); @@ -520,6 +552,21 @@ int DiskAnnIndexer::keys_search(const std::vector &keys, auto &topk_heap = ctx->topk_heap(); topk_heap.clear(); + auto &group_topk_heaps = ctx->group_topk_heaps(); + group_topk_heaps.clear(); + auto emplace_candidate = [&](diskann_id_t id, VectorInfo info) { + if (ctx->group_by_search() && ctx->group_by().is_valid()) { + topk_heap.emplace(id, info); + std::string group_id = ctx->group_by()(get_key(id)); + auto &group_topk_heap = group_topk_heaps[group_id]; + if (group_topk_heap.empty()) { + group_topk_heap.limit(ctx->group_topk()); + } + group_topk_heap.emplace(id, std::move(info)); + } else { + topk_heap.emplace(id, std::move(info)); + } + }; IOContext &io_ctx = ctx->io_ctx(); void *aligned_query_raw = ctx->query(); @@ -557,6 +604,13 @@ int DiskAnnIndexer::keys_search(const std::vector &keys, while (frontier.size() < beam_width_) { if (!ctx->filter().is_valid() || !ctx->filter()(keys[idx])) { diskann_id_t id = get_id(keys[idx]); + if (id == kInvalidId) { + ++idx; + if (idx >= keys.size()) { + break; + } + continue; + } auto iter = neighbor_cache_.find(id); if (iter != neighbor_cache_.end()) { @@ -616,7 +670,7 @@ int DiskAnnIndexer::keys_search(const std::vector &keys, float cur_expanded_dist = dc.dist(aligned_query_raw, node_fp_coords_copy); - topk_heap.emplace( + emplace_candidate( std::get<0>(cached_neighbor), VectorInfo(cur_expanded_dist, make_vector_copy(node_fp_coords_copy))); } @@ -631,7 +685,7 @@ int DiskAnnIndexer::keys_search(const std::vector &keys, float cur_expanded_dist = dc.dist(aligned_query_raw, data_buf); - topk_heap.emplace( + emplace_candidate( frontier_neighbor.first, VectorInfo(cur_expanded_dist, make_vector_copy(data_buf))); @@ -709,8 +763,13 @@ int DiskAnnIndexer::get_vector(diskann_id_t id, IndexContext::Pointer &context, io_timer.reset(); - reader_->read(frontier_read_reqs, io_ctx); + int read_ret = reader_->read(frontier_read_reqs, io_ctx); stats.io_us += io_timer.micro_seconds(); + if (read_ret != 0) { + LOG_ERROR("get_vector: reader_->read failed, ret=%d", read_ret); + ctx->set_error(true); + return IndexError_Runtime; + } uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( node_per_sector_, max_node_size_, frontier_neighbor.second, id); @@ -770,11 +829,12 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { candidates.reserve(ctx->list_size()); - diskann_id_t best_medoid = 0; + diskann_id_t best_medoid = entrypoints_.front(); float best_dist = (std::numeric_limits::max)(); for (uint64_t cur_m = 0; cur_m < entrypoints_.size(); cur_m++) { - float cur_expanded_dist = - dc.dist(ctx->query(), centroid_data_ + aligned_dim_ * cur_m); + const void *entrypoint = + static_cast(centroid_data_) + centroid_stride_ * cur_m; + float cur_expanded_dist = dc.dist(ctx->query(), entrypoint); if (cur_expanded_dist < best_dist) { best_medoid = entrypoints_[cur_m]; @@ -956,35 +1016,39 @@ int DiskAnnIndexer::cached_beam_search_in_mem(DiskAnnContext * /*ctx*/) { return IndexError_NotImplemented; } -int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { +void DiskAnnIndexer::populate_group_topk_heaps(DiskAnnContext *ctx) { + auto &group_topk_heaps = ctx->group_topk_heaps(); + group_topk_heaps.clear(); if (!ctx->group_by().is_valid()) { - return 0; + return; } - std::function group_by = [&](diskann_id_t id) { - return ctx->group_by()(get_key(id)); - }; - - // devide into groups auto &topk_heap = ctx->topk_heap(); - auto &visit_filter = ctx->visit_filter(); - - std::map &group_topk_heaps = ctx->group_topk_heaps(); - for (uint32_t i = 0; i < topk_heap.size(); ++i) { diskann_id_t id = topk_heap[i].first; - auto info = topk_heap[i].second; - - std::string group_id = group_by(id); + const auto &info = topk_heap[i].second; + std::string group_id = ctx->group_by()(get_key(id)); auto &group_topk_heap = group_topk_heaps[group_id]; if (group_topk_heap.empty()) { group_topk_heap.limit(ctx->group_topk()); } + group_topk_heap.emplace(id, info); + } +} - topk_heap.emplace(id, info); +int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { + if (!ctx->group_by().is_valid()) { + ctx->group_topk_heaps().clear(); + return 0; } + // Divide the initial candidates into groups. + auto &topk_heap = ctx->topk_heap(); + auto &visit_filter = ctx->visit_filter(); + populate_group_topk_heaps(ctx); + auto &group_topk_heaps = ctx->group_topk_heaps(); + // stage 2, expand to reach group num as possible if (group_topk_heaps.size() < ctx->group_num()) { NeighborPriorityQueue candidates; @@ -1086,8 +1150,14 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { io_timer.reset(); - reader_->read(frontier_read_reqs, io_ctx); // synchronous IO linux + int read_ret = reader_->read(frontier_read_reqs, io_ctx); stats.io_us += io_timer.micro_seconds(); + if (read_ret != 0) { + LOG_ERROR("cached_beam_search_by_group: reader_->read failed, ret=%d", + read_ret); + ctx->set_error(true); + return IndexError_Runtime; + } } for (auto &cached_neighbor : cached_neighbors) { @@ -1099,7 +1169,8 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { if (!ctx->filter().is_valid() || !ctx->filter()(get_key(std::get<0>(cached_neighbor)))) { - std::string group_id = group_by(std::get<0>(cached_neighbor)); + std::string group_id = + ctx->group_by()(get_key(std::get<0>(cached_neighbor))); auto &group_topk_heap = group_topk_heaps[group_id]; if (group_topk_heap.empty()) { @@ -1153,7 +1224,8 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { if (!ctx->filter().is_valid() || !ctx->filter()(get_key(frontier_neighbor.first))) { - std::string group_id = group_by(frontier_neighbor.first); + std::string group_id = + ctx->group_by()(get_key(frontier_neighbor.first)); auto &group_topk_heap = group_topk_heaps[group_id]; if (group_topk_heap.empty()) { diff --git a/src/core/algorithm/diskann/diskann_indexer.h b/src/core/algorithm/diskann/diskann_indexer.h index c372d288f..51d39dc21 100644 --- a/src/core/algorithm/diskann/diskann_indexer.h +++ b/src/core/algorithm/diskann/diskann_indexer.h @@ -66,6 +66,7 @@ class DiskAnnIndexer { protected: int use_medroids_data_as_centroids(); + void populate_group_topk_heaps(DiskAnnContext *ctx); private: DiskAnnSearcherEntity *entity_; @@ -82,7 +83,8 @@ class DiskAnnIndexer { uint64_t index_segment_offset_{0}; uint64_t sector_num_per_node_{0}; - float *centroid_data_{nullptr}; + void *centroid_data_{nullptr}; + size_t centroid_stride_{0}; diskann_id_t medoid_; std::vector entrypoints_; @@ -91,7 +93,13 @@ class DiskAnnIndexer { PQTable::Pointer pq_table_; - IOContext init_ctx_{0}; + IOContext init_ctx_{ +#if defined(__APPLE__) || defined(__MACH__) + -1 +#else + 0 +#endif + }; std::vector neighbor_cache_buffer_; void *coord_cache_buf_{nullptr}; diff --git a/src/core/algorithm/diskann/diskann_pq_trainer.cc b/src/core/algorithm/diskann/diskann_pq_trainer.cc index 73ca01656..c84744cb8 100644 --- a/src/core/algorithm/diskann/diskann_pq_trainer.cc +++ b/src/core/algorithm/diskann/diskann_pq_trainer.cc @@ -153,7 +153,7 @@ int DiskAnnPqTrainer::convert_pivot_data( cluster * dim + chunk_offsets[chunk]; const T *feature_ptr = reinterpret_cast(centroids[idx].feature()); - for (size_t d = 0; d <= chunk_dims[chunk]; ++d) { + for (size_t d = 0; d < chunk_dims[chunk]; ++d) { pivot_data_ptr[d] = feature_ptr[d]; } } diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index a34c546e5..630f5d294 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -20,15 +20,39 @@ namespace zvec { namespace core { +namespace { + +bool query_meta_matches(const IndexMeta &meta, const IndexQueryMeta &qmeta) { + return qmeta.data_type() == meta.data_type() && + qmeta.dimension() == meta.dimension() && + qmeta.element_size() == meta.element_size(); +} + +bool group_options_valid(const DiskAnnContext *ctx) { + return !ctx->group_by_search() || + (ctx->group_topk() > 0 && ctx->group_by().is_valid()); +} + +} // namespace + DiskAnnSearcher::DiskAnnSearcher() {} DiskAnnSearcher::~DiskAnnSearcher() {} int DiskAnnSearcher::init(const ailego::Params &search_params) { + if (state_ == STATE_LOADED) { + LOG_ERROR("Unload DiskAnnSearcher before reinitializing it"); + return IndexError_NoReady; + } + + params_ = search_params; + list_size_ = 200; + cache_nodes_num_ = 0; log_diskann_io_backend(); - search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); - search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + params_.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); + params_.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + state_ = STATE_INITED; return 0; } @@ -37,6 +61,12 @@ void DiskAnnSearcher::print_debug_info() {} int DiskAnnSearcher::cleanup() { LOG_INFO("Begin DiskAnnSearcher:cleanup"); + unload(); + params_.clear(); + list_size_ = 200; + cache_nodes_num_ = 0; + state_ = STATE_INIT; + LOG_INFO("End DiskAnnSearcher:cleanup"); return 0; @@ -46,6 +76,21 @@ int DiskAnnSearcher::load(IndexStorage::Pointer storage, IndexMetric::Pointer measure) { LOG_INFO("DiskAnnSearcher::load Begin"); + if (!storage) { + LOG_ERROR("Invalid storage"); + return IndexError_InvalidArgument; + } + if (state_ != STATE_INITED) { + LOG_ERROR("Initialize and unload DiskAnnSearcher before loading an index"); + return IndexError_NoReady; + } + + diskann_indexer_.reset(); + entity_.clear(); + measure_.reset(); + meta_.clear(); + stats_.clear(); + auto start_time = ailego::Monotime::MilliSeconds(); int ret = IndexHelper::DeserializeFromStorage(storage.get(), &meta_); @@ -73,7 +118,10 @@ int DiskAnnSearcher::load(IndexStorage::Pointer storage, diskann_indexer_->cache_bfs_levels(cache_nodes_num_, node_list); - diskann_indexer_->load_cache_list(node_list); + ret = diskann_indexer_->load_cache_list(node_list); + if (ret != 0) { + return ret; + } node_list.clear(); node_list.shrink_to_fit(); @@ -98,6 +146,7 @@ int DiskAnnSearcher::load(IndexStorage::Pointer storage, } stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time); + stats_.set_loaded_count(entity_.doc_cnt()); state_ = STATE_LOADED; magic_ = IndexContext::GenerateMagic(); @@ -110,7 +159,13 @@ int DiskAnnSearcher::load(IndexStorage::Pointer storage, int DiskAnnSearcher::unload() { LOG_INFO("DiskAnnSearcher unload index"); - state_ = STATE_INITED; + const State next_state = state_ == STATE_INIT ? STATE_INIT : STATE_INITED; + diskann_indexer_.reset(); + entity_.clear(); + measure_.reset(); + meta_.clear(); + stats_.clear(); + state_ = next_state; return 0; } @@ -126,14 +181,43 @@ int DiskAnnSearcher::update_context(DiskAnnContext *ctx) const { entity, magic_); } +int DiskAnnSearcher::ensure_compatible_context(ContextPointer &context, + DiskAnnContext *&ctx) const { + if (ctx->magic() == magic_) { + return 0; + } + + auto replacement = create_context(); + if (!replacement) { + LOG_ERROR("Failed to recreate context for current searcher"); + return IndexError_Runtime; + } + auto *replacement_ctx = dynamic_cast(replacement.get()); + if (!replacement_ctx) { + LOG_ERROR("Failed to cast recreated DiskAnn context"); + return IndexError_Cast; + } + replacement_ctx->copy_query_options_from(*ctx); + context = std::move(replacement); + ctx = replacement_ctx; + return 0; +} + int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { - // do search + if (ailego_unlikely(state_ != STATE_LOADED)) { + LOG_ERROR("Load DiskAnnSearcher before searching"); + return IndexError_NoReady; + } if (ailego_unlikely(!query || !context)) { LOG_ERROR("The context is not created by this searcher"); return IndexError_Mismatch; } + if (ailego_unlikely(!query_meta_matches(meta_, qmeta))) { + LOG_ERROR("Query meta does not match DiskAnn index meta"); + return IndexError_Mismatch; + } DiskAnnContext *ctx = dynamic_cast(context.get()); ailego_do_if_false(ctx) { @@ -141,18 +225,13 @@ int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, return IndexError_Cast; } - // Context is pooled per index type. When switching between DiskAnn indexes - // with different element sizes (e.g., fp16 vs fp32), the cached context has - // undersized buffers. Recreate it to ensure correct buffer allocations. - if (ctx->magic() != magic_) { - uint32_t saved_topk = ctx->topk(); - context = create_context(); - if (!context) { - LOG_ERROR("Failed to recreate context for current streamer"); - return IndexError_Runtime; - } - ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + int ret = ensure_compatible_context(context, ctx); + if (ret != 0) { + return ret; + } + if (ailego_unlikely(!group_options_valid(ctx))) { + LOG_ERROR("Group search requires a callback and a positive group topk"); + return IndexError_InvalidArgument; } ctx->clear(); @@ -161,7 +240,10 @@ int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, for (uint32_t i = 0; i < count; i++) { ctx->reset_query(query); - diskann_indexer_->knn_search(ctx); + ret = diskann_indexer_->knn_search(ctx); + if (ailego_unlikely(ret != 0)) { + return ret; + } if (ailego_unlikely(ctx->error())) { return IndexError_Runtime; @@ -178,10 +260,18 @@ int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, int DiskAnnSearcher::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (ailego_unlikely(state_ != STATE_LOADED)) { + LOG_ERROR("Load DiskAnnSearcher before searching"); + return IndexError_NoReady; + } if (ailego_unlikely(!query || !context)) { LOG_ERROR("The context is not created by this searcher"); return IndexError_Mismatch; } + if (ailego_unlikely(!query_meta_matches(meta_, qmeta))) { + LOG_ERROR("Query meta does not match DiskAnn index meta"); + return IndexError_Mismatch; + } DiskAnnContext *ctx = dynamic_cast(context.get()); ailego_do_if_false(ctx) { @@ -189,17 +279,13 @@ int DiskAnnSearcher::search_bf_impl(const void *query, return IndexError_Cast; } - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer, recreate it - //! to ensure buffers are correctly sized for this index's parameters. - uint32_t saved_topk = ctx->topk(); - context = create_context(); - if (!context) { - LOG_ERROR("Failed to recreate context for current streamer"); - return IndexError_Runtime; - } - ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + int ret = ensure_compatible_context(context, ctx); + if (ret != 0) { + return ret; + } + if (ailego_unlikely(!group_options_valid(ctx))) { + LOG_ERROR("Group search requires a callback and a positive group topk"); + return IndexError_InvalidArgument; } ctx->clear(); @@ -208,7 +294,10 @@ int DiskAnnSearcher::search_bf_impl(const void *query, for (size_t i = 0; i < count; ++i) { ctx->reset_query(query); - diskann_indexer_->linear_search(ctx); + ret = diskann_indexer_->linear_search(ctx); + if (ailego_unlikely(ret != 0)) { + return ret; + } ctx->topk_to_result(i); @@ -226,10 +315,18 @@ int DiskAnnSearcher::search_bf_by_p_keys_impl( const void *query, const std::vector> &p_keys, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (ailego_unlikely(state_ != STATE_LOADED)) { + LOG_ERROR("Load DiskAnnSearcher before searching"); + return IndexError_NoReady; + } if (ailego_unlikely(!query || !context)) { LOG_ERROR("The context is not created by this searcher"); return IndexError_Mismatch; } + if (ailego_unlikely(!query_meta_matches(meta_, qmeta))) { + LOG_ERROR("Query meta does not match DiskAnn index meta"); + return IndexError_Mismatch; + } DiskAnnContext *ctx = dynamic_cast(context.get()); ailego_do_if_false(ctx) { @@ -242,17 +339,13 @@ int DiskAnnSearcher::search_bf_by_p_keys_impl( return IndexError_InvalidArgument; } - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer, recreate it - //! to ensure buffers are correctly sized for this index's parameters. - uint32_t saved_topk = ctx->topk(); - context = create_context(); - if (!context) { - LOG_ERROR("Failed to recreate context for current streamer"); - return IndexError_Runtime; - } - ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + int ret = ensure_compatible_context(context, ctx); + if (ret != 0) { + return ret; + } + if (ailego_unlikely(!group_options_valid(ctx))) { + LOG_ERROR("Group search requires a callback and a positive group topk"); + return IndexError_InvalidArgument; } ctx->clear(); @@ -261,7 +354,10 @@ int DiskAnnSearcher::search_bf_by_p_keys_impl( for (size_t i = 0; i < count; ++i) { ctx->reset_query(query); - diskann_indexer_->keys_search(p_keys[i], ctx); + ret = diskann_indexer_->keys_search(p_keys[i], ctx); + if (ailego_unlikely(ret != 0)) { + return ret; + } ctx->topk_to_result(i); @@ -277,10 +373,38 @@ int DiskAnnSearcher::search_bf_by_p_keys_impl( int DiskAnnSearcher::get_vector(uint64_t key, Context::Pointer &context, std::string &vector) const { - return diskann_indexer_->get_vector(key, context, vector); + vector.clear(); + if (state_ != STATE_LOADED) { + LOG_ERROR("Load DiskAnnSearcher before fetching vectors"); + return IndexError_NoReady; + } + if (!context) { + LOG_ERROR("Invalid context for get_vector"); + return IndexError_Mismatch; + } + auto *ctx = dynamic_cast(context.get()); + if (!ctx) { + LOG_ERROR("Cast context to DiskAnnContext failed"); + return IndexError_Cast; + } + int ret = ensure_compatible_context(context, ctx); + if (ret != 0) { + return ret; + } + + diskann_id_t id = diskann_indexer_->get_id(key); + if (id == kInvalidId) { + LOG_ERROR("Vector key does not exist: %lu", (unsigned long)key); + return IndexError_NoExist; + } + return diskann_indexer_->get_vector(id, context, vector); } IndexSearcher::Context::Pointer DiskAnnSearcher::create_context() const { + if (state_ != STATE_LOADED) { + LOG_ERROR("Load DiskAnnSearcher before creating a context"); + return Context::Pointer(); + } const DiskAnnEntity::Pointer search_ctx_entity = entity_.clone(); if (!search_ctx_entity) { LOG_ERROR("Failed to create search context entity"); diff --git a/src/core/algorithm/diskann/diskann_searcher.h b/src/core/algorithm/diskann/diskann_searcher.h index 99584fa35..ee9900d56 100644 --- a/src/core/algorithm/diskann/diskann_searcher.h +++ b/src/core/algorithm/diskann/diskann_searcher.h @@ -139,6 +139,8 @@ class DiskAnnSearcher : public IndexSearcher { //! To share ctx across streamer/searcher, we need to update the context for //! current streamer/searcher int update_context(DiskAnnContext *ctx) const; + int ensure_compatible_context(ContextPointer &context, + DiskAnnContext *&ctx) const; private: enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.cc b/src/core/algorithm/diskann/diskann_searcher_entity.cc index c9e49deba..7f4074fef 100644 --- a/src/core/algorithm/diskann/diskann_searcher_entity.cc +++ b/src/core/algorithm/diskann/diskann_searcher_entity.cc @@ -17,6 +17,24 @@ namespace zvec { namespace core { +void DiskAnnSearcherEntity::clear() { + storage_.reset(); + meta_segment_.reset(); + pq_meta_segment_.reset(); + pq_data_segment_.reset(); + vector_segment_.reset(); + key_segment_.reset(); + key_mapping_segment_.reset(); + entrypoint_segment_.reset(); + pq_table_.reset(); + key_buffer_.clear(); + key_mapping_buffer_.clear(); + entrypoints_.clear(); + meta_.clear(); + meta_header_ = {}; + pq_meta_ = {}; +} + const DiskAnnEntity::Pointer DiskAnnSearcherEntity::clone() const { auto meta_segment = meta_segment_->clone(); if (ailego_unlikely(!meta_segment)) { @@ -150,6 +168,12 @@ int DiskAnnSearcherEntity::load_pq_segment() { memcpy(reinterpret_cast(&pq_meta_), data, sizeof(DiskAnnPqMeta)); offset += read_size; + if (pq_meta_.chunk_num == 0 || meta_header_.doc_cnt == 0 || + pq_meta_.full_pivot_data_size == 0 || pq_meta_.centroid_data_size == 0 || + pq_meta_.chunk_num > meta_.dimension()) { + LOG_ERROR("Invalid empty DiskAnn PQ metadata"); + return IndexError_InvalidFormat; + } // 2. read full pivot data std::vector full_pivot_data; @@ -222,9 +246,7 @@ int DiskAnnSearcherEntity::load_pq_segment() { pq_table_ = std::make_shared(meta_, pq_meta_.chunk_num); - pq_table_->init(full_pivot_data, centroid, chunk_offsets, pq_data); - - return 0; + return pq_table_->init(full_pivot_data, centroid, chunk_offsets, pq_data); } int DiskAnnSearcherEntity::load_header_segment() { @@ -398,8 +420,8 @@ const void *DiskAnnSearcherEntity::get_vector(diskann_id_t id) const { const void *vec; if (ailego_unlikely(vector_segment_->read(total_offset, &vec, read_size) != read_size)) { - LOG_ERROR("Read vector from segment failed, id: %u, offset: %lu", id, - total_offset); + LOG_ERROR("Read vector from segment failed, id: %u, offset: %llu", id, + (unsigned long long)total_offset); return nullptr; } diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.h b/src/core/algorithm/diskann/diskann_searcher_entity.h index 953d1117f..439b2b067 100644 --- a/src/core/algorithm/diskann/diskann_searcher_entity.h +++ b/src/core/algorithm/diskann/diskann_searcher_entity.h @@ -34,6 +34,7 @@ class DiskAnnSearcherEntity : public DiskAnnEntity { public: const DiskAnnEntity::Pointer clone() const override; + void clear(); int load(const IndexMeta &meta, IndexStorage::Pointer storage); int load_pq_segment(); int load_header_segment(); diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 82e97dcd6..1752161bd 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -21,18 +21,42 @@ namespace zvec { namespace core { +namespace { + +bool query_meta_matches(const IndexMeta &meta, const IndexQueryMeta &qmeta) { + return qmeta.data_type() == meta.data_type() && + qmeta.dimension() == meta.dimension() && + qmeta.element_size() == meta.element_size(); +} + +bool group_options_valid(const DiskAnnContext *ctx) { + return !ctx->group_by_search() || + (ctx->group_topk() > 0 && ctx->group_by().is_valid()); +} + +} // namespace + DiskAnnStreamer::DiskAnnStreamer() {} DiskAnnStreamer::~DiskAnnStreamer() {} int DiskAnnStreamer::init(const IndexMeta &meta, const ailego::Params &search_params) { + if (state_ == STATE_LOADED) { + LOG_ERROR("Close DiskAnnStreamer before reinitializing it"); + return IndexError_NoReady; + } + meta_ = meta; + params_ = search_params; + list_size_ = 200; + cache_nodes_num_ = 0; log_diskann_io_backend(); - search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); - search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + params_.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); + params_.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + state_ = STATE_INITED; return 0; } @@ -41,6 +65,12 @@ void DiskAnnStreamer::print_debug_info() {} int DiskAnnStreamer::cleanup() { LOG_INFO("Begin DiskAnnStreamer:cleanup"); + unload(); + params_.clear(); + list_size_ = 200; + cache_nodes_num_ = 0; + state_ = STATE_INIT; + LOG_INFO("End DiskAnnStreamer:cleanup"); return 0; @@ -49,6 +79,25 @@ int DiskAnnStreamer::cleanup() { int DiskAnnStreamer::open(IndexStorage::Pointer storage) { LOG_INFO("DiskAnnStreamer::load Begin"); + if (!storage) { + LOG_ERROR("Invalid storage"); + return IndexError_InvalidArgument; + } + if (state_ != STATE_INITED) { + LOG_ERROR("Initialize and close DiskAnnStreamer before opening an index"); + return IndexError_NoReady; + } + + { + std::lock_guard lock(fetch_mutex_); + fetch_ctx_.reset(); + fetch_vector_buffer_.clear(); + diskann_indexer_.reset(); + entity_.clear(); + measure_.reset(); + stats_.clear(); + } + auto start_time = ailego::Monotime::MilliSeconds(); int ret = IndexHelper::DeserializeFromStorage(storage.get(), &meta_); @@ -76,7 +125,10 @@ int DiskAnnStreamer::open(IndexStorage::Pointer storage) { diskann_indexer_->cache_bfs_levels(cache_nodes_num_, node_list); - diskann_indexer_->load_cache_list(node_list); + ret = diskann_indexer_->load_cache_list(node_list); + if (ret != 0) { + return ret; + } node_list.clear(); node_list.shrink_to_fit(); @@ -97,6 +149,7 @@ int DiskAnnStreamer::open(IndexStorage::Pointer storage) { } stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time); + stats_.set_loaded_count(entity_.doc_cnt()); state_ = STATE_LOADED; magic_ = IndexContext::GenerateMagic(); @@ -109,7 +162,19 @@ int DiskAnnStreamer::open(IndexStorage::Pointer storage) { int DiskAnnStreamer::unload() { LOG_INFO("DiskAnnStreamer unload index"); - state_ = STATE_INITED; + const State next_state = state_ == STATE_INIT ? STATE_INIT : STATE_INITED; + { + std::lock_guard lock(fetch_mutex_); + fetch_ctx_.reset(); + fetch_vector_buffer_.clear(); + diskann_indexer_.reset(); + entity_.clear(); + measure_.reset(); + meta_.clear(); + stats_.clear(); + } + + state_ = next_state; return 0; } @@ -125,14 +190,43 @@ int DiskAnnStreamer::update_context(DiskAnnContext *ctx) const { entity, magic_); } +int DiskAnnStreamer::ensure_compatible_context(ContextPointer &context, + DiskAnnContext *&ctx) const { + if (ctx->magic() == magic_) { + return 0; + } + + auto replacement = create_context(); + if (!replacement) { + LOG_ERROR("Failed to recreate context for current streamer"); + return IndexError_Runtime; + } + auto *replacement_ctx = dynamic_cast(replacement.get()); + if (!replacement_ctx) { + LOG_ERROR("Failed to cast recreated DiskAnn context"); + return IndexError_Cast; + } + replacement_ctx->copy_query_options_from(*ctx); + context = std::move(replacement); + ctx = replacement_ctx; + return 0; +} + int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { - // do search + if (ailego_unlikely(state_ != STATE_LOADED)) { + LOG_ERROR("Open DiskAnnStreamer before searching"); + return IndexError_NoReady; + } if (ailego_unlikely(!query || !context)) { LOG_ERROR("The context is not created by this searcher"); return IndexError_Mismatch; } + if (ailego_unlikely(!query_meta_matches(meta_, qmeta))) { + LOG_ERROR("Query meta does not match DiskAnn index meta"); + return IndexError_Mismatch; + } DiskAnnContext *ctx = dynamic_cast(context.get()); ailego_do_if_false(ctx) { @@ -140,18 +234,13 @@ int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, return IndexError_Cast; } - // Context is pooled per index type. When switching between DiskAnn indexes - // with different element sizes (e.g., fp16 vs fp32), the cached context has - // undersized buffers. Recreate it to ensure correct buffer allocations. - if (ctx->magic() != magic_) { - uint32_t saved_topk = ctx->topk(); - context = create_context(); - if (!context) { - LOG_ERROR("Failed to recreate context for current streamer"); - return IndexError_Runtime; - } - ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + int ret = ensure_compatible_context(context, ctx); + if (ret != 0) { + return ret; + } + if (ailego_unlikely(!group_options_valid(ctx))) { + LOG_ERROR("Group search requires a callback and a positive group topk"); + return IndexError_InvalidArgument; } ctx->clear(); @@ -160,7 +249,14 @@ int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, for (uint32_t i = 0; i < count; i++) { ctx->reset_query(query); - diskann_indexer_->knn_search(ctx); + ret = diskann_indexer_->knn_search(ctx); + if (ailego_unlikely(ret != 0)) { + return ret; + } + + if (ailego_unlikely(ctx->error())) { + return IndexError_Runtime; + } ctx->topk_to_result(i); @@ -173,10 +269,18 @@ int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, int DiskAnnStreamer::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (ailego_unlikely(state_ != STATE_LOADED)) { + LOG_ERROR("Open DiskAnnStreamer before searching"); + return IndexError_NoReady; + } if (ailego_unlikely(!query || !context)) { LOG_ERROR("The context is not created by this searcher"); return IndexError_Mismatch; } + if (ailego_unlikely(!query_meta_matches(meta_, qmeta))) { + LOG_ERROR("Query meta does not match DiskAnn index meta"); + return IndexError_Mismatch; + } DiskAnnContext *ctx = dynamic_cast(context.get()); ailego_do_if_false(ctx) { @@ -184,17 +288,13 @@ int DiskAnnStreamer::search_bf_impl(const void *query, return IndexError_Cast; } - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer, recreate it - //! to ensure buffers are correctly sized for this index's parameters. - uint32_t saved_topk = ctx->topk(); - context = create_context(); - if (!context) { - LOG_ERROR("Failed to recreate context for current streamer"); - return IndexError_Runtime; - } - ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + int ret = ensure_compatible_context(context, ctx); + if (ret != 0) { + return ret; + } + if (ailego_unlikely(!group_options_valid(ctx))) { + LOG_ERROR("Group search requires a callback and a positive group topk"); + return IndexError_InvalidArgument; } ctx->clear(); @@ -203,7 +303,10 @@ int DiskAnnStreamer::search_bf_impl(const void *query, for (size_t i = 0; i < count; ++i) { ctx->reset_query(query); - diskann_indexer_->linear_search(ctx); + ret = diskann_indexer_->linear_search(ctx); + if (ailego_unlikely(ret != 0)) { + return ret; + } ctx->topk_to_result(i); @@ -221,10 +324,18 @@ int DiskAnnStreamer::search_bf_by_p_keys_impl( const void *query, const std::vector> &p_keys, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (ailego_unlikely(state_ != STATE_LOADED)) { + LOG_ERROR("Open DiskAnnStreamer before searching"); + return IndexError_NoReady; + } if (ailego_unlikely(!query || !context)) { LOG_ERROR("The context is not created by this searcher"); return IndexError_Mismatch; } + if (ailego_unlikely(!query_meta_matches(meta_, qmeta))) { + LOG_ERROR("Query meta does not match DiskAnn index meta"); + return IndexError_Mismatch; + } DiskAnnContext *ctx = dynamic_cast(context.get()); ailego_do_if_false(ctx) { @@ -237,17 +348,13 @@ int DiskAnnStreamer::search_bf_by_p_keys_impl( return IndexError_InvalidArgument; } - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer, recreate it - //! to ensure buffers are correctly sized for this index's parameters. - uint32_t saved_topk = ctx->topk(); - context = create_context(); - if (!context) { - LOG_ERROR("Failed to recreate context for current streamer"); - return IndexError_Runtime; - } - ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + int ret = ensure_compatible_context(context, ctx); + if (ret != 0) { + return ret; + } + if (ailego_unlikely(!group_options_valid(ctx))) { + LOG_ERROR("Group search requires a callback and a positive group topk"); + return IndexError_InvalidArgument; } ctx->clear(); @@ -256,7 +363,10 @@ int DiskAnnStreamer::search_bf_by_p_keys_impl( for (size_t i = 0; i < count; ++i) { ctx->reset_query(query); - diskann_indexer_->keys_search(p_keys[i], ctx); + ret = diskann_indexer_->keys_search(p_keys[i], ctx); + if (ailego_unlikely(ret != 0)) { + return ret; + } ctx->topk_to_result(i); @@ -272,7 +382,31 @@ int DiskAnnStreamer::search_bf_by_p_keys_impl( int DiskAnnStreamer::get_vector(uint64_t key, Context::Pointer &context, std::string &vector) const { - return diskann_indexer_->get_vector(key, context, vector); + vector.clear(); + if (state_ != STATE_LOADED) { + LOG_ERROR("Open DiskAnnStreamer before fetching vectors"); + return IndexError_NoReady; + } + if (!context) { + LOG_ERROR("Invalid context for get_vector"); + return IndexError_Mismatch; + } + auto *ctx = dynamic_cast(context.get()); + if (!ctx) { + LOG_ERROR("Cast context to DiskAnnContext failed"); + return IndexError_Cast; + } + int ret = ensure_compatible_context(context, ctx); + if (ret != 0) { + return ret; + } + + diskann_id_t id = diskann_indexer_->get_id(key); + if (id == kInvalidId) { + LOG_ERROR("Vector key does not exist: %lu", (unsigned long)key); + return IndexError_NoExist; + } + return diskann_indexer_->get_vector(id, context, vector); } const void *DiskAnnStreamer::get_vector_by_id(uint32_t /*id*/) const { @@ -284,20 +418,48 @@ const void *DiskAnnStreamer::get_vector_by_id(uint32_t /*id*/) const { int DiskAnnStreamer::get_vector_by_id(const uint32_t id, IndexStorage::MemoryBlock &block) const { - // Lazily create a reusable context for fetch operations + if (state_ != STATE_LOADED) { + LOG_ERROR("Open DiskAnnStreamer before fetching vectors"); + return IndexError_NoReady; + } + if (id >= entity_.doc_cnt()) { + LOG_ERROR("Vector id out of range: %u", id); + return IndexError_OutOfRange; + } + + std::lock_guard lock(fetch_mutex_); if (!fetch_ctx_) { fetch_ctx_ = create_context(); if (!fetch_ctx_) { LOG_ERROR("Failed to create context for get_vector_by_id"); return IndexError_Runtime; } + } else { + auto *ctx = dynamic_cast(fetch_ctx_.get()); + if (!ctx) { + LOG_ERROR("Cast fetch context to DiskAnnContext failed"); + return IndexError_Cast; + } + int ret = ensure_compatible_context(fetch_ctx_, ctx); + if (ret != 0) { + return ret; + } } + int ret = diskann_indexer_->get_vector(id, fetch_ctx_, fetch_vector_buffer_); if (ret != 0) { LOG_ERROR("Failed to get vector by id: %u", id); - return IndexError_Runtime; + return ret; } - block.reset((void *)fetch_vector_buffer_.data()); + + void *owned = ailego_malloc(fetch_vector_buffer_.size()); + if (!owned) { + LOG_ERROR("Failed to allocate fetched vector buffer"); + return IndexError_NoMemory; + } + std::memcpy(owned, fetch_vector_buffer_.data(), fetch_vector_buffer_.size()); + block = + IndexStorage::MemoryBlock::MakeOwned(owned, fetch_vector_buffer_.size()); return 0; } @@ -316,6 +478,10 @@ IndexSearcher::Provider::Pointer DiskAnnStreamer::create_provider(void) const { } IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { + if (state_ != STATE_LOADED) { + LOG_ERROR("Open DiskAnnStreamer before creating a context"); + return Context::Pointer(); + } const DiskAnnEntity::Pointer search_ctx_entity = entity_.clone(); if (!search_ctx_entity) { LOG_ERROR("Failed to create search context entity"); @@ -338,6 +504,7 @@ IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { } ctx->set_list_size(list_size_); + ctx->set_magic(magic_); return Context::Pointer(ctx); } diff --git a/src/core/algorithm/diskann/diskann_streamer.h b/src/core/algorithm/diskann/diskann_streamer.h index ddb159ae7..43c36712a 100644 --- a/src/core/algorithm/diskann/diskann_streamer.h +++ b/src/core/algorithm/diskann/diskann_streamer.h @@ -13,6 +13,7 @@ // limitations under the License. #pragma once +#include #include #include "diskann_context.h" #include "diskann_indexer.h" @@ -149,6 +150,8 @@ class DiskAnnStreamer : public IndexStreamer { //! To share ctx across streamer/searcher, we need to update the context for //! current streamer/searcher int update_context(DiskAnnContext *ctx) const; + int ensure_compatible_context(ContextPointer &context, + DiskAnnContext *&ctx) const; private: enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; @@ -166,7 +169,9 @@ class DiskAnnStreamer : public IndexStreamer { DiskAnnIndexer::Pointer diskann_indexer_{nullptr}; DiskAnnSearcherEntity entity_{}; - // Mutable members for get_vector_by_id (caches context and buffer) + // Fetches share the expensive I/O context, while returned MemoryBlocks own + // independent copies so their lifetime does not depend on this buffer. + mutable std::mutex fetch_mutex_; mutable ContextPointer fetch_ctx_{}; mutable std::string fetch_vector_buffer_; diff --git a/src/core/algorithm/diskann/diskann_util.h b/src/core/algorithm/diskann/diskann_util.h index a02130bf0..3e2ecd652 100644 --- a/src/core/algorithm/diskann/diskann_util.h +++ b/src/core/algorithm/diskann/diskann_util.h @@ -35,7 +35,13 @@ class DiskAnnUtil { } static inline void alloc_aligned(void **ptr, size_t size, size_t align) { - *ptr = ::aligned_alloc(align, size); + // C11 aligned_alloc() requires size to be an integral multiple of + // alignment. This is true on Linux (glibc relaxes the requirement) + // but NOT on macOS, where aligned_alloc(32, 16) returns NULL. + // Use posix_memalign() which is portable and has no such restriction. + if (::posix_memalign(ptr, align, size) != 0) { + *ptr = nullptr; + } } static inline void free_aligned(void *ptr) { diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index e377f5342..51f746f6d 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -291,6 +291,17 @@ int DiskAnnIndex::_prepare_for_search( } context->set_topk(diskann_search_param->topk); + context->set_fetch_vector(diskann_search_param->fetch_vector); + if (diskann_search_param->filter) { + context->set_filter(*diskann_search_param->filter); + } else { + context->reset_filter(); + } + if (diskann_search_param->radius > 0.0f) { + context->set_threshold(diskann_search_param->radius); + } else { + context->reset_threshold(); + } // Propagate the query-time beam-search list size into the context. Must be // at least topk to keep enough candidates for a correct result. @@ -338,4 +349,12 @@ int DiskAnnIndex::Merge(const std::vector &indexes, #endif // DISKANN_SUPPORTED +ailego::IOBackendType DiskAnnIndex::io_backend_type() const { + ailego::IOBackendType type = ailego::current_io_backend_type(); + if (type == ailego::IOBackendType::kPread) { + LOG_WARN("%s", ailego::current_io_backend_description().c_str()); + } + return type; +} + } // namespace zvec::core_interface diff --git a/src/db/collection.cc b/src/db/collection.cc index 165e8b717..4cb24a518 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(); @@ -1837,6 +1839,16 @@ Result CollectionImpl::DebugGetHnswStorageMode( Status::NotFound("No HNSW index found for column '", column_name, "'")); } +Result CollectionImpl::DebugGetIoBackendType() const { + switch (ailego::current_io_backend_type()) { + case ailego::IOBackendType::kLibAio: + return std::string("libaio"); + case ailego::IOBackendType::kPread: + return std::string("pread"); + } + return std::string("unknown"); +} + 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..82f4ac7d7 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -198,19 +198,20 @@ Status FieldSchema::validate() const { } if (index_params_->type() == IndexType::DISKANN) { - // DiskAnn requires Linux x86_64/i686/i386. The CMake variable + // 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. + // On Linux, libaio is loaded eagerly (via dlopen); if it is missing, + // DiskAnn falls back to synchronous pread() with degraded performance. + // On macOS, DiskAnn uses kqueue. #if !DISKANN_SUPPORTED return Status::NotSupported( - "DiskAnn is not supported on this platform (Linux x86_64 only)"); + "DiskAnn is not supported on this platform. It is available on " + "Linux (x86_64/ARM64 with libaio) and macOS (with kqueue)."); #endif } diff --git a/src/include/zvec/core/framework/index_context.h b/src/include/zvec/core/framework/index_context.h index 141421453..618005a1e 100644 --- a/src/include/zvec/core/framework/index_context.h +++ b/src/include/zvec/core/framework/index_context.h @@ -233,23 +233,66 @@ class IndexContext { //! Set threshold for RNN void set_threshold(float val) { - if (index_metric_ && index_metric_->support_normalize()) { - index_metric_->denormalize(&val); - } + raw_threshold_ = val; + threshold_set_ = true; + apply_threshold(); + } - threshold_ = val; + //! Retrieve threshold in the caller-facing metric space. + float raw_threshold(void) const { + return raw_threshold_; + } + + //! Retrieve whether an RNN threshold was explicitly configured. + bool threshold_set(void) const { + return threshold_set_; } - //! Retrieve value of threshold for RNN + //! Retrieve the internal threshold used by the search implementation. float threshold(void) const { return threshold_; } //! Reset value of threshold for RNN void reset_threshold(void) { + raw_threshold_ = std::numeric_limits::max(); threshold_ = std::numeric_limits::max(); + threshold_set_ = false; + } + + protected: + //! Replace the metric associated with this context and recompute any + //! configured threshold in the new metric's internal distance space. + void update_index_metric(IndexMetric::Pointer index_metric) { + index_metric_ = std::move(index_metric); + if (threshold_set_) { + apply_threshold(); + } } + //! Copy common per-query options. Thresholds are copied in caller-facing + //! metric space so a pooled context can safely move between indexes that use + //! different internal score normalization. + void copy_common_query_options_from(const IndexContext &rhs) { + filter_ = rhs.filter_; + group_by_ = rhs.group_by_; + if (rhs.threshold_set_) { + set_threshold(rhs.raw_threshold_); + } else { + reset_threshold(); + } + } + + private: + void apply_threshold() { + float val = raw_threshold_; + if (index_metric_ && index_metric_->support_normalize()) { + index_metric_->denormalize(&val); + } + threshold_ = val; + } + + public: //! Generate a global magic number static uint32_t GenerateMagic(void); @@ -262,8 +305,9 @@ class IndexContext { //! Members IndexFilter filter_{}; IndexGroupBy group_by_{}; + float raw_threshold_{std::numeric_limits::max()}; float threshold_{std::numeric_limits::max()}; - + bool threshold_set_{false}; Profiler profiler_{}; diff --git a/src/include/zvec/core/framework/index_storage.h b/src/include/zvec/core/framework/index_storage.h index 049c79917..7f5f533ae 100644 --- a/src/include/zvec/core/framework/index_storage.h +++ b/src/include/zvec/core/framework/index_storage.h @@ -15,6 +15,7 @@ #pragma once #include +#include #include #include #include @@ -123,6 +124,7 @@ class IndexStorage : public IndexModule { deep_copy_from(rhs); break; default: + release_current(); break; } } @@ -151,6 +153,7 @@ class IndexStorage : public IndexModule { rhs.type_ = MemoryBlockType::MBT_UNKNOWN; break; default: + release_current(); break; } } @@ -253,6 +256,11 @@ class IndexStorage : public IndexModule { scratch_size_ = rhs.scratch_size_; if (scratch_size_ > 0 && rhs.data_) { data_ = ailego_malloc(scratch_size_); + if (!data_) { + type_ = MemoryBlockType::MBT_UNKNOWN; + scratch_size_ = 0; + throw std::bad_alloc(); + } std::memcpy(data_, rhs.data_, scratch_size_); } else { data_ = nullptr; diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 935a6569d..36c7e76f8 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: int CreateAndInitStreamer(const BaseIndexParam ¶m) override; diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index c4b77575f..8b1f680ce 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 + //! "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/core/algorithm/diskann/diskann_builder_test.cc b/tests/core/algorithm/diskann/diskann_builder_test.cc index 752759e9d..0f9a8ea7e 100644 --- a/tests/core/algorithm/diskann/diskann_builder_test.cc +++ b/tests/core/algorithm/diskann/diskann_builder_test.cc @@ -17,11 +17,13 @@ #include #include #include +#include #include #include #include #include #include "diskann_holder.h" +#include "diskann_params.h" using namespace zvec::core; using namespace zvec::ailego; @@ -143,6 +145,56 @@ TEST_F(DiskAnnBuilderTest, SmallDatasetBuildTime) { << " ms — likely a lost-wakeup regression in progress loops."; } +TEST_F(DiskAnnBuilderTest, MemoryLimitCapsPqChunkCount) { + constexpr size_t kTestDim = 8; + constexpr size_t kDocCnt = 16; + + IndexMeta meta(IndexMeta::DataType::DT_FP32, kTestDim); + meta.set_metric("SquaredEuclidean", 0, Params()); + + auto holder = + make_shared>(kTestDim); + for (size_t i = 0; i < kDocCnt; ++i) { + NumericalVector vec(kTestDim, static_cast(i)); + ASSERT_TRUE(holder->emplace(i, vec)); + } + + Params params; + params.set(PARAM_DISKANN_BUILDER_MAX_DEGREE, 16); + params.set(PARAM_DISKANN_BUILDER_LIST_SIZE, 20); + params.set(PARAM_DISKANN_BUILDER_MAX_PQ_CHUNK_NUM, 4); + params.set(PARAM_DISKANN_BUILDER_THREAD_COUNT, 2); + // 40 total bytes gives a two-byte PQ code budget per vector. Before the + // fix this value was calculated and then discarded. + params.set(PARAM_DISKANN_BUILDER_MEMORY_LIMIT, + 40.0 / (1024.0 * 1024.0 * 1024.0)); + + auto builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); + ASSERT_NE(builder, nullptr); + ASSERT_EQ(0, builder->init(meta, params)); + ASSERT_EQ(0, builder->train(holder)); + ASSERT_EQ(0, builder->build(holder)); + + const string path = _dir + "/MemoryLimitCapsPqChunkCount"; + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(dumper, nullptr); + ASSERT_EQ(0, dumper->create(path)); + ASSERT_EQ(0, builder->dump(dumper)); + ASSERT_EQ(0, dumper->close()); + + auto storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(path, false)); + auto pq_meta_segment = storage->get(DiskAnnEntity::kDiskAnnPqMetaSegmentId); + ASSERT_NE(pq_meta_segment, nullptr); + const void *data = nullptr; + ASSERT_EQ(sizeof(DiskAnnPqMeta), + pq_meta_segment->read(0, &data, sizeof(DiskAnnPqMeta))); + DiskAnnPqMeta pq_meta{}; + std::memcpy(&pq_meta, data, sizeof(pq_meta)); + EXPECT_EQ(2U, pq_meta.chunk_num); +} + TEST_F(DiskAnnBuilderTest, TestImplicitFactoryRegistration) { IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); ASSERT_NE(builder, nullptr) diff --git a/tests/core/algorithm/diskann/diskann_file_reader_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_test.cc new file mode 100644 index 000000000..8cc853671 --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_file_reader_test.cc @@ -0,0 +1,169 @@ +// 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. + +#include "diskann_file_reader.h" +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec::core; + +namespace { + +constexpr size_t kPageSize = 4096; + +class TemporaryFile { + public: + TemporaryFile() : fd_(::mkstemp(path_)) {} + + ~TemporaryFile() { + close(); + ::unlink(path_); + } + + int fd() const { + return fd_; + } + + const char *path() const { + return path_; + } + + void close() { + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } + } + + bool write_all(const void *data, size_t size) { + const auto *bytes = static_cast(data); + size_t written = 0; + while (written < size) { + ssize_t ret = ::pwrite(fd_, bytes + written, size - written, written); + if (ret < 0 && errno == EINTR) { + continue; + } + if (ret <= 0) { + return false; + } + written += static_cast(ret); + } + return ::fsync(fd_) == 0; + } + + private: + char path_[64] = "/tmp/zvec-diskann-reader-XXXXXX"; + int fd_; +}; + +using AlignedBuffer = std::unique_ptr; + +AlignedBuffer make_aligned_buffer(size_t size) { + void *buffer = nullptr; + if (::posix_memalign(&buffer, kPageSize, size) != 0) { + return AlignedBuffer(nullptr, &std::free); + } + std::memset(buffer, 0, size); + return AlignedBuffer(buffer, &std::free); +} + +} // namespace + +TEST(DiskAnnFileReaderTest, BatchAlignedReadsPreserveRequestOrder) { + constexpr size_t kPageCount = 32; + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + std::vector source(kPageSize * kPageCount); + for (size_t page = 0; page < kPageCount; ++page) { + std::memset(source.data() + page * kPageSize, + static_cast(page + 1), kPageSize); + } + ASSERT_TRUE(file.write_all(source.data(), source.size())); + file.close(); + + AlignedBuffer output = make_aligned_buffer(source.size()); + ASSERT_NE(output, nullptr); + + std::vector requests; + requests.reserve(kPageCount); + for (size_t i = 0; i < kPageCount; ++i) { + const size_t source_page = (i * 7) % kPageCount; + requests.emplace_back( + source_page * kPageSize, kPageSize, + static_cast(output.get()) + i * kPageSize); + } + + LinuxAlignedFileReader reader; + reader.open(file.path()); + + IOContext ctx{}; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_EQ(reader.read(requests, ctx, false), 0); + + for (size_t i = 0; i < kPageCount; ++i) { + const size_t source_page = (i * 7) % kPageCount; + const auto *page = + static_cast(output.get()) + i * kPageSize; + for (size_t byte = 0; byte < kPageSize; ++byte) { + ASSERT_EQ(page[byte], static_cast(source_page + 1)); + } + } + + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); +} + +TEST(DiskAnnFileReaderTest, ShortReadReturnsError) { + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + std::vector source(kPageSize, 0x5a); + ASSERT_TRUE(file.write_all(source.data(), source.size())); + file.close(); + + AlignedBuffer output = make_aligned_buffer(kPageSize * 2); + ASSERT_NE(output, nullptr); + std::vector requests{ + {0, kPageSize * 2, output.get()}, + }; + + LinuxAlignedFileReader reader; + reader.open(file.path()); + + IOContext ctx{}; + ASSERT_EQ(setup_io_ctx(ctx), 0); + EXPECT_NE(reader.read(requests, ctx, false), 0); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); +} + +TEST(DiskAnnFileReaderTest, ReadBeforeOpenReturnsError) { + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector requests{ + {0, kPageSize, output.get()}, + }; + + LinuxAlignedFileReader reader; + IOContext ctx{}; + EXPECT_NE(reader.read(requests, ctx, false), 0); +} diff --git a/tests/core/algorithm/diskann/diskann_searcher_test.cc b/tests/core/algorithm/diskann/diskann_searcher_test.cc index ff432886f..a560788d1 100644 --- a/tests/core/algorithm/diskann/diskann_searcher_test.cc +++ b/tests/core/algorithm/diskann/diskann_searcher_test.cc @@ -16,6 +16,10 @@ #include #include #include +#include +#include +#include +#include #include #include #include @@ -220,6 +224,85 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { EXPECT_GT(recall, 0.90f); EXPECT_GT(topk1Recall, 0.80f); EXPECT_GT(cost, 2.0f); + + // A context created by the streamer must carry the streamer's magic so it + // can be reused instead of being recreated on every search. + IndexStreamer::Pointer streamer = + IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(streamer, nullptr); + ASSERT_EQ(0, streamer->init(*_index_meta_ptr, search_params)); + + auto streamer_storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_EQ(0, streamer_storage->open(path, false)); + ASSERT_EQ(0, streamer->open(streamer_storage)); + + auto streamer_ctx = streamer->create_context(); + ASSERT_NE(streamer_ctx, nullptr); + streamer_ctx->set_topk(topk); + auto *original_ctx = streamer_ctx.get(); + + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); + EXPECT_EQ(original_ctx, streamer_ctx.get()); + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); + EXPECT_EQ(original_ctx, streamer_ctx.get()); + + // Concurrent fetches must not share an I/O context or return pointers into + // a mutable streamer-owned string. + std::atomic fetch_ok{true}; + std::vector fetch_threads; + for (uint32_t thread_id = 0; thread_id < 4; ++thread_id) { + fetch_threads.emplace_back([&, thread_id]() { + for (uint32_t i = thread_id; i < 40; i += 4) { + IndexStorage::MemoryBlock block; + if (streamer->get_vector_by_id(i, block) != 0 || + block.data() == nullptr || + *static_cast(block.data()) != + static_cast(i)) { + fetch_ok.store(false); + return; + } + } + }); + } + for (auto &thread : fetch_threads) { + thread.join(); + } + EXPECT_TRUE(fetch_ok.load()); + + // Query metadata controls both the copy size and the batch stride, so a + // mismatch must be rejected before either searcher touches the input. + IndexQueryMeta wrong_type(IndexMeta::DataType::DT_FP16, dim); + EXPECT_EQ(IndexError_Mismatch, + searcher->search_impl(vec.data(), wrong_type, knnCtx)); + EXPECT_EQ(IndexError_Mismatch, + streamer->search_impl(vec.data(), wrong_type, streamer_ctx)); + IndexQueryMeta wrong_dimension(IndexMeta::DataType::DT_FP32, dim - 1); + EXPECT_EQ(IndexError_Mismatch, + searcher->search_bf_impl(vec.data(), wrong_dimension, linearCtx)); + + // Group parameters without a grouping callback are invalid instead of a + // successful search with an empty result. + auto invalid_group_ctx = searcher->create_context(); + ASSERT_NE(invalid_group_ctx, nullptr); + invalid_group_ctx->set_group_params(2, 3); + EXPECT_EQ(IndexError_InvalidArgument, + searcher->search_impl(vec.data(), qmeta, invalid_group_ctx)); + + // I/O failures from the indexer must be propagated by the streamer instead + // of being converted into a successful search with incomplete results. + ASSERT_EQ(0, ::truncate(path.c_str(), 0)); + EXPECT_NE(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); + + // Closing/unloading releases the index and makes all query entry points + // reject work until another index is loaded. + ASSERT_EQ(0, streamer->close()); + EXPECT_EQ(nullptr, streamer->create_context()); + EXPECT_EQ(IndexError_NoReady, + streamer->search_impl(vec.data(), qmeta, streamer_ctx)); + ASSERT_EQ(0, searcher->unload()); + EXPECT_EQ(nullptr, searcher->create_context()); + EXPECT_EQ(IndexError_NoReady, + searcher->search_impl(vec.data(), qmeta, knnCtx)); } TEST_F(DiskAnnSearcherTest, TestNodeCache) { @@ -539,7 +622,6 @@ TEST_F(DiskAnnSearcherTest, TestGroup) { NumericalVector vec(dim); IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); size_t group_topk = 20; - uint64_t total_time = 0; auto groupbyFunc = [](uint64_t key) { uint32_t group_id = key / 10 % 10; @@ -559,29 +641,62 @@ TEST_F(DiskAnnSearcherTest, TestGroup) { vec[j] = query_value / 10 + 0.1f; } - auto t1 = Realtime::MicroSeconds(); ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); - auto t2 = Realtime::MicroSeconds(); - - total_time += t2 - t1; auto &group_result = ctx->group_result(); + ASSERT_EQ(group_num, group_result.size()); + std::set seen_group_ids; for (uint32_t i = 0; i < group_result.size(); ++i) { const std::string &group_id = group_result[i].group_id(); auto &result = group_result[i].docs(); + ASSERT_TRUE(seen_group_ids.insert(group_id).second); ASSERT_GT(result.size(), 0); + ASSERT_LE(result.size(), group_topk); std::cout << "Group ID: " << group_id << std::endl; for (uint32_t j = 0; j < result.size(); ++j) { + EXPECT_EQ(group_id, groupbyFunc(result[j].key())); std::cout << "\tKey: " << result[j].key() << std::fixed << std::setprecision(3) << ", Score: " << result[j].score() << std::endl; } } -#if 0 + // Reusing a group context must not retain scores or documents from the + // previous query. + query_value = doc_cnt / 10; + for (size_t j = 0; j < dim; ++j) { + vec[j] = query_value / 10 + 0.1f; + } + ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); + const auto &reused_group_result = ctx->group_result(); + ASSERT_EQ(group_num, reused_group_result.size()); + for (const auto &group : reused_group_result) { + for (const auto &doc : group.docs()) { + float delta = static_cast(doc.key()) / 10.0f - vec[0]; + float expected_score = dim * delta * delta; + EXPECT_NEAR(expected_score, doc.score(), + std::max(1e-3f, expected_score * 1e-4f)); + } + } + + // Full linear group search must maintain a heap for every group while it + // scans, rather than grouping only the global top-k afterward. + auto linear_ctx = searcher->create_context(); + linear_ctx->set_group_params(group_num, group_topk); + linear_ctx->set_group_by(groupbyFunc); + ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linear_ctx)); + const auto &linear_group_result = linear_ctx->group_result(); + ASSERT_EQ(group_num, linear_group_result.size()); + for (const auto &group : linear_group_result) { + ASSERT_EQ(group_topk, group.docs().size()); + for (const auto &doc : group.docs()) { + EXPECT_EQ(group.group_id(), groupbyFunc(doc.key())); + } + } + // do linear search by p_keys test auto groupbyFuncLinear = [](uint64_t key) { uint32_t group_id = key % 10; @@ -618,7 +733,6 @@ TEST_F(DiskAnnSearcherTest, TestGroup) { ASSERT_EQ(10 - i, result[0].key()); } -#endif } TEST_F(DiskAnnSearcherTest, TestFetchVector) { @@ -628,12 +742,13 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) { auto holder = make_shared>(dim); size_t doc_cnt = 10000UL; + auto key_for_id = [](size_t id) { return 100000UL + id * 3; }; for (size_t i = 0; i < doc_cnt; i++) { NumericalVector vec(dim); for (size_t j = 0; j < dim; ++j) { vec[j] = i; } - ASSERT_TRUE(holder->emplace(i, vec)); + ASSERT_TRUE(holder->emplace(key_for_id(i), vec)); } Params params; @@ -687,7 +802,7 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) { for (size_t i = 0; i < doc_cnt; i += doc_cnt / 10) { std::string vec_value; - ASSERT_EQ(0, searcher->get_vector(i, linearCtx, vec_value)); + ASSERT_EQ(0, searcher->get_vector(key_for_id(i), linearCtx, vec_value)); float vector_value = *(const float *)(vec_value.data()); ASSERT_EQ(vector_value, i); @@ -696,8 +811,6 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) { size_t topk = 200; linearCtx->set_topk(topk); knnCtx->set_topk(topk); - uint64_t knnTotalTime = 0; - uint64_t linearTotalTime = 0; IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); @@ -707,25 +820,87 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) { vec[j] = i; } - auto t1 = Realtime::MicroSeconds(); ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); - auto t2 = Realtime::MicroSeconds(); ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); - auto t3 = Realtime::MicroSeconds(); - knnTotalTime += t2 - t1; - linearTotalTime += t3 - t2; auto &knnResult = knnCtx->result(); ASSERT_EQ(topk, knnResult.size()); auto &linearResult = linearCtx->result(); ASSERT_EQ(topk, linearResult.size()); - ASSERT_EQ(i, linearResult[0].key()); + ASSERT_EQ(key_for_id(i), linearResult[0].key()); ASSERT_NE(knnResult[0].vector_string(), ""); float vector_value = *((float *)(knnResult[0].vector_string().data())); ASSERT_EQ(vector_value, i); } + + std::string missing_vector; + EXPECT_EQ(IndexError_NoExist, + searcher->get_vector(42, linearCtx, missing_vector)); + EXPECT_TRUE(missing_vector.empty()); + + ASSERT_EQ(0, ::truncate(path.c_str(), 0)); + std::string vector_after_truncate; + EXPECT_EQ(IndexError_Runtime, + searcher->get_vector(key_for_id(doc_cnt - 1), linearCtx, + vector_after_truncate)); + EXPECT_TRUE(vector_after_truncate.empty()); +} + +TEST_F(DiskAnnSearcherTest, TestFp16Entrypoint) { + IndexMeta fp16_meta(IndexMeta::DataType::DT_FP16, dim); + fp16_meta.set_metric("SquaredEuclidean", 0, Params()); + + auto holder = + make_shared>(dim); + constexpr size_t doc_cnt = 2000; + for (size_t i = 0; i < doc_cnt; ++i) { + NumericalVector vec(dim); + for (size_t j = 0; j < dim; ++j) { + vec[j] = static_cast(i) / 10.0f; + } + ASSERT_TRUE(holder->emplace(i, vec)); + } + + Params params; + params.set("zvec.diskann.builder.max_degree", 32); + params.set("zvec.diskann.builder.list_size", 100); + params.set("zvec.diskann.builder.max_pq_chunk_num", 32); + params.set("zvec.diskann.builder.threads", 2); + + auto builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); + ASSERT_NE(builder, nullptr); + ASSERT_EQ(0, builder->init(fp16_meta, params)); + ASSERT_EQ(0, builder->train(holder)); + ASSERT_EQ(0, builder->build(holder)); + + const string path = _dir + "/TestFp16Entrypoint"; + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(dumper, nullptr); + ASSERT_EQ(0, dumper->create(path)); + ASSERT_EQ(0, builder->dump(dumper)); + ASSERT_EQ(0, dumper->close()); + + auto searcher = IndexFactory::CreateSearcher("DiskAnnSearcher"); + ASSERT_NE(searcher, nullptr); + ASSERT_EQ(0, searcher->init(params)); + auto storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_EQ(0, storage->open(path, false)); + ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); + + auto ctx = searcher->create_context(); + ASSERT_NE(ctx, nullptr); + ctx->set_topk(10); + + NumericalVector query(dim); + for (size_t j = 0; j < dim; ++j) { + query[j] = 123.1f; + } + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP16, dim); + ASSERT_EQ(0, searcher->search_impl(query.data(), qmeta, ctx)); + ASSERT_EQ(10, ctx->result().size()); + EXPECT_EQ(1231UL, ctx->result()[0].key()); } TEST_F(DiskAnnSearcherTest, TestRnnSearch) { diff --git a/tests/core/interface/index_group_by_test.cc b/tests/core/interface/index_group_by_test.cc index bb41938fb..3c6001ab2 100644 --- a/tests/core/interface/index_group_by_test.cc +++ b/tests/core/interface/index_group_by_test.cc @@ -537,3 +537,142 @@ TEST_F(GroupByInterfaceTest, UnsupportedIndexTypes) { RunRejected(test_case); } } + +#if DISKANN_SUPPORTED +TEST(DiskAnnInterfaceTest, PropagatesCommonQueryOptions) { + const std::string source_name = "test_diskann_query_options_source"; + const std::string index_name = "test_diskann_query_options"; + zvec::test_util::RemoveTestFiles(source_name + "*"); + zvec::test_util::RemoveTestFiles(index_name + "*"); + + auto source_param = FlatIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .Build(); + auto source = IndexFactory::CreateAndInitIndex(*source_param); + ASSERT_NE(source, nullptr); + ASSERT_EQ( + 0, source->Open(source_name, {StorageOptions::StorageType::kMMAP, true})); + for (uint32_t key = 0; key < kNumDocs; ++key) { + std::vector values(kDimension, static_cast(key)); + ASSERT_EQ(0, source->Add(VectorData{DenseVector{values.data()}}, key)); + } + ASSERT_EQ(0, source->Train()); + + auto diskann_param = DiskAnnIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithListSize(kSearchTopk) + .WithPqChunkNum(0) + .Build(); + auto index = IndexFactory::CreateAndInitIndex(*diskann_param); + ASSERT_NE(index, nullptr); + ASSERT_EQ( + 0, index->Open(index_name, {StorageOptions::StorageType::kMMAP, true})); + ASSERT_EQ(0, index->Merge({source}, IndexFilter())); + + std::vector query_values(kDimension, 3.0f); + VectorData query{DenseVector{query_values.data()}}; + + auto filtered_query = std::make_shared(); + filtered_query->topk = 6; + filtered_query->list_size = kSearchTopk; + filtered_query->fetch_vector = true; + filtered_query->filter = std::make_shared(); + filtered_query->filter->set([](uint64_t key) { return key == 3; }); + + SearchResult filtered_result; + ASSERT_EQ(0, index->Search(query, filtered_query, &filtered_result)); + ASSERT_FALSE(filtered_result.doc_list_.empty()); + for (const auto &doc : filtered_result.doc_list_) { + EXPECT_NE(3UL, doc.key()); + ASSERT_EQ(kDimension * sizeof(float), doc.vector_string().size()); + const auto *vector = + reinterpret_cast(doc.vector_string().data()); + for (uint32_t i = 0; i < kDimension; ++i) { + EXPECT_FLOAT_EQ(static_cast(doc.key()), vector[i]); + } + } + + auto radius_query = std::make_shared(); + radius_query->topk = kNumDocs; + radius_query->list_size = kSearchTopk; + radius_query->radius = 1.0f; + + SearchResult radius_result; + ASSERT_EQ(0, index->Search(query, radius_query, &radius_result)); + ASSERT_EQ(1, radius_result.doc_list_.size()); + EXPECT_EQ(3UL, radius_result.doc_list_[0].key()); + EXPECT_LE(radius_result.doc_list_[0].score(), radius_query->radius); + + ASSERT_EQ(0, index->Close()); + ASSERT_EQ(0, source->Close()); + + // The thread-local context pool is shared by all DiskAnn indexes. Reusing + // the prior L2 context for an InnerProduct index must copy the caller-facing + // radius and denormalize it with the new metric. + const std::string ip_source_name = "test_diskann_query_options_ip_source"; + const std::string ip_index_name = "test_diskann_query_options_ip"; + zvec::test_util::RemoveTestFiles(ip_source_name + "*"); + zvec::test_util::RemoveTestFiles(ip_index_name + "*"); + + auto ip_source_param = FlatIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .Build(); + auto ip_source = IndexFactory::CreateAndInitIndex(*ip_source_param); + ASSERT_NE(ip_source, nullptr); + ASSERT_EQ(0, ip_source->Open(ip_source_name, + {StorageOptions::StorageType::kMMAP, true})); + for (uint32_t key = 0; key < kNumDocs; ++key) { + std::vector values(kDimension, static_cast(key)); + ASSERT_EQ(0, ip_source->Add(VectorData{DenseVector{values.data()}}, key)); + } + ASSERT_EQ(0, ip_source->Train()); + + auto ip_diskann_param = DiskAnnIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithListSize(kNumDocs) + .WithPqChunkNum(0) + .Build(); + auto ip_index = IndexFactory::CreateAndInitIndex(*ip_diskann_param); + ASSERT_NE(ip_index, nullptr); + ASSERT_EQ(0, ip_index->Open(ip_index_name, + {StorageOptions::StorageType::kMMAP, true})); + ASSERT_EQ(0, ip_index->Merge({ip_source}, IndexFilter())); + + std::vector ip_query_values(kDimension, 1.0f); + VectorData ip_query{DenseVector{ip_query_values.data()}}; + auto ip_radius_query = std::make_shared(); + ip_radius_query->topk = kNumDocs; + ip_radius_query->list_size = kNumDocs; + ip_radius_query->is_linear = true; + ip_radius_query->radius = static_cast(kDimension * 8); + + SearchResult ip_radius_result; + ASSERT_EQ(0, ip_index->Search(ip_query, ip_radius_query, &ip_radius_result)); + ASSERT_EQ(kNumDocs - 8, ip_radius_result.doc_list_.size()); + for (const auto &doc : ip_radius_result.doc_list_) { + EXPECT_GE(doc.key(), 8UL); + EXPECT_GE(doc.score(), ip_radius_query->radius); + } + + ASSERT_EQ(0, ip_index->Close()); + ASSERT_EQ(0, ip_source->Close()); + zvec::test_util::RemoveTestFiles(source_name + "*"); + zvec::test_util::RemoveTestFiles(index_name + "*"); + zvec::test_util::RemoveTestFiles(ip_source_name + "*"); + zvec::test_util::RemoveTestFiles(ip_index_name + "*"); +} +#endif diff --git a/tests/db/crash_recovery/write_recovery_test.cc b/tests/db/crash_recovery/write_recovery_test.cc index 702e9e498..9f878b1eb 100644 --- a/tests/db/crash_recovery/write_recovery_test.cc +++ b/tests/db/crash_recovery/write_recovery_test.cc @@ -188,7 +188,7 @@ TEST_F(CrashRecoveryTest, CrashRecoveryDuringInsertion) { collection.reset(); } - RunGeneratorAndCrash("0", "10000", "insert", "0", 3); + RunGeneratorAndCrash("0", "10000", "insert", "0", 5); auto result = Collection::Open(dir_path_, options_); ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. " @@ -196,7 +196,7 @@ TEST_F(CrashRecoveryTest, CrashRecoveryDuringInsertion) { auto collection = result.value(); uint64_t doc_count{collection->Stats().value().doc_count}; ASSERT_GT(doc_count, 800) - << "Document count is too low after 3s of insertion and recovery"; + << "Document count is too low after 5s of insertion and recovery"; for (uint64_t doc_id = 0; doc_id < doc_count; doc_id++) { const auto expected_doc = CreateTestDoc(doc_id, 0);