xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

chore: rm hnsw-rabitq searcher/streamer (#367) 

* chore: rm hnsw-rabitq searcher/streamer

* fix clang-tidy

* fix
This commit is contained in:
egolearner 2026-04-23 10:20:46 +08:00 committed by GitHub
parent d9cf689ce1
commit 2b16ed39cd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 1 additions and 3339 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/clang_tidy.yml vendored
View File

@ -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: |

2
src/core/algorithm/hnsw/hnsw_streamer.cc
View File

@ -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.");

554
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.cc
View File

@ -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 <cstdlib>
#include <iostream>
#include <thread>
#include <ailego/pattern/defer.h>
#include <zvec/ailego/container/params.h>
#include <zvec/ailego/utility/string_helper.h>
#include <zvec/ailego/utility/time_helper.h>
#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 &params) {
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<size_t>(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<RabitqConverter>();
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<RabitqReformer>();
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<IndexProvider>(holder);
if (!provider) {
LOG_ERROR("Rabitq builder expect IndexProvider");
return IndexError_InvalidArgument;
}
if (!threads) {
threads = std::make_shared<SingleQueueIndexThreads>(thread_cnt_, false);
}
auto start_time = ailego::Monotime::MilliSeconds();
LOG_INFO("Begin HnswRabitqBuilder::build");
if (holder->count() != static_cast<size_t>(-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<node_id_t> 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<std::mutex> 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<size_t>(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<node_id_t> *finished) {
AILEGO_DEFER([&]() {
std::lock_guard<std::mutex> latch(mutex_);
cond_.notify_one();
});
HnswRabitqContext *ctx = new (std::nothrow) HnswRabitqContext(
meta_.dimension(), metric_,
std::shared_ptr<HnswRabitqEntity>(&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<size_t>(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

101
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder.h
View File

@ -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 <zvec/ailego/parallel/thread_pool.h>
#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 &params) 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<node_id_t> *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

199
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.cc
View File

@ -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 <iostream>
#include <zvec/ailego/hash/crc32c.h>
#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<size_t>(doc_cnt()));
return IndexError_NoMemory;
}
vectors_buffer_.append(reinterpret_cast<const char *>(vec), vector_size());
vectors_buffer_.append(padding_size_, '\0');
keys_buffer_.append(reinterpret_cast<const char *>(&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<const key_t *>(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<IndexStorage::MemoryBlock> &vec_blocks) const {
std::vector<const void *> 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<std::pair<node_id_t, ResultRecord>> &neighbors) {
NeighborsHeader *hd =
const_cast<NeighborsHeader *>(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<NeighborsHeader *>(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<key_t *>(const_cast<char *>(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

139
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_entity.h
View File

@ -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 <zvec/ailego/internal/platform.h>
#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<IndexStorage::MemoryBlock> &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<const NeighborsHeader *>(
neighbors_buffer_.data() + neighbors_size_ * id);
} else {
size_t offset = neighbors_index_[id].offset;
return reinterpret_cast<const NeighborsHeader *>(
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<std::pair<node_id_t, ResultRecord>> &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<NeighborIndex> 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

4
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_register.cc
View File

@ -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

516
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.cc
View File

@ -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 <rabitqlib/utils/rotator.hpp>
#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<uint32_t>(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<HnswRabitqContext *>(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<const char *>(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<HnswRabitqContext *>(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<std::string(node_id_t)> 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<const char *>(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<const char *>(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<std::vector<uint64_t>> &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<HnswRabitqContext *>(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<std::string(node_id_t)> 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<const char *>(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<const char *>(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

142
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher.h
View File

@ -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 &params) 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<std::vector<uint64_t>> &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<std::vector<uint64_t>> &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 &params(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

515
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.cc
View File

@ -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 <zvec/ailego/hash/crc32c.h>
#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<const key_t *>(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<const node_id_t *>(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<IndexStorage::MemoryBlock> &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<const NeighborsHeader *>(
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<uint32_t>(m->neighbor_cnt),
reinterpret_cast<const node_id_t *>(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<const NeighborsHeader *>(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<const void **>(&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<const void **>(&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<SegmentPointer> 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<char[]>());
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<const void **>(&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<const void **>(&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<uint32_t *>(dst) = neighbors_index[id].neighbor_cnt;
memcpy(dst + sizeof(uint32_t), neighbor_data, rd_size);
}
return 0;
}
int HnswRabitqSearcherEntity::get_fixed_neighbors(
std::vector<uint32_t> *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<const void **>(&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<const void **>(&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<SegmentPointer> &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

158
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_entity.h
View File

@ -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<HnswRabitqSearcherEntity>;
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<IndexStorage::MemoryBlock> &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<uint32_t> *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<char> &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<SegmentPointer> &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<IndexStorage::SegmentData> segment_datas_{};
std::shared_ptr<char> fixed_neighbors_{}; // level 0 fixed size neighbors
bool neighbors_in_memory_enabled_{false};
bool loaded_{false};
};
} // namespace core
} // namespace zvec

3
src/core/algorithm/hnsw_rabitq/hnsw_rabitq_streamer.cc
View File

@ -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.");

428
tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_builder_test.cc
View File

@ -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 <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <future>
#include <gtest/gtest.h>
#include <zvec/ailego/container/params.h>
#include <zvec/ailego/container/vector.h>
#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<IndexMeta> _index_meta_ptr;
};
std::string HnswRabitqBuilderTest::_dir("hnswRabitqBuilderTest");
shared_ptr<IndexMeta> 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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 1000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<float> query_vec(dim);
for (size_t j = 0; j < dim; ++j) {
query_vec[j] = static_cast<float>(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<size_t> 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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(
test_dim);
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(test_dim);
for (size_t j = 0; j < test_dim; ++j) {
vec[j] = static_cast<float>(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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 1000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 1000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<SingleQueueIndexThreads>(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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 1000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 1000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt2 = 2000UL;
for (size_t i = 0; i < doc_cnt2; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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

573
tests/core/algorithm/hnsw_rabitq/hnsw_rabitq_searcher_test.cc
View File

@ -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 <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <cstdio>
#include <random>
#include <gtest/gtest.h>
#include <zvec/ailego/container/params.h>
#include <zvec/ailego/container/vector.h>
#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<IndexMeta> _index_meta_ptr;
};
std::string HnswRabitqSearcherTest::_dir("HnswRabitqSearcherTest");
shared_ptr<IndexMeta> 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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 10000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<float> query_vec(dim);
for (size_t j = 0; j < dim; ++j) {
query_vec[j] = static_cast<float>(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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 10000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<float> 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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 10000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<float> 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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 10000UL;
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_real_distribution<float> dist(-1.0, 1.0);
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> 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<float> 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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 10000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<float> query_vec(dim);
for (size_t j = 0; j < dim; ++j) {
query_vec[j] = static_cast<float>(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<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
size_t doc_cnt = 10000UL;
for (size_t i = 0; i < doc_cnt; i++) {
NumericalVector<float> vec(dim);
for (size_t j = 0; j < dim; ++j) {
vec[j] = static_cast<float>(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<float> query_vec(dim);
for (size_t j = 0; j < dim; ++j) {
query_vec[j] = static_cast<float>(j) / 1000.0f;
}
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
// Test with different topk values
std::vector<size_t> 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

2
tools/core/local_builder.cc
View File

@ -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

2
tools/core/local_builder_original.cc
View File

@ -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