chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,14 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_hnsw_sparse
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/hnsw_sparse
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,469 @@
|
||||
// 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_sparse_builder.h"
|
||||
#include <sys/stat.h>
|
||||
#include <sys/types.h>
|
||||
#include <fcntl.h>
|
||||
#include <future>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
#include "tests/test_util.h"
|
||||
#include "zvec/core/framework/index_framework.h"
|
||||
#include "hnsw_sparse_params.h"
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
using namespace std;
|
||||
using namespace testing;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
class HnswSparseBuilderTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
|
||||
static std::string _dir;
|
||||
static shared_ptr<IndexMeta> _index_meta_ptr;
|
||||
};
|
||||
|
||||
std::string HnswSparseBuilderTest::_dir("HnswSparseBuilderTest/");
|
||||
shared_ptr<IndexMeta> HnswSparseBuilderTest::_index_meta_ptr;
|
||||
|
||||
void HnswSparseBuilderTest::SetUp(void) {
|
||||
_index_meta_ptr.reset(new (nothrow) IndexMeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32));
|
||||
_index_meta_ptr->set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
}
|
||||
|
||||
void HnswSparseBuilderTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(_dir);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestGeneral) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<OnePassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, 1);
|
||||
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);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
|
||||
auto holder2 =
|
||||
make_shared<MultiPassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
size_t doc_cnt2 = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt2; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder2->emplace(i, 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);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestMemquota) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<OnePassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
size_t doc_cnt = 1000UL;
|
||||
uint32_t sparse_count = 32;
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw.sparse_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(HnswSparseBuilderTest, TestIndexThreads) {
|
||||
IndexBuilder::Pointer builder1 =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder1, nullptr);
|
||||
IndexBuilder::Pointer builder2 =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder2, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
|
||||
size_t doc_cnt = 1000UL;
|
||||
uint32_t sparse_count = 32;
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
ailego::Params params;
|
||||
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(HnswSparseBuilderTest, TestHalfFloatConverter) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<OnePassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
ailego::Params converter_params;
|
||||
auto converter = IndexFactory::CreateConverter("HalfFloatSparseConverter");
|
||||
converter->init(*_index_meta_ptr, converter_params);
|
||||
|
||||
IndexMeta index_meta = converter->meta();
|
||||
|
||||
converter->transform(holder);
|
||||
|
||||
auto converted_holder = converter->sparse_result();
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, 1);
|
||||
ASSERT_EQ(0, builder->init(index_meta, converter_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 + "TestHalFloatConverter";
|
||||
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);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
|
||||
auto holder2 =
|
||||
make_shared<MultiPassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
size_t doc_cnt2 = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt2; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder2->emplace(i, 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);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestIndptr) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
|
||||
std::vector<uint64_t> keys;
|
||||
keys.reserve(doc_cnt);
|
||||
|
||||
std::vector<uint64_t> sparse_indptr;
|
||||
sparse_indptr.reserve(doc_cnt + 1);
|
||||
|
||||
std::vector<uint32_t> sparse_indices;
|
||||
sparse_indices.reserve(doc_cnt * sparse_count);
|
||||
|
||||
std::vector<float> sparse_values;
|
||||
sparse_values.reserve(doc_cnt * sparse_count);
|
||||
|
||||
size_t sparse_count_total = 0;
|
||||
sparse_indptr.push_back(0);
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices.push_back(20 * j);
|
||||
sparse_values.push_back(i);
|
||||
}
|
||||
|
||||
keys.push_back(i);
|
||||
|
||||
sparse_count_total += sparse_count;
|
||||
sparse_indptr.push_back(sparse_count_total);
|
||||
}
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, 1);
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->build(doc_cnt, keys.data(), sparse_indptr.data(),
|
||||
sparse_indices.data(), sparse_values.data()));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestIndptr";
|
||||
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);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestIndptrFp16) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
|
||||
std::vector<uint64_t> keys;
|
||||
keys.reserve(doc_cnt);
|
||||
|
||||
std::vector<uint64_t> sparse_indptr;
|
||||
sparse_indptr.reserve(doc_cnt + 1);
|
||||
|
||||
std::vector<uint32_t> sparse_indices;
|
||||
sparse_indices.reserve(doc_cnt * sparse_count);
|
||||
|
||||
std::vector<float> sparse_values;
|
||||
sparse_values.reserve(doc_cnt * sparse_count);
|
||||
|
||||
size_t sparse_count_total = 0;
|
||||
sparse_indptr.push_back(0);
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices.push_back(20 * j);
|
||||
sparse_values.push_back(i);
|
||||
}
|
||||
|
||||
keys.push_back(i);
|
||||
|
||||
sparse_count_total += sparse_count;
|
||||
sparse_indptr.push_back(sparse_count_total);
|
||||
}
|
||||
|
||||
IndexMeta meta(IndexMeta::MetaType::MT_SPARSE, IndexMeta::DataType::DT_FP16);
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, 1);
|
||||
ASSERT_EQ(0, builder->init(meta, params));
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
ASSERT_EQ(0, builder->build(qmeta, doc_cnt, keys.data(), sparse_indptr.data(),
|
||||
sparse_indices.data(), sparse_values.data()));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestIndptrFp16";
|
||||
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);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,360 @@
|
||||
// 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 <sys/stat.h>
|
||||
#include <sys/types.h>
|
||||
#include <fcntl.h>
|
||||
#include <future>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <ailego/math/norm_matrix.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
#include "tests/test_util.h"
|
||||
#include "hnsw_sparse_streamer.h"
|
||||
|
||||
using namespace std;
|
||||
using namespace testing;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
class HnswSparseStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void);
|
||||
void TearDown(void);
|
||||
void generate_sparse_data(
|
||||
size_t cnt, uint32_t sparse_dim_count,
|
||||
std::vector<NumericalVector<uint32_t>> &sparse_indices_list,
|
||||
std::vector<NumericalVector<float>> &sparse_vec_list, bool norm);
|
||||
|
||||
static std::string dir_;
|
||||
static shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string HnswSparseStreamerTest::dir_(
|
||||
"hnsw_sparse_streamer_buffer_test_dir/");
|
||||
shared_ptr<IndexMeta> HnswSparseStreamerTest::index_meta_ptr_;
|
||||
|
||||
void HnswSparseStreamerTest::generate_sparse_data(
|
||||
size_t cnt, uint32_t sparse_dim_count,
|
||||
std::vector<NumericalVector<uint32_t>> &sparse_indices_list,
|
||||
std::vector<NumericalVector<float>> &sparse_vec_list, bool norm) {
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(-1.0, 1.0);
|
||||
|
||||
for (size_t i = 0; i < cnt; ++i) {
|
||||
// prepare sparse
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
|
||||
NumericalVector<float> sparse_vec(sparse_dim_count);
|
||||
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_vec[j] = dist(gen);
|
||||
}
|
||||
|
||||
float norm;
|
||||
ailego::Norm2Matrix<float, 1>::Compute(sparse_vec.data(), sparse_dim_count,
|
||||
&norm);
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_vec[j] = sparse_vec[j] / norm;
|
||||
}
|
||||
|
||||
sparse_indices_list.push_back(sparse_indices);
|
||||
sparse_vec_list.push_back(sparse_vec);
|
||||
}
|
||||
}
|
||||
|
||||
void HnswSparseStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (nothrow) IndexMeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32));
|
||||
index_meta_ptr_->set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void HnswSparseStreamerTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseStreamerTest, TestGeneral) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("HnswSparseStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
size_t sparse_dim_count = 32;
|
||||
|
||||
IndexMeta index_meta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
index_meta.set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT, 20);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR, 16);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION, 10);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_EF, 5);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U);
|
||||
|
||||
ailego::Params stg_params;
|
||||
auto write_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_EQ(0, write_storage->init(stg_params));
|
||||
ASSERT_EQ(0, write_storage->open(dir_ + "Test/HnswSparseSearch", true));
|
||||
ASSERT_EQ(0, write_streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, write_streamer->open(write_storage));
|
||||
|
||||
size_t cnt = 20000U;
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
std::vector<NumericalVector<uint32_t>> sparse_indices_list;
|
||||
std::vector<NumericalVector<float>> sparse_vec_list;
|
||||
|
||||
generate_sparse_data(cnt, sparse_dim_count, sparse_indices_list,
|
||||
sparse_vec_list, true);
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(0, write_streamer->add_impl(
|
||||
i, sparse_dim_count, sparse_indices_list[i].data(),
|
||||
sparse_vec_list[i].data(), qmeta, ctx));
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
write_storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("HnswSparseStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/HnswSparseSearch", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
|
||||
auto linearCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!linearCtx);
|
||||
|
||||
auto knnCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!knnCtx);
|
||||
|
||||
// streamer->print_debug_info();
|
||||
size_t topk = 200;
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
|
||||
uint64_t knnTotalTime = 0;
|
||||
uint64_t linearTotalTime = 0;
|
||||
|
||||
int totalHits = 0;
|
||||
int totalCnts = 0;
|
||||
int topk1Hits = 0;
|
||||
|
||||
for (size_t i = 0; i < cnt; i += 100) {
|
||||
const auto &sparse_indices = sparse_indices_list[i];
|
||||
const auto &sparse_vec = sparse_vec_list[i];
|
||||
|
||||
auto t1 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(
|
||||
0, read_streamer->search_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_vec.data(), qmeta, knnCtx));
|
||||
|
||||
auto t2 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(
|
||||
sparse_dim_count, sparse_indices.data(), sparse_vec.data(),
|
||||
qmeta, linearCtx));
|
||||
|
||||
auto t3 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
knnTotalTime += t2 - t1;
|
||||
linearTotalTime += t3 - t2;
|
||||
|
||||
// std::cout << "i: " << i << std::endl;
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
topk1Hits += i == knnResult[0].key();
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
float recall = totalHits * 1.0f / totalCnts;
|
||||
float topk1Recall = topk1Hits * 100.0f / cnt;
|
||||
float cost = linearTotalTime * 1.0f / knnTotalTime;
|
||||
#if 0
|
||||
printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d "
|
||||
"R@%zd=%f R@1=%f cost=%f\n",
|
||||
knnTotalTime, linearTotalTime, totalHits, totalCnts, topk, recall,
|
||||
topk1Recall, cost);
|
||||
#endif
|
||||
EXPECT_GT(recall, 0.80f);
|
||||
EXPECT_GT(topk1Recall, 0.80f);
|
||||
// EXPECT_GT(cost, 2.0f);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseStreamerTest, TestHnswSearchMMap) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("HnswSparseStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
size_t sparse_dim_count = 32;
|
||||
|
||||
IndexMeta index_meta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
index_meta.set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT, 20);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR, 16);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION, 10);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_EF, 5);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U);
|
||||
|
||||
ailego::Params stg_params;
|
||||
auto write_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_EQ(0, write_storage->init(stg_params));
|
||||
ASSERT_EQ(0, write_storage->open(dir_ + "Test/HnswSparseSearch", true));
|
||||
ASSERT_EQ(0, write_streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, write_streamer->open(write_storage));
|
||||
|
||||
size_t cnt = 20000U;
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
std::vector<NumericalVector<uint32_t>> sparse_indices_list;
|
||||
std::vector<NumericalVector<float>> sparse_vec_list;
|
||||
|
||||
generate_sparse_data(cnt, sparse_dim_count, sparse_indices_list,
|
||||
sparse_vec_list, true);
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(0, write_streamer->add_impl(
|
||||
i, sparse_dim_count, sparse_indices_list[i].data(),
|
||||
sparse_vec_list[i].data(), qmeta, ctx));
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
write_storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("HnswSparseStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/HnswSparseSearch", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
|
||||
auto linearCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!linearCtx);
|
||||
|
||||
auto knnCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!knnCtx);
|
||||
|
||||
// streamer->print_debug_info();
|
||||
size_t topk = 200;
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
|
||||
uint64_t knnTotalTime = 0;
|
||||
uint64_t linearTotalTime = 0;
|
||||
|
||||
int totalHits = 0;
|
||||
int totalCnts = 0;
|
||||
int topk1Hits = 0;
|
||||
|
||||
for (size_t i = 0; i < cnt; i += 100) {
|
||||
const auto &sparse_indices = sparse_indices_list[i];
|
||||
const auto &sparse_vec = sparse_vec_list[i];
|
||||
|
||||
auto t1 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(
|
||||
0, read_streamer->search_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_vec.data(), qmeta, knnCtx));
|
||||
|
||||
auto t2 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(
|
||||
sparse_dim_count, sparse_indices.data(), sparse_vec.data(),
|
||||
qmeta, linearCtx));
|
||||
|
||||
auto t3 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
knnTotalTime += t2 - t1;
|
||||
linearTotalTime += t3 - t2;
|
||||
|
||||
// std::cout << "i: " << i << std::endl;
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
topk1Hits += i == knnResult[0].key();
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
float recall = totalHits * 1.0f / totalCnts;
|
||||
float topk1Recall = topk1Hits * 100.0f / cnt;
|
||||
float cost = linearTotalTime * 1.0f / knnTotalTime;
|
||||
#if 0
|
||||
printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d "
|
||||
"R@%zd=%f R@1=%f cost=%f\n",
|
||||
knnTotalTime, linearTotalTime, totalHits, totalCnts, topk, recall,
|
||||
topk1Recall, cost);
|
||||
#endif
|
||||
EXPECT_GT(recall, 0.80f);
|
||||
EXPECT_GT(topk1Recall, 0.80f);
|
||||
// EXPECT_GT(cost, 2.0f);
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user