From 2b16ed39cd4368581225061583e7c83518b3c53d Mon Sep 17 00:00:00 2001 From: egolearner Date: Thu, 23 Apr 2026 10:20:46 +0800 Subject: [PATCH] chore: rm hnsw-rabitq searcher/streamer (#367) * chore: rm hnsw-rabitq searcher/streamer * fix clang-tidy * fix --- .github/workflows/clang_tidy.yml | 2 +- src/core/algorithm/hnsw/hnsw_streamer.cc | 2 - .../hnsw_rabitq/hnsw_rabitq_builder.cc | 554 ----------------- .../hnsw_rabitq/hnsw_rabitq_builder.h | 101 --- .../hnsw_rabitq/hnsw_rabitq_builder_entity.cc | 199 ------ .../hnsw_rabitq/hnsw_rabitq_builder_entity.h | 139 ----- .../hnsw_rabitq/hnsw_rabitq_register.cc | 4 - .../hnsw_rabitq/hnsw_rabitq_searcher.cc | 516 ---------------- .../hnsw_rabitq/hnsw_rabitq_searcher.h | 142 ----- .../hnsw_rabitq_searcher_entity.cc | 515 ---------------- .../hnsw_rabitq/hnsw_rabitq_searcher_entity.h | 158 ----- .../hnsw_rabitq/hnsw_rabitq_streamer.cc | 3 - .../hnsw_rabitq/hnsw_rabitq_builder_test.cc | 428 ------------- .../hnsw_rabitq/hnsw_rabitq_searcher_test.cc | 573 ------------------ tools/core/local_builder.cc | 2 - tools/core/local_builder_original.cc | 2 - 16 files changed, 1 insertion(+), 3339 deletions(-) delete mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.cc delete mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.h delete mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.cc delete mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.h delete mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.cc delete mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.h delete mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.cc delete mode 100644 src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.h delete mode 100644 tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_test.cc delete mode 100644 tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_test.cc diff --git a/.github/workflows/clang_tidy.yml b/.github/workflows/clang_tidy.yml index 49e85eb..0b595b6 100644 --- a/.github/workflows/clang_tidy.yml +++ b/.github/workflows/clang_tidy.yml @@ -27,7 +27,7 @@ jobs: - name: Install dependencies run: | sudo apt-get update - sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build + sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build libomp-dev - name: Configure CMake and export compile commands run: | diff --git a/src/core/algorithm/hnsw/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc index 5804c7d..82319e2 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -342,8 +342,6 @@ int HnswStreamer::dump(const IndexDumper::Pointer &dumper) { shared_mutex_.lock(); AILEGO_DEFER([&]() { shared_mutex_.unlock(); }); - meta_.set_searcher("HnswSearcher", HnswEntity::kRevision, ailego::Params()); - int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); if (ret != 0) { LOG_ERROR("Failed to serialize meta into dumper."); diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.cc deleted file mode 100644 index 235fb20..0000000 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.cc +++ /dev/null @@ -1,554 +0,0 @@ -// 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 "hnsw_rabitq_builder.h" -#include -#include -#include -#include -#include -#include -#include -#include "zvec/core/framework/index_error.h" -#include "zvec/core/framework/index_factory.h" -#include "zvec/core/framework/index_logger.h" -#include "zvec/core/framework/index_memory.h" -#include "zvec/core/framework/index_meta.h" -#include "zvec/core/framework/index_provider.h" -#include "hnsw_rabitq_algorithm.h" -#include "hnsw_rabitq_entity.h" -#include "hnsw_rabitq_params.h" -#include "rabitq_converter.h" -#include "rabitq_params.h" -#include "rabitq_reformer.h" - -namespace zvec { -namespace core { - -HnswRabitqBuilder::HnswRabitqBuilder() {} - -int HnswRabitqBuilder::init(const IndexMeta &meta, - const ailego::Params ¶ms) { - LOG_INFO("Begin HnswRabitqBuilder::init"); - - meta_ = meta; - auto params_copy = params; - meta_.set_builder("HnswRabitqBuilder", HnswRabitqEntity::kRevision, - std::move(params_copy)); - - size_t memory_quota = 0UL; - params.get(PARAM_HNSW_RABITQ_BUILDER_MEMORY_QUOTA, &memory_quota); - params.get(PARAM_HNSW_RABITQ_BUILDER_THREAD_COUNT, &thread_cnt_); - params.get(PARAM_HNSW_RABITQ_BUILDER_MIN_NEIGHBOR_COUNT, &min_neighbor_cnt_); - params.get(PARAM_HNSW_RABITQ_BUILDER_EFCONSTRUCTION, &ef_construction_); - params.get(PARAM_HNSW_RABITQ_BUILDER_CHECK_INTERVAL_SECS, - &check_interval_secs_); - - params.get(PARAM_HNSW_RABITQ_BUILDER_MAX_NEIGHBOR_COUNT, - &upper_max_neighbor_cnt_); - float multiplier = HnswRabitqEntity::kDefaultL0MaxNeighborCntMultiplier; - params.get(PARAM_HNSW_RABITQ_BUILDER_L0_MAX_NEIGHBOR_COUNT_MULTIPLIER, - &multiplier); - l0_max_neighbor_cnt_ = multiplier * upper_max_neighbor_cnt_; - scaling_factor_ = upper_max_neighbor_cnt_; - params.get(PARAM_HNSW_RABITQ_BUILDER_SCALING_FACTOR, &scaling_factor_); - - multiplier = HnswRabitqEntity::kDefaultNeighborPruneMultiplier; - params.get(PARAM_HNSW_RABITQ_BUILDER_NEIGHBOR_PRUNE_MULTIPLIER, &multiplier); - size_t prune_cnt = multiplier * upper_max_neighbor_cnt_; - - if (ef_construction_ == 0) { - ef_construction_ = HnswRabitqEntity::kDefaultEfConstruction; - } - if (upper_max_neighbor_cnt_ == 0) { - upper_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt; - } - if (upper_max_neighbor_cnt_ > kMaxNeighborCnt) { - LOG_ERROR("[%s] must be in range (0,%d]", - PARAM_HNSW_RABITQ_BUILDER_MAX_NEIGHBOR_COUNT.c_str(), - kMaxNeighborCnt); - return IndexError_InvalidArgument; - } - if (min_neighbor_cnt_ > upper_max_neighbor_cnt_) { - LOG_ERROR("[%s]-[%d] must be <= [%s]-[%d]", - PARAM_HNSW_RABITQ_BUILDER_MIN_NEIGHBOR_COUNT.c_str(), - min_neighbor_cnt_, - PARAM_HNSW_RABITQ_BUILDER_MAX_NEIGHBOR_COUNT.c_str(), - upper_max_neighbor_cnt_); - return IndexError_InvalidArgument; - } - if (l0_max_neighbor_cnt_ == 0) { - l0_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt; - } - if (l0_max_neighbor_cnt_ > HnswRabitqEntity::kMaxNeighborCnt) { - LOG_ERROR("L0MaxNeighborCnt must be in range (0,%d)", - HnswRabitqEntity::kMaxNeighborCnt); - return IndexError_InvalidArgument; - } - if (scaling_factor_ == 0U) { - scaling_factor_ = HnswRabitqEntity::kDefaultScalingFactor; - } - if (scaling_factor_ < 5 || scaling_factor_ > 1000) { - LOG_ERROR("[%s] must be in range [5,1000]", - PARAM_HNSW_RABITQ_BUILDER_SCALING_FACTOR.c_str()); - return IndexError_InvalidArgument; - } - if (thread_cnt_ == 0) { - thread_cnt_ = std::thread::hardware_concurrency(); - } - if (thread_cnt_ > std::thread::hardware_concurrency()) { - LOG_WARN("[%s] greater than cpu cores %zu", - PARAM_HNSW_RABITQ_BUILDER_THREAD_COUNT.c_str(), - static_cast(std::thread::hardware_concurrency())); - } - if (prune_cnt == 0UL) { - prune_cnt = upper_max_neighbor_cnt_; - } - - metric_ = IndexFactory::CreateMetric(meta_.metric_name()); - if (!metric_) { - LOG_ERROR("CreateMetric failed, name: %s", meta_.metric_name().c_str()); - return IndexError_NoExist; - } - int ret = metric_->init(meta_, meta_.metric_params()); - if (ret != 0) { - LOG_ERROR("IndexMetric init failed, ret=%d", ret); - return ret; - } - - uint32_t total_bits = 0; - params.get(PARAM_RABITQ_TOTAL_BITS, &total_bits); - if (total_bits == 0) { - total_bits = kDefaultRabitqTotalBits; - } - if (total_bits < 1 || total_bits > 9) { - LOG_ERROR("Invalid total_bits: %zu, must be in [1, 9]", (size_t)total_bits); - return IndexError_InvalidArgument; - } - uint8_t ex_bits = total_bits - 1; - entity_.set_ex_bits(ex_bits); - - uint32_t dimension = 0; - params.get(PARAM_HNSW_RABITQ_GENERAL_DIMENSION, &dimension); - if (dimension == 0) { - LOG_ERROR("%s not set", PARAM_HNSW_RABITQ_GENERAL_DIMENSION.c_str()); - return IndexError_InvalidArgument; - } - if (dimension < kMinRabitqDimSize || dimension > kMaxRabitqDimSize) { - LOG_ERROR("Invalid dimension: %u, must be in [%d, %d]", dimension, - kMinRabitqDimSize, kMaxRabitqDimSize); - return IndexError_InvalidArgument; - } - entity_.update_rabitq_params_and_vector_size(dimension); - - entity_.set_ef_construction(ef_construction_); - entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_); - entity_.set_min_neighbor_cnt(min_neighbor_cnt_); - entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_); - entity_.set_scaling_factor(scaling_factor_); - entity_.set_memory_quota(memory_quota); - entity_.set_prune_cnt(prune_cnt); - - ret = entity_.init(); - if (ret != 0) { - return ret; - } - - alg_ = HnswRabitqAlgorithm::UPointer(new HnswRabitqAlgorithm(entity_)); - - ret = alg_->init(); - if (ret != 0) { - return ret; - } - - // Create and initialize RaBitQ converter - converter_ = std::make_shared(); - - IndexMeta converter_meta = meta_; - converter_meta.set_dimension(dimension); - ret = converter_->init(converter_meta, params); - if (ret != 0) { - LOG_ERROR("Failed to initialize RabitqConverter: %d", ret); - return ret; - } - - state_ = BUILD_STATE_INITED; - LOG_INFO( - "End HnswRabitqBuilder::init, params: rawVectorSize=%u vectorSize=%zu " - "efConstruction=%u " - "l0NeighborCnt=%u upperNeighborCnt=%u scalingFactor=%u " - "memoryQuota=%zu neighborPruneCnt=%zu metricName=%s ", - meta_.element_size(), entity_.vector_size(), ef_construction_, - l0_max_neighbor_cnt_, upper_max_neighbor_cnt_, scaling_factor_, - memory_quota, prune_cnt, meta_.metric_name().c_str()); - - return 0; -} - -int HnswRabitqBuilder::cleanup(void) { - LOG_INFO("Begin HnswRabitqBuilder::cleanup"); - - l0_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultL0MaxNeighborCnt; - min_neighbor_cnt_ = 0; - upper_max_neighbor_cnt_ = HnswRabitqEntity::kDefaultUpperMaxNeighborCnt; - ef_construction_ = HnswRabitqEntity::kDefaultEfConstruction; - scaling_factor_ = HnswRabitqEntity::kDefaultScalingFactor; - check_interval_secs_ = kDefaultLogIntervalSecs; - errcode_ = 0; - error_ = false; - entity_.cleanup(); - if (alg_) { - alg_->cleanup(); - } - meta_.clear(); - metric_.reset(); - stats_.clear_attributes(); - stats_.set_trained_count(0UL); - stats_.set_built_count(0UL); - stats_.set_dumped_count(0UL); - stats_.set_discarded_count(0UL); - stats_.set_trained_costtime(0UL); - stats_.set_built_costtime(0UL); - stats_.set_dumped_costtime(0UL); - state_ = BUILD_STATE_INIT; - - LOG_INFO("End HnswRabitqBuilder::cleanup"); - - return 0; -} - -int HnswRabitqBuilder::train(IndexThreads::Pointer, - IndexHolder::Pointer holder) { - if (state_ != BUILD_STATE_INITED) { - LOG_ERROR("Init the builder before HnswRabitqBuilder::train"); - return IndexError_NoReady; - } - - if (!holder) { - LOG_ERROR("Input holder is nullptr while training index"); - return IndexError_InvalidArgument; - } - if (!holder->is_matched(meta_)) { - LOG_ERROR("Input holder doesn't match index meta while training index"); - return IndexError_Mismatch; - } - LOG_INFO("Begin HnswRabitqBuilder::train"); - size_t trained_cost_time = 0; - size_t trained_count = 0; - - int ret = train_converter_and_load_reformer(holder); - if (ret != 0) { - return ret; - } - - if (metric_->support_train()) { - auto start_time = ailego::Monotime::MilliSeconds(); - auto iter = holder->create_iterator(); - if (!iter) { - LOG_ERROR("Create iterator for holder failed"); - return IndexError_Runtime; - } - while (iter->is_valid()) { - ret = metric_->train(iter->data(), meta_.dimension()); - if (ailego_unlikely(ret != 0)) { - LOG_ERROR("Hnsw build measure train failed, ret=%d", ret); - return ret; - } - iter->next(); - ++trained_count; - } - trained_cost_time = ailego::Monotime::MilliSeconds() - start_time; - } - stats_.set_trained_count(trained_count); - stats_.set_trained_costtime(trained_cost_time); - state_ = BUILD_STATE_TRAINED; - - LOG_INFO("End HnswRabitqBuilder::train"); - - return 0; -} - -int HnswRabitqBuilder::train_converter_and_load_reformer( - IndexHolder::Pointer holder) { - // Train converter (KMeans clustering) - int ret = converter_->train(holder); - if (ret != 0) { - LOG_ERROR("Failed to train RabitqConverter: %d", ret); - return ret; - } - auto memory_dumper = IndexFactory::CreateDumper("MemoryDumper"); - memory_dumper->init(ailego::Params()); - std::string file_id = ailego::StringHelper::Concat( - "rabitq_converter_", ailego::Monotime::MilliSeconds(), rand()); - ret = memory_dumper->create(file_id); - if (ret != 0) { - LOG_ERROR("Failed to create memory dumper: %d", ret); - return ret; - } - // Release memory - AILEGO_DEFER([&file_id]() { IndexMemory::Instance()->remove(file_id); }); - ret = converter_->dump(memory_dumper); - if (ret != 0) { - LOG_ERROR("Failed to dump RabitqConverter: %d", ret); - return ret; - } - ret = memory_dumper->close(); - if (ret != 0) { - LOG_ERROR("Failed to close memory dumper: %d", ret); - return ret; - } - - reformer_ = std::make_shared(); - ailego::Params reformer_params; - reformer_params.set(PARAM_RABITQ_METRIC_NAME, meta_.metric_name()); - ret = reformer_->init(reformer_params); - if (ret != 0) { - LOG_ERROR("Failed to initialize RabitqReformer: %d", ret); - return ret; - } - auto memory_storage = IndexFactory::CreateStorage("MemoryReadStorage"); - ret = memory_storage->open(file_id, false); - if (ret != 0) { - LOG_ERROR("Failed to open memory storage: %d", ret); - return ret; - } - ret = reformer_->load(memory_storage); - if (ret != 0) { - LOG_ERROR("Failed to load RabitqReformer: %d", ret); - return ret; - } - return 0; -} - -int HnswRabitqBuilder::train(const IndexTrainer::Pointer & /*trainer*/) { - if (state_ != BUILD_STATE_INITED) { - LOG_ERROR("Init the builder before HnswRabitqBuilder::train"); - return IndexError_NoReady; - } - - LOG_INFO("Begin HnswRabitqBuilder::train by trainer"); - - stats_.set_trained_count(0UL); - stats_.set_trained_costtime(0UL); - state_ = BUILD_STATE_TRAINED; - - LOG_INFO("End HnswRabitqBuilder::train by trainer"); - - return 0; -} - -int HnswRabitqBuilder::build(IndexThreads::Pointer threads, - IndexHolder::Pointer holder) { - if (state_ != BUILD_STATE_TRAINED) { - LOG_ERROR("Train the index before HnswRabitqBuilder::build"); - return IndexError_NoReady; - } - - if (!holder) { - LOG_ERROR("Input holder is nullptr while building index"); - return IndexError_InvalidArgument; - } - if (!holder->is_matched(meta_)) { - LOG_ERROR("Input holder doesn't match index meta while building index"); - return IndexError_Mismatch; - } - IndexProvider::Pointer provider = - std::dynamic_pointer_cast(holder); - if (!provider) { - LOG_ERROR("Rabitq builder expect IndexProvider"); - return IndexError_InvalidArgument; - } - - if (!threads) { - threads = std::make_shared(thread_cnt_, false); - } - - auto start_time = ailego::Monotime::MilliSeconds(); - LOG_INFO("Begin HnswRabitqBuilder::build"); - - if (holder->count() != static_cast(-1)) { - LOG_DEBUG("HnswRabitqBuilder holder documents count %lu", holder->count()); - int ret = entity_.reserve_space(holder->count()); - if (ret != 0) { - LOG_ERROR("HnswBuilde reserver space failed"); - return ret; - } - } - auto iter = holder->create_iterator(); - if (!iter) { - LOG_ERROR("Create iterator for holder failed"); - return IndexError_Runtime; - } - int ret; - error_ = false; - IndexQueryMeta ometa; - ometa.set_meta(holder->data_type(), holder->dimension()); - while (iter->is_valid()) { - const void *vec = iter->data(); - // quantize vector - std::string converted_vector; - IndexQueryMeta converted_meta; - ret = reformer_->convert(vec, ometa, &converted_vector, &converted_meta); - if (ret != 0) { - LOG_ERROR("Rabitq hnsw convert failed, ret=%d", ret); - return ret; - } - - - level_t level = alg_->get_random_level(); - node_id_t id; - - if (converted_vector.size() != entity_.vector_size()) { - LOG_ERROR( - "Converted vector size %zu is not equal to entity vector size %zu", - converted_vector.size(), entity_.vector_size()); - return IndexError_InvalidArgument; - } - ret = entity_.add_vector(level, iter->key(), converted_vector.data(), &id); - if (ailego_unlikely(ret != 0)) { - return ret; - } - iter->next(); - } - - LOG_INFO("Finished save vector, start build graph..."); - - auto task_group = threads->make_group(); - if (!task_group) { - LOG_ERROR("Failed to create task group"); - return IndexError_Runtime; - } - - std::atomic finished{0}; - for (size_t i = 0; i < threads->count(); ++i) { - task_group->submit(ailego::Closure ::New(this, &HnswRabitqBuilder::do_build, - i, threads->count(), provider, - &finished)); - } - - while (!task_group->is_finished()) { - std::unique_lock lk(mutex_); - cond_.wait_until(lk, std::chrono::system_clock::now() + - std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to build index while waiting finish"); - return errcode_; - } - LOG_INFO("Built cnt %zu, finished percent %.3f%%", - static_cast(finished.load()), - finished.load() * 100.0f / entity_.doc_cnt()); - } - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to build index while waiting finish"); - return errcode_; - } - task_group->wait_finish(); - - stats_.set_built_count(finished.load()); - stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time); - - state_ = BUILD_STATE_BUILT; - LOG_INFO("End HnswRabitqBuilder::build with RaBitQ quantization"); - return 0; -} - -void HnswRabitqBuilder::do_build(node_id_t idx, size_t step_size, - IndexProvider::Pointer provider, - std::atomic *finished) { - AILEGO_DEFER([&]() { - std::lock_guard latch(mutex_); - cond_.notify_one(); - }); - HnswRabitqContext *ctx = new (std::nothrow) HnswRabitqContext( - meta_.dimension(), metric_, - std::shared_ptr(&entity_, [](HnswRabitqEntity *) {})); - if (ailego_unlikely(ctx == nullptr)) { - if (!error_.exchange(true)) { - LOG_ERROR("Failed to create context"); - errcode_ = IndexError_NoMemory; - } - return; - } - HnswRabitqContext::Pointer auto_ptr(ctx); - ctx->set_provider(std::move(provider)); - ctx->set_max_scan_num(entity_.doc_cnt()); - int ret = ctx->init(HnswRabitqContext::kBuilderContext); - if (ret != 0) { - if (!error_.exchange(true)) { - LOG_ERROR("Failed to init context"); - errcode_ = IndexError_Runtime; - } - return; - } - - for (node_id_t id = idx; id < entity_.doc_cnt(); id += step_size) { - ctx->reset_query(ctx->dist_calculator().get_vector(id)); - ret = alg_->add_node(id, entity_.get_level(id), ctx); - if (ailego_unlikely(ret != 0)) { - if (!error_.exchange(true)) { - LOG_ERROR("Hnsw graph add node failed"); - errcode_ = ret; - } - return; - } - ctx->clear(); - (*finished)++; - } -} - -int HnswRabitqBuilder::dump(const IndexDumper::Pointer &dumper) { - if (state_ != BUILD_STATE_BUILT) { - LOG_INFO("Build the index before HnswRabitqBuilder::dump"); - return IndexError_NoReady; - } - - LOG_INFO("Begin HnswRabitqBuilder::dump"); - - meta_.set_searcher("HnswRabitqSearcher", HnswRabitqEntity::kRevision, - ailego::Params()); - auto start_time = ailego::Monotime::MilliSeconds(); - - int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); - if (ret != 0) { - LOG_ERROR("Failed to serialize meta into dumper."); - return ret; - } - - // Dump RaBitQ centroids first - if (converter_) { - ret = converter_->dump(dumper); - if (ret != 0) { - LOG_ERROR("Failed to dump RabitqConverter: %d", ret); - return ret; - } - LOG_INFO("RaBitQ centroids dumped: %zu bytes, cost %zu ms", - converter_->stats().dumped_size(), - static_cast(converter_->stats().dumped_costtime())); - } - - ret = entity_.dump(dumper); - if (ret != 0) { - LOG_ERROR("HnswRabitqBuilder dump index failed"); - return ret; - } - - stats_.set_dumped_count(entity_.doc_cnt()); - stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time); - - LOG_INFO("End HnswRabitqBuilder::dump"); - return 0; -} - - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.h deleted file mode 100644 index ffb5abe..0000000 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.h +++ /dev/null @@ -1,101 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -#pragma once - -#include -#include "zvec/core/framework/index_builder.h" -#include "zvec/core/framework/index_converter.h" -#include "zvec/core/framework/index_reformer.h" -#include "hnsw_rabitq_algorithm.h" -#include "hnsw_rabitq_builder_entity.h" - -namespace zvec { -namespace core { - -class HnswRabitqBuilder : public IndexBuilder { - public: - //! Constructor - HnswRabitqBuilder(); - - //! Initialize the builder - virtual int init(const IndexMeta &meta, - const ailego::Params ¶ms) override; - - //! Cleanup the builder - virtual int cleanup(void) override; - - //! Train the data - virtual int train(IndexThreads::Pointer, - IndexHolder::Pointer holder) override; - - //! Train the data - virtual int train(const IndexTrainer::Pointer &trainer) override; - - - //! Build the index - virtual int build(IndexThreads::Pointer threads, - IndexHolder::Pointer holder) override; - - //! Dump index into storage - virtual int dump(const IndexDumper::Pointer &dumper) override; - - //! Retrieve statistics - virtual const Stats &stats(void) const override { - return stats_; - } - - private: - void do_build(node_id_t idx, size_t step_size, - IndexProvider::Pointer provider, - std::atomic *finished); - - int train_converter_and_load_reformer(IndexHolder::Pointer holder); - - constexpr static uint32_t kDefaultLogIntervalSecs = 15U; - constexpr static uint32_t kMaxNeighborCnt = 65535; - - private: - enum BUILD_STATE { - BUILD_STATE_INIT = 0, - BUILD_STATE_INITED = 1, - BUILD_STATE_TRAINED = 2, - BUILD_STATE_BUILT = 3 - }; - - HnswRabitqBuilderEntity entity_{}; - HnswRabitqAlgorithm::UPointer alg_; // impl graph algorithm - uint32_t thread_cnt_{0}; - uint32_t min_neighbor_cnt_{0}; - uint32_t upper_max_neighbor_cnt_{ - HnswRabitqEntity::kDefaultUpperMaxNeighborCnt}; - uint32_t l0_max_neighbor_cnt_{HnswRabitqEntity::kDefaultL0MaxNeighborCnt}; - uint32_t ef_construction_{HnswRabitqEntity::kDefaultEfConstruction}; - uint32_t scaling_factor_{HnswRabitqEntity::kDefaultScalingFactor}; - uint32_t check_interval_secs_{kDefaultLogIntervalSecs}; - - int errcode_{0}; - std::atomic_bool error_{false}; - IndexMeta meta_{}; - IndexMetric::Pointer metric_{}; - IndexConverter::Pointer converter_{}; // RaBitQ converter - IndexReformer::Pointer reformer_{}; // RaBitQ reformer - std::mutex mutex_{}; - std::condition_variable cond_{}; - Stats stats_{}; - - BUILD_STATE state_{BUILD_STATE_INIT}; -}; - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.cc deleted file mode 100644 index 12ef7bd..0000000 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.cc +++ /dev/null @@ -1,199 +0,0 @@ -// 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 "hnsw_rabitq_builder_entity.h" -#include -#include -#include "utility/sparse_utility.h" - -namespace zvec { -namespace core { - -HnswRabitqBuilderEntity::HnswRabitqBuilderEntity() { - update_ep_and_level(kInvalidNodeId, 0U); -} - -int HnswRabitqBuilderEntity::cleanup() { - memory_quota_ = 0UL; - neighbors_size_ = 0U; - upper_neighbors_size_ = 0U; - padding_size_ = 0U; - vectors_buffer_.clear(); - keys_buffer_.clear(); - neighbors_buffer_.clear(); - upper_neighbors_buffer_.clear(); - neighbors_index_.clear(); - - vectors_buffer_.shrink_to_fit(); - keys_buffer_.shrink_to_fit(); - neighbors_buffer_.shrink_to_fit(); - upper_neighbors_buffer_.shrink_to_fit(); - neighbors_index_.shrink_to_fit(); - - this->HnswRabitqEntity::cleanup(); - - return 0; -} - -int HnswRabitqBuilderEntity::init() { - size_t size = vector_size(); - - //! aligned size to 32 - set_node_size(AlignSize(size)); - //! if node size is aligned to 1k, the build performance will downgrade - if (node_size() % 1024 == 0) { - set_node_size(AlignSize(node_size() + 1)); - } - - padding_size_ = node_size() - size; - - neighbors_size_ = neighbors_size(); - upper_neighbors_size_ = upper_neighbors_size(); - - return 0; -} - -int HnswRabitqBuilderEntity::reserve_space(size_t docs) { - if (memory_quota_ > 0 && (node_size() * docs + neighbors_size_ * docs + - sizeof(NeighborIndex) * docs > - memory_quota_)) { - return IndexError_NoMemory; - } - - vectors_buffer_.reserve(node_size() * docs); - keys_buffer_.reserve(sizeof(key_t) * docs); - neighbors_buffer_.reserve(neighbors_size_ * docs); - neighbors_index_.reserve(docs); - - return 0; -} - -int HnswRabitqBuilderEntity::add_vector(level_t level, key_t key, - const void *vec, node_id_t *id) { - if (memory_quota_ > 0 && - (vectors_buffer_.capacity() + keys_buffer_.capacity() + - neighbors_buffer_.capacity() + upper_neighbors_buffer_.capacity() + - neighbors_index_.capacity() * sizeof(NeighborIndex)) > memory_quota_) { - LOG_ERROR("Add vector failed, used memory exceed quota, cur_doc=%zu", - static_cast(doc_cnt())); - return IndexError_NoMemory; - } - - vectors_buffer_.append(reinterpret_cast(vec), vector_size()); - vectors_buffer_.append(padding_size_, '\0'); - keys_buffer_.append(reinterpret_cast(&key), sizeof(key)); - - // init level 0 neighbors - neighbors_buffer_.append(neighbors_size_, '\0'); - - neighbors_index_.emplace_back(upper_neighbors_buffer_.size(), level); - - // init upper layer neighbors - for (level_t cur_level = 1; cur_level <= level; ++cur_level) { - upper_neighbors_buffer_.append(upper_neighbors_size_, '\0'); - } - - *id = (*mutable_doc_cnt())++; - - return 0; -} - -key_t HnswRabitqBuilderEntity::get_key(node_id_t id) const { - return *(reinterpret_cast(keys_buffer_.data() + - id * sizeof(key_t))); -} - -const void *HnswRabitqBuilderEntity::get_vector(node_id_t id) const { - return vectors_buffer_.data() + id * node_size(); -} - -int HnswRabitqBuilderEntity::get_vector( - const node_id_t id, IndexStorage::MemoryBlock &block) const { - const void *vec = get_vector(id); - block.reset((void *)vec); - return 0; -} - -int HnswRabitqBuilderEntity::get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const { - for (uint32_t i = 0; i < count; ++i) { - vecs[i] = vectors_buffer_.data() + ids[i] * node_size(); - } - - return 0; -} - -int HnswRabitqBuilderEntity::get_vector( - const node_id_t *ids, uint32_t count, - std::vector &vec_blocks) const { - std::vector vecs(count); - get_vector(ids, count, vecs.data()); - for (uint32_t i = 0; i < count; ++i) { - vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i])); - } - return 0; -} - -const Neighbors HnswRabitqBuilderEntity::get_neighbors(level_t level, - node_id_t id) const { - const NeighborsHeader *hd = get_neighbor_header(level, id); - return {hd->neighbor_cnt, hd->neighbors}; -} - -int HnswRabitqBuilderEntity::update_neighbors( - level_t level, node_id_t id, - const std::vector> &neighbors) { - NeighborsHeader *hd = - const_cast(get_neighbor_header(level, id)); - for (size_t i = 0; i < neighbors.size(); ++i) { - hd->neighbors[i] = neighbors[i].first; - } - hd->neighbor_cnt = neighbors.size(); - - // std::cout << "id: " << id << ", neighbour, id: "; - // for (size_t i = 0; i < neighbors.size(); ++i) { - // if (i == neighbors.size()-1) - // std::cout << neighbors[i].first << ", score:" << neighbors[i].second << - // std::endl; - // else - // std::cout << neighbors[i].first << ", score:" << neighbors[i].second << - // ", id: "; - // } - - return 0; -} - -void HnswRabitqBuilderEntity::add_neighbor(level_t level, node_id_t id, - uint32_t /*size*/, - node_id_t neighbor_id) { - NeighborsHeader *hd = - const_cast(get_neighbor_header(level, id)); - hd->neighbors[hd->neighbor_cnt++] = neighbor_id; - - return; -} - -int HnswRabitqBuilderEntity::dump(const IndexDumper::Pointer &dumper) { - key_t *keys = - reinterpret_cast(const_cast(keys_buffer_.data())); - auto ret = - dump_segments(dumper, keys, [&](node_id_t id) { return get_level(id); }); - if (ailego_unlikely(ret < 0)) { - return ret; - } - - return 0; -} - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.h deleted file mode 100644 index 7460674..0000000 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.h +++ /dev/null @@ -1,139 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -#pragma once - -#include -#include "hnsw_rabitq_entity.h" - -namespace zvec { -namespace core { - -class HnswRabitqBuilderEntity : public HnswRabitqEntity { - public: - //! Add vector and key to hnsw entity, and local id will be saved to id - virtual int add_vector(level_t level, key_t key, const void *vec, - node_id_t *id) override; - - //! Get primary key of the node id - virtual key_t get_key(node_id_t id) const override; - - //! Get vector feature data by key - virtual const void *get_vector(node_id_t id) const override; - - //! Batch get vectors feature data by keys - virtual int get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const override; - - virtual int get_vector(const node_id_t id, - IndexStorage::MemoryBlock &block) const override; - virtual int get_vector( - const node_id_t *ids, uint32_t count, - std::vector &vec_blocks) const override; - - //! Get the node id's neighbors on graph level - const NeighborsHeader *get_neighbor_header(level_t level, - node_id_t id) const { - if (level == 0) { - return reinterpret_cast( - neighbors_buffer_.data() + neighbors_size_ * id); - } else { - size_t offset = neighbors_index_[id].offset; - return reinterpret_cast( - upper_neighbors_buffer_.data() + offset + - (level - 1) * upper_neighbors_size_); - } - } - - //! Get the node id's neighbors on graph level - virtual const Neighbors get_neighbors(level_t level, - node_id_t id) const override; - - //! Replace node id in level's neighbors - virtual int update_neighbors( - level_t level, node_id_t id, - const std::vector> &neighbors) - override; - - //! add a neighbor to id in graph level - virtual void add_neighbor(level_t level, node_id_t id, uint32_t size, - node_id_t neighbor_id) override; - - //! Dump the hnsw graph to dumper - virtual int dump(const IndexDumper::Pointer &dumper) override; - - //! Cleanup the entity - virtual int cleanup(void) override; - - public: - //! Constructor - HnswRabitqBuilderEntity(); - - //! Get the node graph level by id - level_t get_level(node_id_t id) const { - return neighbors_index_[id].level; - } - - //! Init builerEntity - int init(); - - //! reserve buffer space for documents - //! @param docs number of documents - int reserve_space(size_t docs); - - //! Set memory quota params - inline void set_memory_quota(size_t memory_quota) { - memory_quota_ = memory_quota; - } - - //! Get neighbors size - inline size_t neighbors_size() const { - return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); - } - - //! Get upper neighbors size - inline size_t upper_neighbors_size() const { - return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t); - } - - public: - HnswRabitqBuilderEntity(const HnswRabitqBuilderEntity &) = delete; - HnswRabitqBuilderEntity &operator=(const HnswRabitqBuilderEntity &) = delete; - - private: - friend class HnswRabitqSearcherEntity; - //! class internal used only - struct NeighborIndex { - NeighborIndex(size_t off, level_t l) : offset(off), level(l) {} - uint64_t offset : 48; - uint64_t level : 16; - }; - - std::string vectors_buffer_{}; // aligned vectors - std::string keys_buffer_{}; // aligned vectors - std::string neighbors_buffer_{}; // level 0 neighbors buffer - std::string upper_neighbors_buffer_{}; // upper layer neighbors buffer - - std::string sparse_data_buffer_{}; // aligned spase data buffer - size_t sparse_data_offset_{0}; // - - // upper layer offset + level in upper_neighbors_buffer_ - std::vector neighbors_index_{}; - size_t memory_quota_{0UL}; - size_t neighbors_size_{0U}; // level 0 neighbors size - size_t upper_neighbors_size_{0U}; // level 0 neighbors size - size_t padding_size_{}; // padding size for each vector element -}; - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc index e1fbbbb..4f48668 100644 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc @@ -11,8 +11,6 @@ // 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 "hnsw_rabitq_builder.h" -#include "hnsw_rabitq_searcher.h" #include "hnsw_rabitq_streamer.h" #include "rabitq_converter.h" #include "rabitq_reformer.h" @@ -21,8 +19,6 @@ namespace zvec::core { INDEX_FACTORY_REGISTER_STREAMER(HnswRabitqStreamer); INDEX_FACTORY_REGISTER_REFORMER_ALIAS(RabitqReformer, RabitqReformer); -INDEX_FACTORY_REGISTER_SEARCHER(HnswRabitqSearcher); INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(RabitqConverter, RabitqConverter); -INDEX_FACTORY_REGISTER_BUILDER(HnswRabitqBuilder); } // namespace zvec::core \ No newline at end of file diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.cc deleted file mode 100644 index ae9f628..0000000 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.cc +++ /dev/null @@ -1,516 +0,0 @@ -// 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 "hnsw_rabitq_searcher.h" -#include -#include "hnsw_rabitq_algorithm.h" -#include "hnsw_rabitq_entity.h" -#include "hnsw_rabitq_index_provider.h" -#include "hnsw_rabitq_params.h" -#include "hnsw_rabitq_query_entity.h" -#include "hnsw_rabitq_searcher_entity.h" -#include "rabitq_params.h" - -namespace zvec { -namespace core { - -HnswRabitqSearcher::HnswRabitqSearcher() {} - -HnswRabitqSearcher::~HnswRabitqSearcher() {} - -int HnswRabitqSearcher::init(const ailego::Params &search_params) { - params_ = search_params; - params_.get(PARAM_HNSW_RABITQ_SEARCHER_EF, &ef_); - params_.get(PARAM_HNSW_RABITQ_SEARCHER_MAX_SCAN_RATIO, &max_scan_ratio_); - params_.get(PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_ENABLE, - &bf_enabled_); - params_.get(PARAM_HNSW_RABITQ_SEARCHER_CHECK_CRC_ENABLE, &check_crc_enabled_); - params_.get(PARAM_HNSW_RABITQ_SEARCHER_NEIGHBORS_IN_MEMORY_ENABLE, - &neighbors_in_memory_enabled_); - params_.get(PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB, - &bf_negative_probability_); - params_.get(PARAM_HNSW_RABITQ_SEARCHER_BRUTE_FORCE_THRESHOLD, - &bruteforce_threshold_); - params_.get(PARAM_HNSW_RABITQ_SEARCHER_FORCE_PADDING_RESULT_ENABLE, - &force_padding_topk_enabled_); - - if (ef_ == 0) { - ef_ = HnswRabitqEntity::kDefaultEf; - } - if (bf_negative_probability_ <= 0.0f || bf_negative_probability_ >= 1.0f) { - LOG_ERROR( - "[%s] must be in range (0,1)", - PARAM_HNSW_RABITQ_SEARCHER_VISIT_BLOOMFILTER_NEGATIVE_PROB.c_str()); - return IndexError_InvalidArgument; - } - - entity_.set_neighbors_in_memory(neighbors_in_memory_enabled_); - - ailego::Params reformer_params; - reformer_params.set(PARAM_RABITQ_METRIC_NAME, meta_.metric_name()); - int ret = reformer_.init(reformer_params); - if (ret != 0) { - LOG_ERROR("Failed to initialize RabitqReformer: %d", ret); - return ret; - } - - state_ = STATE_INITED; - - LOG_DEBUG( - "Init params: ef=%u maxScanRatio=%f bfEnabled=%u checkCrcEnabled=%u " - "neighborsInMemoryEnabled=%u bfNagtiveProb=%f bruteForceThreshold=%u " - "forcePadding=%u", - ef_, max_scan_ratio_, bf_enabled_, check_crc_enabled_, - neighbors_in_memory_enabled_, bf_negative_probability_, - bruteforce_threshold_, force_padding_topk_enabled_); - - return 0; -} - -void HnswRabitqSearcher::print_debug_info() { - for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { - Neighbors neighbours = entity_.get_neighbors(0, id); - std::cout << "node: " << id << "; "; - for (uint32_t i = 0; i < neighbours.size(); ++i) { - std::cout << neighbours[i]; - - if (i == neighbours.size() - 1) { - std::cout << std::endl; - } else { - std::cout << ", "; - } - } - } -} - -int HnswRabitqSearcher::cleanup() { - LOG_INFO("Begin HnswRabitqSearcher:cleanup"); - - metric_.reset(); - meta_.clear(); - stats_.clear_attributes(); - stats_.set_loaded_count(0UL); - stats_.set_loaded_costtime(0UL); - max_scan_ratio_ = HnswRabitqEntity::kDefaultScanRatio; - max_scan_num_ = 0U; - ef_ = HnswRabitqEntity::kDefaultEf; - bf_enabled_ = false; - bf_negative_probability_ = HnswRabitqEntity::kDefaultBFNegativeProbability; - bruteforce_threshold_ = HnswRabitqEntity::kDefaultBruteForceThreshold; - check_crc_enabled_ = false; - neighbors_in_memory_enabled_ = false; - entity_.cleanup(); - state_ = STATE_INIT; - - LOG_INFO("End HnswRabitqSearcher:cleanup"); - - return 0; -} - -int HnswRabitqSearcher::load(IndexStorage::Pointer container, - IndexMetric::Pointer metric) { - if (state_ != STATE_INITED) { - LOG_ERROR("Init the searcher first before load index"); - return IndexError_Runtime; - } - - LOG_INFO("Begin HnswRabitqSearcher:load"); - - auto start_time = ailego::Monotime::MilliSeconds(); - - int ret = IndexHelper::DeserializeFromStorage(container.get(), &meta_); - if (ret != 0) { - LOG_ERROR("Failed to deserialize meta from container"); - return ret; - } - - ret = reformer_.load(container); - if (ret != 0) { - LOG_ERROR("Failed to load reformer from container: %d", ret); - return ret; - } - - ret = entity_.load(container, check_crc_enabled_); - if (ret != 0) { - LOG_ERROR("HnswRabitqSearcher load index failed"); - return ret; - } - - alg_ = HnswRabitqQueryAlgorithm::UPointer(new HnswRabitqQueryAlgorithm( - entity_, reformer_.num_clusters(), reformer_.rabitq_metric_type())); - - if (metric) { - metric_ = metric; - } else { - metric_ = IndexFactory::CreateMetric(meta_.metric_name()); - if (!metric_) { - LOG_ERROR("CreateMetric failed, name: %s", meta_.metric_name().c_str()); - return IndexError_NoExist; - } - ret = metric_->init(meta_, meta_.metric_params()); - if (ret != 0) { - LOG_ERROR("IndexMetric init failed, ret=%d", ret); - return ret; - } - if (metric_->query_metric()) { - metric_ = metric_->query_metric(); - } - } - - if (!metric_->is_matched(meta_)) { - LOG_ERROR("IndexMetric not match index meta"); - return IndexError_Mismatch; - } - - max_scan_num_ = static_cast(max_scan_ratio_ * entity_.doc_cnt()); - max_scan_num_ = std::max(4096U, max_scan_num_); - - stats_.set_loaded_count(entity_.doc_cnt()); - stats_.set_loaded_costtime(ailego::Monotime::MilliSeconds() - start_time); - state_ = STATE_LOADED; - magic_ = IndexContext::GenerateMagic(); - - LOG_INFO("End HnswRabitqSearcher::load"); - - return 0; -} - -int HnswRabitqSearcher::unload() { - LOG_INFO("HnswRabitqSearcher unload index"); - - meta_.clear(); - entity_.cleanup(); - metric_.reset(); - max_scan_num_ = 0; - stats_.set_loaded_count(0UL); - stats_.set_loaded_costtime(0UL); - state_ = STATE_INITED; - - return 0; -} - -int HnswRabitqSearcher::update_context(HnswRabitqContext *ctx) const { - const HnswRabitqEntity::Pointer entity = entity_.clone(); - if (!entity) { - LOG_ERROR("Failed to clone search context entity"); - return IndexError_Runtime; - } - ctx->set_max_scan_num(max_scan_num_); - ctx->set_bruteforce_threshold(bruteforce_threshold_); - - return ctx->update_context(HnswRabitqContext::kSearcherContext, meta_, - metric_, entity, magic_); -} - -int HnswRabitqSearcher::search_impl(const void *query, - const IndexQueryMeta &qmeta, uint32_t count, - Context::Pointer &context) const { - if (ailego_unlikely(!query || !context)) { - LOG_ERROR("The context is not created by this searcher"); - return IndexError_Mismatch; - } - HnswRabitqContext *ctx = dynamic_cast(context.get()); - ailego_do_if_false(ctx) { - LOG_ERROR("Cast context to HnswRabitqContext failed"); - return IndexError_Cast; - } - - if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) { - return search_bf_impl(query, qmeta, count, context); - } - - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer - int ret = update_context(ctx); - if (ret != 0) { - return ret; - } - } - - ctx->clear(); - ctx->resize_results(count); - for (size_t q = 0; q < count; ++q) { - HnswRabitqQueryEntity entity; - int ret = reformer_.transform_to_entity(query, &entity); - if (ailego_unlikely(ret != 0)) { - LOG_ERROR("Hnsw searcher transform failed"); - return ret; - } - ctx->reset_query(query); - ret = alg_->search(&entity, ctx); - if (ailego_unlikely(ret != 0)) { - LOG_ERROR("Hnsw searcher fast search failed"); - return ret; - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - - if (ailego_unlikely(ctx->error())) { - return IndexError_Runtime; - } - - return 0; -} - -int HnswRabitqSearcher::search_bf_impl(const void *query, - const IndexQueryMeta &qmeta, - uint32_t count, - Context::Pointer &context) const { - if (ailego_unlikely(!query || !context)) { - LOG_ERROR("The context is not created by this searcher"); - return IndexError_Mismatch; - } - HnswRabitqContext *ctx = dynamic_cast(context.get()); - ailego_do_if_false(ctx) { - LOG_ERROR("Cast context to HnswRabitqContext failed"); - return IndexError_Cast; - } - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer - int ret = update_context(ctx); - if (ret != 0) { - return ret; - } - } - - ctx->clear(); - ctx->resize_results(count); - - if (ctx->group_by_search()) { - if (!ctx->group_by().is_valid()) { - LOG_ERROR("Invalid group-by function"); - return IndexError_InvalidArgument; - } - - std::function group_by = [&](node_id_t id) { - return ctx->group_by()(entity_.get_key(id)); - }; - - for (size_t q = 0; q < count; ++q) { - HnswRabitqQueryEntity entity; - int ret = reformer_.transform_to_entity(query, &entity); - if (ailego_unlikely(ret != 0)) { - LOG_ERROR("Hnsw searcher transform failed"); - return ret; - } - ctx->reset_query(query); - ctx->group_topk_heaps().clear(); - - for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { - if (entity_.get_key(id) == kInvalidKey) { - continue; - } - if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { - EstimateRecord dist; - alg_->get_full_est(id, dist, entity); - - std::string group_id = group_by(id); - - auto &topk_heap = ctx->group_topk_heaps()[group_id]; - if (topk_heap.empty()) { - topk_heap.limit(ctx->group_topk()); - } - topk_heap.emplace_back(id, dist); - } - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - } else { - for (size_t q = 0; q < count; ++q) { - HnswRabitqQueryEntity entity; - int ret = reformer_.transform_to_entity(query, &entity); - if (ailego_unlikely(ret != 0)) { - LOG_ERROR("Hnsw searcher transform failed"); - return ret; - } - ctx->reset_query(query); - ctx->topk_heap().clear(); - for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { - if (entity_.get_key(id) == kInvalidKey) { - continue; - } - if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { - EstimateRecord dist; - alg_->get_full_est(id, dist, entity); - ctx->topk_heap().emplace(id, dist); - } - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - } - - if (ailego_unlikely(ctx->error())) { - return IndexError_Runtime; - } - - return 0; -} - -int HnswRabitqSearcher::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(!query || !context)) { - LOG_ERROR("The context is not created by this searcher"); - return IndexError_Mismatch; - } - - if (ailego_unlikely(p_keys.size() != count)) { - LOG_ERROR("The size of p_keys is not equal to count"); - return IndexError_InvalidArgument; - } - - HnswRabitqContext *ctx = dynamic_cast(context.get()); - ailego_do_if_false(ctx) { - LOG_ERROR("Cast context to HnswRabitqContext failed"); - return IndexError_Cast; - } - if (ctx->magic() != magic_) { - //! context is created by another searcher or streamer - int ret = update_context(ctx); - if (ret != 0) { - return ret; - } - } - - ctx->clear(); - ctx->resize_results(count); - - if (ctx->group_by_search()) { - if (!ctx->group_by().is_valid()) { - LOG_ERROR("Invalid group-by function"); - return IndexError_InvalidArgument; - } - - std::function group_by = [&](node_id_t id) { - return ctx->group_by()(entity_.get_key(id)); - }; - - for (size_t q = 0; q < count; ++q) { - HnswRabitqQueryEntity entity; - int ret = reformer_.transform_to_entity(query, &entity); - if (ailego_unlikely(ret != 0)) { - LOG_ERROR("Hnsw searcher transform failed"); - return ret; - } - ctx->reset_query(query); - ctx->group_topk_heaps().clear(); - - for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { - uint64_t pk = p_keys[q][idx]; - if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { - node_id_t id = entity_.get_id(pk); - if (id != kInvalidNodeId) { - EstimateRecord dist; - alg_->get_full_est(id, dist, entity); - std::string group_id = group_by(id); - - auto &topk_heap = ctx->group_topk_heaps()[group_id]; - if (topk_heap.empty()) { - topk_heap.limit(ctx->group_topk()); - } - topk_heap.emplace_back(id, dist); - } - } - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - } else { - for (size_t q = 0; q < count; ++q) { - HnswRabitqQueryEntity entity; - int ret = reformer_.transform_to_entity(query, &entity); - if (ailego_unlikely(ret != 0)) { - LOG_ERROR("Hnsw searcher transform failed"); - return ret; - } - ctx->reset_query(query); - ctx->topk_heap().clear(); - for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { - uint64_t pk = p_keys[q][idx]; - if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { - node_id_t id = entity_.get_id(pk); - if (id != kInvalidNodeId) { - EstimateRecord dist; - alg_->get_full_est(id, dist, entity); - ctx->topk_heap().emplace(id, dist); - } - } - } - ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); - } - } - - if (ailego_unlikely(ctx->error())) { - return IndexError_Runtime; - } - - return 0; -} - -IndexSearcher::Context::Pointer HnswRabitqSearcher::create_context() const { - if (ailego_unlikely(state_ != STATE_LOADED)) { - LOG_ERROR("Load the index first before create context"); - return Context::Pointer(); - } - const HnswRabitqEntity::Pointer search_ctx_entity = entity_.clone(); - if (!search_ctx_entity) { - LOG_ERROR("Failed to create search context entity"); - return Context::Pointer(); - } - HnswRabitqContext *ctx = new (std::nothrow) - HnswRabitqContext(meta_.dimension(), metric_, search_ctx_entity); - if (ailego_unlikely(ctx == nullptr)) { - LOG_ERROR("Failed to new HnswRabitqContext"); - return Context::Pointer(); - } - ctx->set_ef(ef_); - ctx->set_max_scan_num(max_scan_num_); - uint32_t filter_mode = - bf_enabled_ ? VisitFilter::BloomFilter : VisitFilter::ByteMap; - ctx->set_filter_mode(filter_mode); - ctx->set_filter_negative_probability(bf_negative_probability_); - ctx->set_magic(magic_); - ctx->set_force_padding_topk(force_padding_topk_enabled_); - ctx->set_bruteforce_threshold(bruteforce_threshold_); - if (ailego_unlikely(ctx->init(HnswRabitqContext::kSearcherContext)) != 0) { - LOG_ERROR("Init HnswRabitqContext failed"); - delete ctx; - return Context::Pointer(); - } - - return Context::Pointer(ctx); -} - -IndexProvider::Pointer HnswRabitqSearcher::create_provider(void) const { - LOG_DEBUG("HnswRabitqSearcher create provider"); - - auto entity = entity_.clone(); - if (ailego_unlikely(!entity)) { - LOG_ERROR("Clone HnswRabitqEntity failed"); - return Provider::Pointer(); - } - return Provider::Pointer(new (std::nothrow) HnswRabitqIndexProvider( - meta_, entity, "HnswRabitqSearcher")); -} - -const void *HnswRabitqSearcher::get_vector(uint64_t key) const { - return entity_.get_vector_by_key(key); -} - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.h deleted file mode 100644 index 062498e..0000000 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.h +++ /dev/null @@ -1,142 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -#pragma once - -#include "zvec/core/framework/index_framework.h" -#include "hnsw_rabitq_query_algorithm.h" -#include "hnsw_rabitq_searcher_entity.h" -#include "rabitq_reformer.h" - -namespace zvec { -namespace core { - -class HnswRabitqSearcher : public IndexSearcher { - public: - using ContextPointer = IndexSearcher::Context::Pointer; - - public: - HnswRabitqSearcher(void); - ~HnswRabitqSearcher(void); - - HnswRabitqSearcher(const HnswRabitqSearcher &) = delete; - HnswRabitqSearcher &operator=(const HnswRabitqSearcher &) = delete; - - protected: - //! Initialize Searcher - virtual int init(const ailego::Params ¶ms) override; - - //! Cleanup Searcher - virtual int cleanup(void) override; - - //! Load Index from storage - virtual int load(IndexStorage::Pointer container, - IndexMetric::Pointer metric) override; - - //! Unload index from storage - virtual int unload(void) override; - - //! KNN Search - virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, - ContextPointer &context) const override { - return search_impl(query, qmeta, 1, context); - } - - //! KNN Search - virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, - uint32_t count, - ContextPointer &context) const override; - - //! Linear Search - virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, - ContextPointer &context) const override { - return search_bf_impl(query, qmeta, 1, context); - } - - //! Linear Search - virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, - uint32_t count, - ContextPointer &context) const override; - - //! Linear search by primary keys - virtual int search_bf_by_p_keys_impl( - const void *query, const std::vector> &p_keys, - const IndexQueryMeta &qmeta, ContextPointer &context) const override { - return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); - } - - //! Linear search by primary keys - virtual int search_bf_by_p_keys_impl( - const void *query, const std::vector> &p_keys, - const IndexQueryMeta &qmeta, uint32_t count, - ContextPointer &context) const override; - - //! Fetch vector by key - virtual const void *get_vector(uint64_t key) const override; - - //! Create a searcher context - virtual ContextPointer create_context() const override; - - //! Create a new iterator - virtual IndexProvider::Pointer create_provider(void) const override; - - //! Retrieve statistics - virtual const Stats &stats(void) const override { - return stats_; - } - - //! Retrieve meta of index - virtual const IndexMeta &meta(void) const override { - return meta_; - } - - //! Retrieve params of index - virtual const ailego::Params ¶ms(void) const override { - return params_; - } - - virtual void print_debug_info() override; - - private: - //! To share ctx across streamer/searcher, we need to update the context for - //! current streamer/searcher - int update_context(HnswRabitqContext *ctx) const; - - private: - enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; - - HnswRabitqSearcherEntity entity_{}; - HnswRabitqQueryAlgorithm::UPointer alg_; // impl graph algorithm - - IndexMetric::Pointer metric_{}; - IndexMeta meta_{}; - ailego::Params params_{}; - Stats stats_; - uint32_t ef_{HnswRabitqEntity::kDefaultEf}; - uint32_t max_scan_num_{0U}; - uint32_t bruteforce_threshold_{HnswRabitqEntity::kDefaultBruteForceThreshold}; - float max_scan_ratio_{HnswRabitqEntity::kDefaultScanRatio}; - bool bf_enabled_{false}; - bool check_crc_enabled_{false}; - bool neighbors_in_memory_enabled_{false}; - bool force_padding_topk_enabled_{false}; - float bf_negative_probability_{ - HnswRabitqEntity::kDefaultBFNegativeProbability}; - uint32_t magic_{0U}; - RabitqReformer reformer_; - - State state_{STATE_INIT}; -}; - -} // namespace core -} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.cc deleted file mode 100644 index 415060c..0000000 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.cc +++ /dev/null @@ -1,515 +0,0 @@ -// 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 "hnsw_rabitq_searcher_entity.h" -#include -#include "utility/sparse_utility.h" - -namespace zvec { -namespace core { - -HnswRabitqSearcherEntity::HnswRabitqSearcherEntity() {} - -int HnswRabitqSearcherEntity::cleanup(void) { - storage_.reset(); - vectors_.reset(); - keys_.reset(); - neighbors_.reset(); - neighbors_meta_.reset(); - neighbors_in_memory_enabled_ = false; - loaded_ = false; - - this->HnswRabitqEntity::cleanup(); - - return 0; -} - -key_t HnswRabitqSearcherEntity::get_key(node_id_t id) const { - const void *key; - if (ailego_unlikely(keys_->read(id * sizeof(key_t), &key, sizeof(key_t)) != - sizeof(key_t))) { - LOG_ERROR("Read key from segment failed"); - return kInvalidKey; - } - return *(reinterpret_cast(key)); -} - -//! Get vector local id by key -node_id_t HnswRabitqSearcherEntity::get_id(key_t key) const { - if (ailego_unlikely(!mapping_)) { - LOG_ERROR("Index missing mapping segment"); - return kInvalidNodeId; - } - - //! Do binary search - node_id_t start = 0UL; - node_id_t end = doc_cnt(); - const void *data; - node_id_t idx = 0u; - while (start < end) { - idx = start + (end - start) / 2; - if (ailego_unlikely( - mapping_->read(idx * sizeof(node_id_t), &data, sizeof(node_id_t)) != - sizeof(node_id_t))) { - LOG_ERROR("Read key from segment failed"); - return kInvalidNodeId; - } - const key_t *mkey; - node_id_t local_id = *reinterpret_cast(data); - if (ailego_unlikely(keys_->read(local_id * sizeof(key_t), - (const void **)(&mkey), - sizeof(key_t)) != sizeof(key_t))) { - LOG_ERROR("Read key from segment failed"); - return kInvalidNodeId; - } - if (*mkey < key) { - start = idx + 1; - } else if (*mkey > key) { - end = idx; - } else { - return local_id; - } - } - return kInvalidNodeId; -} - -const void *HnswRabitqSearcherEntity::get_vector_by_key(key_t key) const { - node_id_t local_id = get_id(key); - if (ailego_unlikely(local_id == kInvalidNodeId)) { - return nullptr; - } - - return get_vector(local_id); -} - -const void *HnswRabitqSearcherEntity::get_vector(node_id_t id) const { - size_t read_size = vector_size(); - size_t offset = node_size() * id; - - const void *vec; - if (ailego_unlikely(vectors_->read(offset, &vec, read_size) != read_size)) { - LOG_ERROR("Read vector from segment failed"); - return nullptr; - } - return vec; -} - -int HnswRabitqSearcherEntity::get_vector( - const node_id_t id, IndexStorage::MemoryBlock &block) const { - const void *vec = get_vector(id); - block.reset((void *)vec); - return 0; -} - -const void *HnswRabitqSearcherEntity::get_vectors() const { - const void *vec; - size_t len = node_size() * doc_cnt(); - if (vectors_->read(0, &vec, len) != len) { - LOG_ERROR("Read vectors from segment failed"); - return nullptr; - } - return vec; -} - -int HnswRabitqSearcherEntity::get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const { - ailego_assert_with(count <= segment_datas_.size(), "invalid count"); - - size_t read_size = vector_size(); - - for (uint32_t i = 0; i < count; ++i) { - segment_datas_[i].offset = node_size() * ids[i]; - segment_datas_[i].length = read_size; - - ailego_assert_with(segment_datas_[i].offset < vectors_->data_size(), - "invalid offset"); - } - if (ailego_unlikely(!vectors_->read(&segment_datas_[0], count))) { - LOG_ERROR("Read vectors from segment failed"); - return IndexError_ReadData; - } - for (uint32_t i = 0; i < count; ++i) { - vecs[i] = segment_datas_[i].data; - } - - return 0; -} - -int HnswRabitqSearcherEntity::get_vector( - const node_id_t *ids, uint32_t count, - std::vector &vec_blocks) const { - const void *vecs[count]; - get_vector(ids, count, vecs); - for (uint32_t i = 0; i < count; ++i) { - vec_blocks.emplace_back(IndexStorage::MemoryBlock((void *)vecs[i])); - } - return 0; -} - -const Neighbors HnswRabitqSearcherEntity::get_neighbors(level_t level, - node_id_t id) const { - if (level == 0) { - if (neighbors_in_memory_enabled_) { - auto hd = reinterpret_cast( - fixed_neighbors_.get() + neighbors_size() * id); - return {hd->neighbor_cnt, hd->neighbors}; - } - - const GraphNeighborMeta *m; - if (ailego_unlikely(neighbors_meta_->read(id * sizeof(GraphNeighborMeta), - (const void **)(&m), - sizeof(GraphNeighborMeta)) != - sizeof(GraphNeighborMeta))) { - LOG_ERROR("Read neighbors meta from segment failed"); - return {0, nullptr}; - } - - const void *data; - if (ailego_unlikely(neighbors_->read(m->offset, &data, - m->neighbor_cnt * sizeof(node_id_t)) != - m->neighbor_cnt * sizeof(node_id_t))) { - LOG_ERROR("Read neighbors from segment failed"); - return {0, nullptr}; - } - return {static_cast(m->neighbor_cnt), - reinterpret_cast(data)}; - } - - //! Read level > 0 neighbors - const HnswNeighborMeta *m; - if (ailego_unlikely(upper_neighbors_meta_->read(id * sizeof(HnswNeighborMeta), - (const void **)(&m), - sizeof(HnswNeighborMeta)) != - sizeof(HnswNeighborMeta))) { - LOG_ERROR("Read neighbors meta from segment failed"); - return {0, nullptr}; - } - - ailego_assert_with(level <= m->level, "invalid level"); - size_t offset = m->offset + (level - 1) * upper_neighbors_size(); - ailego_assert_with(offset <= upper_neighbors_->data_size(), "invalid offset"); - const void *data; - if (ailego_unlikely( - upper_neighbors_->read(offset, &data, upper_neighbors_size()) != - upper_neighbors_size())) { - LOG_ERROR("Read neighbors from segment failed"); - return {0, nullptr}; - } - - auto hd = reinterpret_cast(data); - return {hd->neighbor_cnt, hd->neighbors}; -} - -int HnswRabitqSearcherEntity::load(const IndexStorage::Pointer &container, - bool check_crc) { - storage_ = container; - - int ret = load_segments(check_crc); - if (ret != 0) { - return ret; - } - - loaded_ = true; - - LOG_INFO( - "Index info: docCnt=%u entryPoint=%u maxLevel=%d efConstruct=%zu " - "l0NeighborCnt=%zu upperNeighborCnt=%zu scalingFactor=%zu " - "vectorSize=%zu nodeSize=%zu vectorSegmentSize=%zu keySegmentSize=%zu " - "neighborsSegmentSize=%zu neighborsMetaSegmentSize=%zu ", - doc_cnt(), entry_point(), cur_max_level(), ef_construction(), - l0_neighbor_cnt(), upper_neighbor_cnt(), scaling_factor(), vector_size(), - node_size(), vectors_->data_size(), keys_->data_size(), - neighbors_ == nullptr ? 0 : neighbors_->data_size(), - neighbors_meta_ == nullptr ? 0 : neighbors_meta_->data_size()); - - return 0; -} - -int HnswRabitqSearcherEntity::load_segments(bool check_crc) { - //! load header - const void *data = nullptr; - HNSWHeader hd; - auto graph_hd_segment = storage_->get(kGraphHeaderSegmentId); - if (!graph_hd_segment || graph_hd_segment->data_size() < sizeof(hd.graph)) { - LOG_ERROR("Miss or invalid segment %s", kGraphHeaderSegmentId.c_str()); - return IndexError_InvalidFormat; - } - if (graph_hd_segment->read(0, reinterpret_cast(&data), - sizeof(hd.graph)) != sizeof(hd.graph)) { - LOG_ERROR("Read segment %s failed", kGraphHeaderSegmentId.c_str()); - return IndexError_ReadData; - } - memcpy(&hd.graph, data, sizeof(hd.graph)); - - auto hnsw_hd_segment = storage_->get(kHnswHeaderSegmentId); - if (!hnsw_hd_segment || hnsw_hd_segment->data_size() < sizeof(hd.hnsw)) { - LOG_ERROR("Miss or invalid segment %s", kHnswHeaderSegmentId.c_str()); - return IndexError_InvalidFormat; - } - if (hnsw_hd_segment->read(0, reinterpret_cast(&data), - sizeof(hd.hnsw)) != sizeof(hd.hnsw)) { - LOG_ERROR("Read segment %s failed", kHnswHeaderSegmentId.c_str()); - return IndexError_ReadData; - } - memcpy(&hd.hnsw, data, sizeof(hd.hnsw)); - *mutable_header() = hd; - segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt())); - - vectors_ = storage_->get(kGraphFeaturesSegmentId); - if (!vectors_) { - LOG_ERROR("IndexStorage get segment %s failed", - kGraphFeaturesSegmentId.c_str()); - return IndexError_InvalidFormat; - } - keys_ = storage_->get(kGraphKeysSegmentId); - if (!keys_) { - LOG_ERROR("IndexStorage get segment %s failed", - kGraphKeysSegmentId.c_str()); - return IndexError_InvalidFormat; - } - - neighbors_ = storage_->get(kGraphNeighborsSegmentId); - if (!neighbors_ || (neighbors_->data_size() == 0 && doc_cnt() > 1)) { - LOG_ERROR("IndexStorage get segment %s failed or empty", - kGraphNeighborsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - neighbors_meta_ = storage_->get(kGraphOffsetsSegmentId); - if (!neighbors_meta_ || - neighbors_meta_->data_size() < sizeof(GraphNeighborMeta) * doc_cnt()) { - LOG_ERROR("IndexStorage get segment %s failed or invalid size", - kGraphOffsetsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - upper_neighbors_ = storage_->get(kHnswNeighborsSegmentId); - if (!upper_neighbors_ || - (upper_neighbors_->data_size() == 0 && cur_max_level() > 0)) { - LOG_ERROR("IndexStorage get segment %s failed or empty", - kHnswNeighborsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - upper_neighbors_meta_ = storage_->get(kHnswOffsetsSegmentId); - if (!upper_neighbors_meta_ || upper_neighbors_meta_->data_size() < - sizeof(HnswNeighborMeta) * doc_cnt()) { - LOG_ERROR("IndexStorage get segment %s failed or invalid size", - kHnswOffsetsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - mapping_ = storage_->get(kGraphMappingSegmentId); - if (!mapping_ || mapping_->data_size() < sizeof(node_id_t) * doc_cnt()) { - LOG_ERROR("IndexStorage get segment %s failed or invalid size", - kGraphMappingSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - if (check_crc) { - std::vector segments; - segments.emplace_back(graph_hd_segment); - segments.emplace_back(hnsw_hd_segment); - segments.emplace_back(vectors_); - segments.emplace_back(keys_); - - segments.emplace_back(neighbors_); - segments.emplace_back(neighbors_meta_); - segments.emplace_back(upper_neighbors_); - segments.emplace_back(upper_neighbors_meta_); - - if (!do_crc_check(segments)) { - LOG_ERROR("Check index crc failed, the index may broken"); - return IndexError_Runtime; - } - } - - if (neighbors_in_memory_enabled_) { - int ret = load_and_flat_neighbors(); - if (ret != 0) { - return ret; - } - } - - return 0; -} - -int HnswRabitqSearcherEntity::load_and_flat_neighbors() { - fixed_neighbors_.reset( - new (std::nothrow) char[neighbors_size() * doc_cnt()]{}, - std::default_delete()); - if (!fixed_neighbors_) { - LOG_ERROR("Malloc memory failed"); - return IndexError_NoMemory; - } - - //! Get a new segemnt to release the buffer after loading neighbors - auto neighbors_meta = storage_->get(kGraphOffsetsSegmentId); - if (!neighbors_meta) { - LOG_ERROR("IndexStorage get segment graph.offsets failed"); - return IndexError_InvalidArgument; - } - - const GraphNeighborMeta *neighbors_index = nullptr; - if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index), - neighbors_meta->data_size()) != - neighbors_meta->data_size()) { - LOG_ERROR("Read segment %s data failed", kGraphOffsetsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - const char *neighbor_data; - for (node_id_t id = 0; id < doc_cnt(); ++id) { - size_t rd_size = neighbors_index[id].neighbor_cnt * sizeof(node_id_t); - if (ailego_unlikely( - neighbors_->read(neighbors_index[id].offset, - reinterpret_cast(&neighbor_data), - rd_size) != rd_size)) { - LOG_ERROR("Read neighbors from segment failed"); - return IndexError_ReadData; - } - // copy level 0 neighbors to fixed size neighbors memory - char *dst = fixed_neighbors_.get() + neighbors_size() * id; - *reinterpret_cast(dst) = neighbors_index[id].neighbor_cnt; - memcpy(dst + sizeof(uint32_t), neighbor_data, rd_size); - } - - return 0; -} - -int HnswRabitqSearcherEntity::get_fixed_neighbors( - std::vector *fixed_neighbors) const { - //! Get a new segemnt to release the buffer after loading neighbors - auto neighbors_meta = storage_->get(kGraphOffsetsSegmentId); - if (!neighbors_meta) { - LOG_ERROR("IndexStorage get segment graph.offsets failed"); - return IndexError_InvalidArgument; - } - - const GraphNeighborMeta *neighbors_index = nullptr; - size_t meta_size = neighbors_meta->data_size(); - if (neighbors_meta->read(0, reinterpret_cast(&neighbors_index), - meta_size) != meta_size) { - LOG_ERROR("Read segment %s data failed", kGraphOffsetsSegmentId.c_str()); - return IndexError_InvalidArgument; - } - - size_t fixed_neighbor_cnt = l0_neighbor_cnt(); - fixed_neighbors->resize((fixed_neighbor_cnt + 1) * doc_cnt(), kInvalidNodeId); - - size_t neighbors_cnt_offset = fixed_neighbor_cnt * doc_cnt(); - size_t total_neighbor_cnt = 0; - for (node_id_t id = 0; id < doc_cnt(); ++id) { - size_t cur_neighbor_cnt = neighbors_index[id].neighbor_cnt; - if (cur_neighbor_cnt == 0) { - (*fixed_neighbors)[neighbors_cnt_offset + id] = 0; - continue; - } - size_t rd_size = cur_neighbor_cnt * sizeof(node_id_t); - const uint32_t *neighbors; - if (neighbors_->read(neighbors_index[id].offset, - reinterpret_cast(&neighbors), - rd_size) != rd_size) { - LOG_ERROR("Read neighbors from segment failed"); - return IndexError_ReadData; - } - - // copy level 0 neighbors to fixed size neighbors memory - auto it = fixed_neighbors->begin() + id * fixed_neighbor_cnt; - std::copy(neighbors, neighbors + cur_neighbor_cnt, it); - - (*fixed_neighbors)[neighbors_cnt_offset + id] = cur_neighbor_cnt; - total_neighbor_cnt += cur_neighbor_cnt; - } - LOG_INFO("total neighbor cnt: %zu, average neighbor cnt: %zu", - total_neighbor_cnt, total_neighbor_cnt / doc_cnt()); - - return 0; -} - -bool HnswRabitqSearcherEntity::do_crc_check( - std::vector &segments) const { - constexpr size_t blk_size = 4096; - const void *data; - for (auto &segment : segments) { - size_t offset = 0; - size_t rd_size; - uint32_t crc = 0; - while (offset < segment->data_size()) { - size_t size = std::min(blk_size, segment->data_size() - offset); - if ((rd_size = segment->read(offset, &data, size)) <= 0) { - break; - } - offset += rd_size; - crc = ailego::Crc32c::Hash(data, rd_size, crc); - } - if (crc != segment->data_crc()) { - return false; - } - } - return true; -} - -const HnswRabitqEntity::Pointer HnswRabitqSearcherEntity::clone() const { - auto vectors = vectors_->clone(); - if (ailego_unlikely(!vectors)) { - LOG_ERROR("clone segment %s failed", kGraphFeaturesSegmentId.c_str()); - return HnswRabitqEntity::Pointer(); - } - auto keys = keys_->clone(); - if (ailego_unlikely(!keys)) { - LOG_ERROR("clone segment %s failed", kGraphKeysSegmentId.c_str()); - return HnswRabitqEntity::Pointer(); - } - - auto mapping = mapping_->clone(); - if (ailego_unlikely(!mapping)) { - LOG_ERROR("clone segment %s failed", kGraphMappingSegmentId.c_str()); - return HnswRabitqEntity::Pointer(); - } - - auto neighbors = neighbors_->clone(); - if (ailego_unlikely(!neighbors)) { - LOG_ERROR("clone segment %s failed", kGraphNeighborsSegmentId.c_str()); - return HnswRabitqEntity::Pointer(); - } - auto upper_neighbors = upper_neighbors_->clone(); - if (ailego_unlikely(!neighbors)) { - LOG_ERROR("clone segment %s failed", kHnswNeighborsSegmentId.c_str()); - return HnswRabitqEntity::Pointer(); - } - auto neighbors_meta = neighbors_meta_->clone(); - if (ailego_unlikely(!neighbors_meta)) { - LOG_ERROR("clone segment %s failed", kGraphOffsetsSegmentId.c_str()); - return HnswRabitqEntity::Pointer(); - } - auto upper_neighbors_meta = upper_neighbors_meta_->clone(); - if (ailego_unlikely(!upper_neighbors_meta)) { - LOG_ERROR("clone segment %s failed", kHnswOffsetsSegmentId.c_str()); - return HnswRabitqEntity::Pointer(); - } - - SegmentGroupParam neighbor_group{neighbors, neighbors_meta, upper_neighbors, - upper_neighbors_meta}; - - HnswRabitqSearcherEntity *entity = new (std::nothrow) - HnswRabitqSearcherEntity(header(), vectors, keys, mapping, neighbor_group, - fixed_neighbors_, neighbors_in_memory_enabled_); - if (ailego_unlikely(!entity)) { - LOG_ERROR("HnswRabitqSearcherEntity new failed"); - } - - return HnswRabitqEntity::Pointer(entity); -} - -} // namespace core -} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.h b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.h deleted file mode 100644 index 164e0fd..0000000 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.h +++ /dev/null @@ -1,158 +0,0 @@ -// Copyright 2025-present the zvec project -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -#pragma once - -#include "hnsw_rabitq_builder_entity.h" -#include "hnsw_rabitq_entity.h" - -namespace zvec { -namespace core { - -class HnswRabitqSearcherEntity : public HnswRabitqEntity { - public: - using Pointer = std::shared_ptr; - using SegmentPointer = IndexStorage::Segment::Pointer; - - public: - struct SegmentGroupParam { - SegmentGroupParam(SegmentPointer neighbors_in, - SegmentPointer neighbors_meta_in, - SegmentPointer upper_neighbors_in, - SegmentPointer upper_neighbors_meta_in) - : neighbors{neighbors_in}, - neighbors_meta{neighbors_meta_in}, - upper_neighbors{upper_neighbors_in}, - upper_neighbors_meta{upper_neighbors_meta_in} {} - - SegmentPointer neighbors{nullptr}; - SegmentPointer neighbors_meta{nullptr}; - SegmentPointer upper_neighbors{nullptr}; - SegmentPointer upper_neighbors_meta{nullptr}; - }; - - //! Constructor - HnswRabitqSearcherEntity(); - - //! Make a copy of searcher entity, to support thread-safe operation. - //! The segment in container cannot be read concurrenly - virtual const HnswRabitqEntity::Pointer clone() const override; - - //! Get primary key of the node id - virtual key_t get_key(node_id_t id) const override; - - //! Get vector local id by key - node_id_t get_id(key_t key) const; - - //! Get vector feature data by key - virtual const void *get_vector_by_key(key_t key) const override; - - //! Get vector feature data by id - virtual const void *get_vector(node_id_t id) const override; - - //! Get vector feature data by id - virtual int get_vector(const node_id_t *ids, uint32_t count, - const void **vecs) const override; - - virtual int get_vector(const node_id_t id, - IndexStorage::MemoryBlock &block) const override; - virtual int get_vector( - const node_id_t *ids, uint32_t count, - std::vector &vec_blocks) const override; - - //! Get all vectors - const void *get_vectors() const; - - //! Get the node id's neighbors on graph level - virtual const Neighbors get_neighbors(level_t level, - node_id_t id) const override; - - virtual int load(const IndexStorage::Pointer &container, - bool check_crc) override; - - int load_segments(bool check_crc); - - virtual int cleanup(void) override; - - public: - bool is_loaded() const { - return loaded_; - } - - void set_neighbors_in_memory(bool enabled) { - neighbors_in_memory_enabled_ = enabled; - } - - //! get fixed length neighbors data - int get_fixed_neighbors(std::vector *fixed_neighbors) const; - - private: - //! Constructor - HnswRabitqSearcherEntity(const HNSWHeader &hd, const SegmentPointer &vectors, - const SegmentPointer &keys, - const SegmentPointer &mapping, - const SegmentGroupParam &neighbor_group, - const std::shared_ptr &fixed_neighbors, - bool neighbors_in_memory_enabled) - : HnswRabitqEntity(hd), - vectors_(vectors), - keys_(keys), - mapping_(mapping), - neighbors_(neighbor_group.neighbors), - neighbors_meta_(neighbor_group.neighbors_meta), - upper_neighbors_(neighbor_group.upper_neighbors), - upper_neighbors_meta_(neighbor_group.upper_neighbors_meta), - neighbors_in_memory_enabled_(neighbors_in_memory_enabled) { - segment_datas_.resize(std::max(l0_neighbor_cnt(), upper_neighbor_cnt()), - IndexStorage::SegmentData(0U, 0U)); - fixed_neighbors_ = fixed_neighbors; - } - - bool do_crc_check(std::vector &segments) const; - - inline size_t neighbors_size() const { - return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); - } - - inline size_t upper_neighbors_size() const { - return sizeof(NeighborsHeader) + upper_neighbor_cnt() * sizeof(node_id_t); - } - - //! If neighbors_in_memory_enabled, load the level0 neighbors to memory - int load_and_flat_neighbors(void); - - public: - HnswRabitqSearcherEntity(const HnswRabitqSearcherEntity &) = delete; - HnswRabitqSearcherEntity &operator=(const HnswRabitqSearcherEntity &) = - delete; - - private: - IndexStorage::Pointer storage_{}; - - SegmentPointer vectors_{}; - SegmentPointer keys_{}; - SegmentPointer mapping_{}; - - SegmentPointer neighbors_{}; - SegmentPointer neighbors_meta_{}; - SegmentPointer upper_neighbors_{}; - SegmentPointer upper_neighbors_meta_{}; - - mutable std::vector segment_datas_{}; - std::shared_ptr fixed_neighbors_{}; // level 0 fixed size neighbors - bool neighbors_in_memory_enabled_{false}; - bool loaded_{false}; -}; - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc index ba79e8a..9eacf0b 100644 --- a/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc +++ b/src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc @@ -415,9 +415,6 @@ int HnswRabitqStreamer::dump(const IndexDumper::Pointer &dumper) { shared_mutex_.lock(); AILEGO_DEFER([&]() { shared_mutex_.unlock(); }); - meta_.set_searcher("HnswRabitqSearcher", HnswRabitqEntity::kRevision, - ailego::Params()); - int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); if (ret != 0) { LOG_ERROR("Failed to serialize meta into dumper."); diff --git a/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_test.cc b/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_test.cc deleted file mode 100644 index ae83c79..0000000 --- a/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_test.cc +++ /dev/null @@ -1,428 +0,0 @@ -// 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 "hnsw_rabitq_builder.h" -#include -#include -#include -#include -#include -#include -#include -#include "zvec/core/framework/index_framework.h" -#include "zvec/core/framework/index_logger.h" -#include "zvec/core/framework/index_provider.h" - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wunused-result" -#endif - -using namespace std; -using namespace zvec::ailego; - -namespace zvec { -namespace core { - -constexpr size_t static dim = 128; - -class HnswRabitqBuilderTest : public testing::Test { - protected: - void SetUp(void); - void TearDown(void); - - static std::string _dir; - static shared_ptr _index_meta_ptr; -}; - -std::string HnswRabitqBuilderTest::_dir("hnswRabitqBuilderTest"); -shared_ptr HnswRabitqBuilderTest::_index_meta_ptr; - -void HnswRabitqBuilderTest::SetUp(void) { - IndexLoggerBroker::SetLevel(0); - _index_meta_ptr.reset(new (nothrow) - IndexMeta(IndexMeta::DataType::DT_FP32, dim)); - _index_meta_ptr->set_metric("SquaredEuclidean", 0, ailego::Params()); -} - -void HnswRabitqBuilderTest::TearDown(void) { - char cmdBuf[100]; - snprintf(cmdBuf, 100, "rm -rf %s", _dir.c_str()); - // system(cmdBuf); -} - -TEST_F(HnswRabitqBuilderTest, TestGeneral) { - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - 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 * dim + j) / 1000.0f; - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - - ASSERT_EQ(0, builder->train(holder)); - - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestGeneral"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats = builder->stats(); - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt, stats.built_count()); - ASSERT_EQ(doc_cnt, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); -} - -TEST_F(HnswRabitqBuilderTest, TestLoad) { - // Load index with searcher and verify search - auto searcher = IndexFactory::CreateSearcher("HnswRabitqSearcher"); - ASSERT_NE(searcher, nullptr); - - ailego::Params search_params; - search_params.set("proxima.hnsw_rabitq.searcher.ef", 100UL); - ASSERT_EQ(0, searcher->init(search_params)); - - auto loader = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_NE(loader, nullptr); - ASSERT_EQ(0, loader->init(ailego::Params())); - string path = _dir + "/TestGeneral"; - ASSERT_EQ(0, loader->open(path, false)); - - ASSERT_EQ(0, searcher->load(loader, nullptr)); - - // Perform search verification - NumericalVector query_vec(dim); - for (size_t j = 0; j < dim; ++j) { - query_vec[j] = static_cast(j) / 1000.0f; - } - - IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); - - auto context = searcher->create_context(); - ASSERT_NE(context, nullptr); - context->set_topk(10); - - ASSERT_EQ(0, searcher->search_impl(query_vec.data(), query_meta, 1, context)); - - const auto &result = context->result(0); - ASSERT_GT(result.size(), 0UL); - ASSERT_LE(result.size(), 10UL); -} - -TEST_F(HnswRabitqBuilderTest, TestDimensions) { - std::vector dimensions = {1, 2, 4, 8, 16, 32, 33, - 63, 64, 128, 256, 512, 1024, 2047, - 2048, 2049, 4095, 4096, 4097, 8192, 16384}; - size_t doc_cnt = 100; - - for (size_t test_dim : dimensions) { - std::cout << "Testing dimension: " << test_dim << std::endl; - - IndexMeta index_meta(IndexMeta::DataType::DT_FP32, test_dim); - index_meta.set_metric("SquaredEuclidean", 0, ailego::Params()); - - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr) << "dim=" << test_dim; - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", test_dim); - - int ret = builder->init(index_meta, params); - - // dimension <= 63 or >= 4096: init() should return -31 - if (test_dim <= 63 || test_dim >= 4096) { - ASSERT_EQ(-31, ret) << "expected init to fail with -31, dim=" << test_dim; - std::cout << "Dimension " << test_dim - << " correctly rejected with ret=" << ret << std::endl; - continue; - } - - // Valid dimensions: verify full build succeeds - ASSERT_EQ(0, ret) << "init failed, dim=" << test_dim; - - auto holder = - make_shared>( - test_dim); - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(test_dim); - for (size_t j = 0; j < test_dim; ++j) { - vec[j] = static_cast(i * test_dim + j) / 1000.0f; - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))) << "dim=" << test_dim; - } - - ret = builder->train(holder); - ASSERT_EQ(0, ret) << "train failed, dim=" << test_dim; - - ret = builder->build(holder); - ASSERT_EQ(0, ret) << "build failed, dim=" << test_dim; - - auto &stats = builder->stats(); - ASSERT_EQ(doc_cnt, stats.built_count()) << "dim=" << test_dim; - - std::cout << "Dimension " << test_dim << " passed, built " - << stats.built_count() << " docs" << std::endl; - } -} - -TEST_F(HnswRabitqBuilderTest, TestMemquota) { - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - 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 * dim + j) / 1000.0f; - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - params.set("proxima.hnsw_rabitq.builder.memory_quota", 100000UL); - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(IndexError_NoMemory, builder->build(holder)); -} - -TEST_F(HnswRabitqBuilderTest, TestIndexThreads) { - IndexBuilder::Pointer builder1 = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder1, nullptr); - IndexBuilder::Pointer builder2 = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder2, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - 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 * dim + j) / 1000.0f; - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - std::srand(ailego::Realtime::MilliSeconds()); - auto threads = - std::make_shared(std::rand() % 4, false); - ASSERT_EQ(0, builder1->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder2->init(*_index_meta_ptr, params)); - - auto build_index1 = [&]() { - ASSERT_EQ(0, builder1->train(threads, holder)); - ASSERT_EQ(0, builder1->build(threads, holder)); - }; - auto build_index2 = [&]() { - ASSERT_EQ(0, builder2->train(threads, holder)); - ASSERT_EQ(0, builder2->build(threads, holder)); - }; - - auto t1 = std::async(std::launch::async, build_index1); - auto t2 = std::async(std::launch::async, build_index2); - t1.wait(); - t2.wait(); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestIndexThreads"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder1->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder2->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats1 = builder1->stats(); - ASSERT_EQ(doc_cnt, stats1.built_count()); - auto &stats2 = builder2->stats(); - ASSERT_EQ(doc_cnt, stats2.built_count()); -} - -TEST_F(HnswRabitqBuilderTest, TestCosine) { - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - 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 * dim + j) / 1000.0f; - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineFp32Converter"); - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - converter->transform(holder); - - auto converted_holder = converter->result(); - converted_holder = convert_holder_to_provider(converted_holder); - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - ASSERT_EQ(0, builder->init(index_meta, params)); - - ASSERT_EQ(0, builder->train(converted_holder)); - - ASSERT_EQ(0, builder->build(converted_holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestCosine"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats = builder->stats(); - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt, stats.built_count()); - ASSERT_EQ(doc_cnt, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); -} - -TEST_F(HnswRabitqBuilderTest, TestCleanupAndRebuild) { - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 1000UL; - 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 * dim + j) / 1000.0f; - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestCleanupAndRebuild"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - auto &stats = builder->stats(); - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt, stats.built_count()); - ASSERT_EQ(doc_cnt, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); - - // Cleanup and rebuild with more documents - ASSERT_EQ(0, builder->cleanup()); - - auto holder2 = - make_shared>(dim); - size_t doc_cnt2 = 2000UL; - for (size_t i = 0; i < doc_cnt2; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = static_cast(i * dim + j) / 1000.0f; - } - ASSERT_TRUE(holder2->emplace(i, std::move(vec))); - } - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder2)); - ASSERT_EQ(0, builder->build(holder2)); - - auto dumper2 = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper2, nullptr); - ASSERT_EQ(0, dumper2->create(path)); - ASSERT_EQ(0, builder->dump(dumper2)); - ASSERT_EQ(0, dumper2->close()); - - ASSERT_EQ(0UL, stats.trained_count()); - ASSERT_EQ(doc_cnt2, stats.built_count()); - ASSERT_EQ(doc_cnt2, stats.dumped_count()); - ASSERT_EQ(0UL, stats.discarded_count()); - ASSERT_EQ(0UL, stats.trained_costtime()); - ASSERT_GT(stats.built_costtime(), 0UL); -} - -} // namespace core -} // namespace zvec - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic pop -#endif diff --git a/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_test.cc b/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_test.cc deleted file mode 100644 index d3b711e..0000000 --- a/tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_test.cc +++ /dev/null @@ -1,573 +0,0 @@ -// 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 "hnsw_rabitq_searcher.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include "zvec/core/framework/index_framework.h" -#include "zvec/core/framework/index_logger.h" -#include "hnsw_rabitq_builder.h" - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wunused-result" -#endif - -using namespace std; -using namespace zvec::ailego; - -namespace zvec { -namespace core { - -constexpr size_t static dim = 128; - -class HnswRabitqSearcherTest : public testing::Test { - protected: - void SetUp(void); - void TearDown(void); - - static std::string _dir; - static shared_ptr _index_meta_ptr; -}; - -std::string HnswRabitqSearcherTest::_dir("HnswRabitqSearcherTest"); -shared_ptr HnswRabitqSearcherTest::_index_meta_ptr; - -void HnswRabitqSearcherTest::SetUp(void) { - IndexLoggerBroker::SetLevel(0); - _index_meta_ptr.reset(new (nothrow) - IndexMeta(IndexMeta::DataType::DT_FP32, dim)); - _index_meta_ptr->set_metric("SquaredEuclidean", 0, ailego::Params()); -} - -void HnswRabitqSearcherTest::TearDown(void) { - char cmdBuf[100]; - snprintf(cmdBuf, 100, "rm -rf %s", _dir.c_str()); - // system(cmdBuf); -} - -TEST_F(HnswRabitqSearcherTest, TestBasicSearch) { - // Build index first - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 10000UL; - 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 * dim + j) / 1000.0f; - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestBasicSearch"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // Test searcher - auto searcher = IndexFactory::CreateSearcher("HnswRabitqSearcher"); - ASSERT_NE(searcher, nullptr); - - ailego::Params search_params; - search_params.set("proxima.hnsw_rabitq.searcher.ef", 100UL); - ASSERT_EQ(0, searcher->init(search_params)); - - auto loader = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_NE(loader, nullptr); - ASSERT_EQ(0, loader->init(ailego::Params())); - ASSERT_EQ(0, loader->open(path, false)); - - ASSERT_EQ(0, searcher->load(loader, nullptr)); - - // Perform search - NumericalVector query_vec(dim); - for (size_t j = 0; j < dim; ++j) { - query_vec[j] = static_cast(j) / 1000.0f; - } - - IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); - - auto context = searcher->create_context(); - ASSERT_TRUE(!!context); - context->set_topk(10); - - ASSERT_EQ(0, searcher->search_impl(query_vec.data(), query_meta, 1, context)); - - const auto &result = context->result(0); - ASSERT_GT(result.size(), 0UL); - ASSERT_LE(result.size(), 10UL); - - // Verify results are sorted by distance - for (size_t i = 1; i < result.size(); ++i) { - ASSERT_LE(result[i - 1].score(), result[i].score()); - } -} - -TEST_F(HnswRabitqSearcherTest, DISABLED_TestRnnSearch) { - // Build index first - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 10000UL; - 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); - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestRnnSearch"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // Test searcher with radius search - auto searcher = IndexFactory::CreateSearcher("HnswRabitqSearcher"); - ASSERT_NE(searcher, nullptr); - - ailego::Params search_params; - search_params.set("proxima.hnsw_rabitq.searcher.ef", 100UL); - ASSERT_EQ(0, searcher->init(search_params)); - - auto loader = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_NE(loader, nullptr); - ASSERT_EQ(0, loader->init(ailego::Params())); - ASSERT_EQ(0, loader->open(path, false)); - - ASSERT_EQ(0, searcher->load(loader, nullptr)); - - NumericalVector query_vec(dim); - for (size_t j = 0; j < dim; ++j) { - query_vec[j] = 0.0f; - } - - IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); - auto context = searcher->create_context(); - ASSERT_NE(context, nullptr); - - size_t topk = 50; - context->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(query_vec.data(), query_meta, 1, context)); - - const auto &results = context->result(0); - ASSERT_EQ(topk, results.size()); - - // Test with radius threshold - float radius = results[topk / 2].score(); - context->set_threshold(radius); - ASSERT_EQ(0, searcher->search_impl(query_vec.data(), query_meta, 1, context)); - ASSERT_GT(topk, results.size()); - for (size_t k = 0; k < results.size(); ++k) { - ASSERT_GE(radius, results[k].score()); - } - - // Test reset threshold - context->reset_threshold(); - ASSERT_EQ(0, searcher->search_impl(query_vec.data(), query_meta, 1, context)); - ASSERT_EQ(topk, results.size()); - ASSERT_LT(radius, results[topk - 1].score()); -} - -TEST_F(HnswRabitqSearcherTest, DISABLED_TestSearchInnerProduct) { - // Build index with InnerProduct metric - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 10000UL; - 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); - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - IndexMeta index_meta(IndexMeta::DataType::DT_FP32, dim); - index_meta.set_metric("InnerProduct", 0, ailego::Params()); - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - ASSERT_EQ(0, builder->init(index_meta, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestSearchInnerProduct"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // Test searcher - auto searcher = IndexFactory::CreateSearcher("HnswRabitqSearcher"); - ASSERT_NE(searcher, nullptr); - - ailego::Params search_params; - search_params.set("proxima.hnsw_rabitq.searcher.ef", 100UL); - ASSERT_EQ(0, searcher->init(search_params)); - - auto loader = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_NE(loader, nullptr); - ASSERT_EQ(0, loader->init(ailego::Params())); - ASSERT_EQ(0, loader->open(path, false)); - - ASSERT_EQ(0, searcher->load(loader, nullptr)); - - NumericalVector query_vec(dim); - for (size_t j = 0; j < dim; ++j) { - query_vec[j] = 1.0f; - } - - IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); - auto context = searcher->create_context(); - ASSERT_TRUE(!!context); - - size_t topk = 50; - context->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(query_vec.data(), query_meta, 1, context)); - - const auto &results = context->result(0); - ASSERT_EQ(topk, results.size()); - - // Test with radius threshold (note: InnerProduct uses negative scores) - float radius = -results[topk / 2].score(); - context->set_threshold(radius); - ASSERT_EQ(0, searcher->search_impl(query_vec.data(), query_meta, 1, context)); - ASSERT_GT(topk, results.size()); - for (size_t k = 0; k < results.size(); ++k) { - LOG_ERROR("radius: %f, score: %f", radius, results[k].score()); - EXPECT_GE(radius, results[k].score()); - } - - // Test reset threshold - context->reset_threshold(); - ASSERT_EQ(0, searcher->search_impl(query_vec.data(), query_meta, 1, context)); - ASSERT_EQ(topk, results.size()); - ASSERT_LT(-radius, results[topk - 1].score()); -} - -TEST_F(HnswRabitqSearcherTest, TestSearchCosine) { - // Build index with Cosine metric - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 10000UL; - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_real_distribution dist(-1.0, 1.0); - - for (size_t i = 0; i < doc_cnt; i++) { - NumericalVector vec(dim); - for (size_t j = 0; j < dim; ++j) { - vec[j] = dist(gen); - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, dim); - index_meta_raw.set_metric("Cosine", 0, ailego::Params()); - - ailego::Params converter_params; - auto converter = IndexFactory::CreateConverter("CosineFp32Converter"); - converter->init(index_meta_raw, converter_params); - - IndexMeta index_meta = converter->meta(); - - converter->transform(holder); - - auto converted_holder = converter->result(); - converted_holder = convert_holder_to_provider(converted_holder); - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - ASSERT_EQ(0, builder->init(index_meta, params)); - ASSERT_EQ(0, builder->train(converted_holder)); - ASSERT_EQ(0, builder->build(converted_holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestSearchCosine"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // Test searcher - auto searcher = IndexFactory::CreateSearcher("HnswRabitqSearcher"); - ASSERT_NE(searcher, nullptr); - - ailego::Params search_params; - search_params.set("proxima.hnsw_rabitq.searcher.ef", 100UL); - ASSERT_EQ(0, searcher->init(search_params)); - - auto loader = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_NE(loader, nullptr); - ASSERT_EQ(0, loader->init(ailego::Params())); - ASSERT_EQ(0, loader->open(path, false)); - - ASSERT_EQ(0, searcher->load(loader, nullptr)); - - NumericalVector query_vec(dim); - for (size_t j = 0; j < dim; ++j) { - query_vec[j] = 1.0f; - } - - IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); - auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name()); - ASSERT_TRUE(reformer != nullptr); - - ASSERT_EQ(0, reformer->init(index_meta.reformer_params())); - - std::string new_query; - IndexQueryMeta new_meta; - ASSERT_EQ(0, reformer->transform(query_vec.data(), query_meta, &new_query, - &new_meta)); - - auto context = searcher->create_context(); - ASSERT_TRUE(!!context); - - size_t topk = 50; - context->set_topk(topk); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, 1, context)); - - const auto &results = context->result(0); - ASSERT_EQ(topk, results.size()); - - // Test with radius threshold - float radius = 0.5f; - context->set_threshold(radius); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, 1, context)); - ASSERT_GT(topk, results.size()); - for (size_t k = 0; k < results.size(); ++k) { - ASSERT_GE(radius, results[k].score()); - } - - // Test reset threshold - context->reset_threshold(); - ASSERT_EQ(0, searcher->search_impl(new_query.data(), new_meta, 1, context)); - ASSERT_EQ(topk, results.size()); - ASSERT_LT(radius, results[topk - 1].score()); -} - -TEST_F(HnswRabitqSearcherTest, TestMultipleQueries) { - // Build index first - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 10000UL; - 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 * dim + j) / 1000.0f; - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestMultipleQueries"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // Test searcher with multiple queries - auto searcher = IndexFactory::CreateSearcher("HnswRabitqSearcher"); - ASSERT_NE(searcher, nullptr); - - ailego::Params search_params; - search_params.set("proxima.hnsw_rabitq.searcher.ef", 100UL); - ASSERT_EQ(0, searcher->init(search_params)); - - auto loader = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_NE(loader, nullptr); - ASSERT_EQ(0, loader->init(ailego::Params())); - ASSERT_EQ(0, loader->open(path, false)); - - ASSERT_EQ(0, searcher->load(loader, nullptr)); - - // Test with different query vectors - for (size_t query_id = 0; query_id < 5; ++query_id) { - NumericalVector query_vec(dim); - for (size_t j = 0; j < dim; ++j) { - query_vec[j] = static_cast(query_id * dim + j) / 1000.0f; - } - - IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); - - auto context = searcher->create_context(); - ASSERT_TRUE(!!context); - context->set_topk(20); - - ASSERT_EQ(0, - searcher->search_impl(query_vec.data(), query_meta, 1, context)); - - const auto &result = context->result(0); - ASSERT_GT(result.size(), 0UL); - ASSERT_LE(result.size(), 20UL); - - // Verify results are sorted - for (size_t i = 1; i < result.size(); ++i) { - ASSERT_LE(result[i - 1].score(), result[i].score()); - } - } -} - -TEST_F(HnswRabitqSearcherTest, TestDifferentTopK) { - // Build index first - IndexBuilder::Pointer builder = - IndexFactory::CreateBuilder("HnswRabitqBuilder"); - ASSERT_NE(builder, nullptr); - - auto holder = - make_shared>(dim); - size_t doc_cnt = 10000UL; - 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 * dim + j) / 1000.0f; - } - ASSERT_TRUE(holder->emplace(i, std::move(vec))); - } - - ailego::Params params; - params.set("proxima.rabitq.num_clusters", 16UL); - params.set("proxima.rabitq.total_bits", 2UL); - params.set("proxima.hnsw_rabitq.general.dimension", dim); - - ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); - ASSERT_EQ(0, builder->train(holder)); - ASSERT_EQ(0, builder->build(holder)); - - auto dumper = IndexFactory::CreateDumper("FileDumper"); - ASSERT_NE(dumper, nullptr); - - string path = _dir + "/TestDifferentTopK"; - ASSERT_EQ(0, dumper->create(path)); - ASSERT_EQ(0, builder->dump(dumper)); - ASSERT_EQ(0, dumper->close()); - - // Test searcher with different topk values - auto searcher = IndexFactory::CreateSearcher("HnswRabitqSearcher"); - ASSERT_NE(searcher, nullptr); - - ailego::Params search_params; - search_params.set("proxima.hnsw_rabitq.searcher.ef", 100UL); - ASSERT_EQ(0, searcher->init(search_params)); - - auto loader = IndexFactory::CreateStorage("FileReadStorage"); - ASSERT_NE(loader, nullptr); - ASSERT_EQ(0, loader->init(ailego::Params())); - ASSERT_EQ(0, loader->open(path, false)); - - ASSERT_EQ(0, searcher->load(loader, nullptr)); - - NumericalVector query_vec(dim); - for (size_t j = 0; j < dim; ++j) { - query_vec[j] = static_cast(j) / 1000.0f; - } - - IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); - - // Test with different topk values - std::vector topk_values = {1, 5, 10, 20, 50, 100}; - for (size_t topk : topk_values) { - auto context = searcher->create_context(); - ASSERT_TRUE(!!context); - context->set_topk(topk); - - ASSERT_EQ(0, - searcher->search_impl(query_vec.data(), query_meta, 1, context)); - - const auto &result = context->result(0); - ASSERT_GT(result.size(), 0UL); - ASSERT_LE(result.size(), topk); - - // Verify results are sorted - for (size_t i = 1; i < result.size(); ++i) { - ASSERT_LE(result[i - 1].score(), result[i].score()); - } - } -} - -} // namespace core -} // namespace zvec - -#if defined(__GNUC__) || defined(__GNUG__) -#pragma GCC diagnostic pop -#endif diff --git a/tools/core/local_builder.cc b/tools/core/local_builder.cc index 575db10..59ba816 100644 --- a/tools/core/local_builder.cc +++ b/tools/core/local_builder.cc @@ -1156,8 +1156,6 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) { converter_name, &cv_build_holder) != 0) { return -1; } - } else if (builder_class == "HnswRabitqBuilder" && !converter_name.empty()) { - cv_build_holder = convert_holder_to_provider(cv_build_holder); } // BUILD diff --git a/tools/core/local_builder_original.cc b/tools/core/local_builder_original.cc index 29f6093..60aa9a2 100644 --- a/tools/core/local_builder_original.cc +++ b/tools/core/local_builder_original.cc @@ -1105,8 +1105,6 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) { converter_name, &cv_build_holder) != 0) { return -1; } - } else if (builder_class == "HnswRabitqBuilder" && !converter_name.empty()) { - cv_build_holder = convert_holder_to_provider(cv_build_holder); } // BUILD