chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:42 +08:00
commit be3ef883e1
1214 changed files with 431743 additions and 0 deletions
@@ -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_flat_sparse
SRCS ${CC_SRCS}
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
)
endforeach()
@@ -0,0 +1,300 @@
// 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 "flat_sparse/flat_sparse_builder.h"
#include <future>
#include <iostream>
#include <vector>
#include <gtest/gtest.h>
#include "tests/test_util.h"
using namespace zvec::core;
using namespace zvec::ailego;
using namespace std;
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
class FlatSparseBuilderTest : 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 FlatSparseBuilderTest::_dir("FlatSparseBuilderTest/");
shared_ptr<IndexMeta> FlatSparseBuilderTest::_index_meta_ptr;
void FlatSparseBuilderTest::SetUp(void) {
_index_meta_ptr.reset(new (nothrow) IndexMeta(IndexMeta::MetaType::MT_SPARSE,
IndexMeta::DataType::DT_FP32));
_index_meta_ptr->set_metric("InnerProductSparse", 0, Params());
}
void FlatSparseBuilderTest::TearDown(void) {
zvec::test_util::RemoveTestPath(_dir);
}
TEST_F(FlatSparseBuilderTest, TestGeneral) {
IndexBuilder::Pointer builder =
IndexFactory::CreateBuilder("FlatSparseBuilder");
ASSERT_NE(builder, nullptr);
auto holder = make_shared<OnePassIndexSparseHolder<IndexMeta::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));
}
Params params;
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_EQ(stats.built_costtime(), 0UL);
// ASSERT_GT(stats.dumped_costtime(), 0UL);
// cleanup and rebuild
ASSERT_EQ(0, builder->cleanup());
auto holder2 = make_shared<MultiPassIndexSparseHolder<IndexMeta::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_EQ(stats.built_costtime(), 0UL);
}
TEST_F(FlatSparseBuilderTest, TestIndexThreads) {
IndexBuilder::Pointer builder1 =
IndexFactory::CreateBuilder("FlatSparseBuilder");
ASSERT_NE(builder1, nullptr);
IndexBuilder::Pointer builder2 =
IndexFactory::CreateBuilder("FlatSparseBuilder");
ASSERT_NE(builder2, nullptr);
auto holder = make_shared<MultiPassIndexSparseHolder<IndexMeta::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));
}
Params params;
std::srand(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(FlatSparseBuilderTest, TestHalfFloatConverter) {
IndexBuilder::Pointer builder =
IndexFactory::CreateBuilder("FlatSparseBuilder");
ASSERT_NE(builder, nullptr);
auto holder = make_shared<OnePassIndexSparseHolder<IndexMeta::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));
}
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();
Params params;
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_EQ(stats.built_costtime(), 0UL);
// ASSERT_GT(stats.dumped_costtime(), 0UL);
// cleanup and rebuild
ASSERT_EQ(0, builder->cleanup());
auto holder2 = make_shared<MultiPassIndexSparseHolder<IndexMeta::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_EQ(stats.built_costtime(), 0UL);
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
@@ -0,0 +1,819 @@
// 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 <future>
#include <iostream>
#include <random>
#include <vector>
#include <ailego/math/norm2_matrix.h>
#include <ailego/utility/math_helper.h>
#include <gtest/gtest.h>
#include <zvec/ailego/container/vector.h>
#include "tests/test_util.h"
#include "zvec/core/framework/index_factory.h"
#include "zvec/core/framework/index_meta.h"
using namespace zvec::core;
using namespace zvec::ailego;
using namespace std;
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
#endif
constexpr size_t static sparse_dim_count = 16;
class FlatSparseSearcherTest : public testing::Test {
protected:
void SetUp(void) override;
void TearDown(void) override;
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 std::shared_ptr<IndexMeta> index_meta_ptr_;
};
std::string FlatSparseSearcherTest::dir_("searcher_test/");
std::shared_ptr<IndexMeta> FlatSparseSearcherTest::index_meta_ptr_;
void FlatSparseSearcherTest::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;
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 FlatSparseSearcherTest::SetUp(void) {
IndexLoggerBroker::SetLevel(2);
index_meta_ptr_.reset(new IndexMeta(IndexMeta::MetaType::MT_SPARSE,
IndexMeta::DataType::DT_FP32));
index_meta_ptr_->set_metric("InnerProductSparse", 0, Params());
zvec::test_util::RemoveTestPath(dir_);
}
void FlatSparseSearcherTest::TearDown(void) {
zvec::test_util::RemoveTestPath(dir_);
}
TEST_F(FlatSparseSearcherTest, TestGeneral) {
// init storage
Params stg_params;
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_TRUE(storage != nullptr);
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "TestGeneral", true));
// init streamer
IndexStreamer::Pointer streamer =
IndexFactory::CreateStreamer("FlatSparseStreamer");
ASSERT_TRUE(streamer != nullptr);
IndexMeta index_meta(IndexMeta::MetaType::MT_SPARSE,
IndexMeta::DataType::DT_FP32);
index_meta.set_metric("InnerProductSparse", 0, Params());
Params params;
ASSERT_EQ(0, streamer->init(index_meta, params));
ASSERT_EQ(0, streamer->open(storage));
// generate sparse data
size_t sparse_dim_count = 32;
size_t cnt = 100U;
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);
// test add data
auto ctx = streamer->create_context();
ASSERT_TRUE(!!ctx);
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
for (size_t i = 0; i < cnt; i++) {
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count,
sparse_indices_list[i].data(),
sparse_vec_list[i].data(), qmeta, ctx));
}
// test get data
uint32_t sparse_count;
std::string sparse_indices_buffer;
std::string sparse_values_buffer;
for (size_t i = 0; i < cnt; i++) {
ASSERT_EQ(
0, streamer->get_sparse_vector(i, &sparse_count, &sparse_indices_buffer,
&sparse_values_buffer));
ASSERT_EQ(sparse_dim_count, sparse_count);
const uint32_t *sparse_indices_ptr =
reinterpret_cast<const uint32_t *>(sparse_indices_buffer.data());
const float *sparse_values_ptr =
reinterpret_cast<const float *>(sparse_values_buffer.data());
for (size_t j = 0; j < sparse_count; ++j) {
ASSERT_EQ(sparse_indices_ptr[j], sparse_indices_list[i][j]);
ASSERT_FLOAT_EQ(sparse_values_ptr[j], sparse_vec_list[i][j]);
// std::cout << "1: " << sparse_values_ptr[j]
// << " 2: " << sparse_vec_list[i][j] << std::endl;
}
// must clear ^_^
sparse_indices_buffer.clear();
sparse_values_buffer.clear();
}
// test dump
auto path = dir_ + "TestGeneral_dump";
auto dumper = IndexFactory::CreateDumper("FileDumper");
ASSERT_NE(dumper, nullptr);
ASSERT_EQ(0, dumper->create(path));
ASSERT_EQ(0, streamer->dump(dumper));
ASSERT_EQ(0, streamer->close());
ASSERT_EQ(0, dumper->close());
// do searcher get vector
IndexSearcher::Pointer searcher =
IndexFactory::CreateSearcher("FlatSparseSearcher");
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
ASSERT_TRUE(read_storage != nullptr);
ASSERT_TRUE(searcher != nullptr);
ASSERT_EQ(0, read_storage->open(path, false));
ASSERT_EQ(0, searcher->init(Params()));
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
// test searcher get data
for (size_t i = 0; i < cnt; i++) {
ASSERT_EQ(
0, searcher->get_sparse_vector(i, &sparse_count, &sparse_indices_buffer,
&sparse_values_buffer));
ASSERT_EQ(sparse_dim_count, sparse_count);
const uint32_t *sparse_indices_ptr =
reinterpret_cast<const uint32_t *>(sparse_indices_buffer.data());
const float *sparse_values_ptr =
reinterpret_cast<const float *>(sparse_values_buffer.data());
for (size_t j = 0; j < sparse_count; ++j) {
ASSERT_EQ(sparse_indices_ptr[j], sparse_indices_list[i][j]);
ASSERT_FLOAT_EQ(sparse_values_ptr[j], sparse_vec_list[i][j]);
// std::cout << "1: " << sparse_values_ptr[j]
// << " 2: " << sparse_vec_list[i][j] << std::endl;
}
// must clear ^_^
sparse_indices_buffer.clear();
sparse_values_buffer.clear();
}
}
TEST_F(FlatSparseSearcherTest, TestStreamerDump) {
IndexStreamer::Pointer streamer =
IndexFactory::CreateStreamer("FlatSparseStreamer");
ASSERT_NE(streamer, nullptr);
Params params;
Params stg_params;
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "TestStreamerDump.index", true));
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
ASSERT_EQ(0, streamer->open(storage));
size_t cnt = 10000U;
auto ctx = 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::DT_FP32);
for (size_t i = 0; i < cnt; i++) {
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count,
sparse_indices_list[i].data(),
sparse_vec_list[i].data(), qmeta, ctx));
}
auto path = dir_ + "TestStreamerDump";
auto dumper = IndexFactory::CreateDumper("FileDumper");
ASSERT_NE(dumper, nullptr);
ASSERT_EQ(0, dumper->create(path));
ASSERT_EQ(0, streamer->dump(dumper));
ASSERT_EQ(0, streamer->close());
ASSERT_EQ(0, dumper->close());
// do searcher knn
IndexSearcher::Pointer searcher =
IndexFactory::CreateSearcher("FlatSparseSearcher");
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
ASSERT_EQ(0, read_storage->open(path, false));
ASSERT_TRUE(searcher != nullptr);
ASSERT_EQ(0, searcher->init(Params()));
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
auto linearCtx = searcher->create_context();
auto knnCtx = searcher->create_context();
size_t topk = 200;
linearCtx->set_topk(topk);
knnCtx->set_topk(topk);
uint64_t knnTotalTime = 0;
uint64_t linearTotalTime = 0;
for (size_t i = 0; i < cnt; i += 50) {
const auto &sparse_indices = sparse_indices_list[i];
const auto &sparse_vec = sparse_vec_list[i];
auto t1 = Realtime::MicroSeconds();
ASSERT_EQ(0, searcher->search_impl(sparse_dim_count, sparse_indices.data(),
sparse_vec.data(), qmeta, knnCtx));
auto t2 = Realtime::MicroSeconds();
ASSERT_EQ(0,
searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
sparse_vec.data(), qmeta, linearCtx));
auto t3 = Realtime::MicroSeconds();
knnTotalTime += t2 - t1;
linearTotalTime += t3 - t2;
auto &knnResult = knnCtx->result();
auto &linearResult = linearCtx->result();
ASSERT_EQ(topk, linearResult.size());
ASSERT_EQ(topk, knnResult.size());
ASSERT_EQ(i, linearResult[0].key());
for (size_t k = 0; k < topk; ++k) {
ASSERT_EQ(linearResult[k].key(), knnResult[k].key());
}
}
printf("linear: %zu, knn: %zu\n", (size_t)linearTotalTime,
(size_t)knnTotalTime);
}
TEST_F(FlatSparseSearcherTest, TestLoadClose) {
Params stg_params;
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_TRUE(storage != nullptr);
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "TestGeneral", true));
// init streamer
IndexStreamer::Pointer streamer =
IndexFactory::CreateStreamer("FlatSparseStreamer");
ASSERT_TRUE(streamer != nullptr);
IndexMeta index_meta(IndexMeta::MetaType::MT_SPARSE,
IndexMeta::DataType::DT_FP32);
index_meta.set_metric("InnerProductSparse", 0, Params());
Params params;
ASSERT_EQ(0, streamer->init(index_meta, params));
ASSERT_EQ(0, streamer->open(storage));
// generate sparse data
size_t sparse_dim_count = 32;
size_t cnt = 100U;
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);
// test add data
auto ctx = streamer->create_context();
ASSERT_TRUE(!!ctx);
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
for (size_t i = 0; i < cnt; i++) {
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count,
sparse_indices_list[i].data(),
sparse_vec_list[i].data(), qmeta, ctx));
}
// test get data
uint32_t sparse_count;
std::string sparse_indices_buffer;
std::string sparse_values_buffer;
for (size_t i = 0; i < cnt; i++) {
ASSERT_EQ(
0, streamer->get_sparse_vector(i, &sparse_count, &sparse_indices_buffer,
&sparse_values_buffer));
ASSERT_EQ(sparse_dim_count, sparse_count);
const uint32_t *sparse_indices_ptr =
reinterpret_cast<const uint32_t *>(sparse_indices_buffer.data());
const float *sparse_values_ptr =
reinterpret_cast<const float *>(sparse_values_buffer.data());
for (size_t j = 0; j < sparse_count; ++j) {
ASSERT_EQ(sparse_indices_ptr[j], sparse_indices_list[i][j]);
ASSERT_FLOAT_EQ(sparse_values_ptr[j], sparse_vec_list[i][j]);
// std::cout << "1: " << sparse_values_ptr[j]
// << " 2: " << sparse_vec_list[i][j] << std::endl;
}
// must clear ^_^
sparse_indices_buffer.clear();
sparse_values_buffer.clear();
}
// test dump
auto path = dir_ + "TestGeneral_dump";
auto dumper = IndexFactory::CreateDumper("FileDumper");
ASSERT_NE(dumper, nullptr);
ASSERT_EQ(0, dumper->create(path));
ASSERT_EQ(0, streamer->dump(dumper));
ASSERT_EQ(0, streamer->close());
ASSERT_EQ(0, dumper->close());
// do searcher get vector
IndexSearcher::Pointer searcher =
IndexFactory::CreateSearcher("FlatSparseSearcher");
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
ASSERT_TRUE(read_storage != nullptr);
ASSERT_TRUE(searcher != nullptr);
ASSERT_EQ(0, read_storage->open(path, false));
ASSERT_EQ(0, searcher->init(Params()));
uint32_t loop = 5;
while (loop--) {
std::cout << "loop: " << loop << std::endl;
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
// test searcher get data
for (size_t i = 0; i < cnt; i++) {
ASSERT_EQ(0, searcher->get_sparse_vector(i, &sparse_count,
&sparse_indices_buffer,
&sparse_values_buffer));
ASSERT_EQ(sparse_dim_count, sparse_count);
const uint32_t *sparse_indices_ptr =
reinterpret_cast<const uint32_t *>(sparse_indices_buffer.data());
const float *sparse_values_ptr =
reinterpret_cast<const float *>(sparse_values_buffer.data());
for (size_t j = 0; j < sparse_count; ++j) {
ASSERT_EQ(sparse_indices_ptr[j], sparse_indices_list[i][j]);
ASSERT_FLOAT_EQ(sparse_values_ptr[j], sparse_vec_list[i][j]);
// std::cout << "1: " << sparse_values_ptr[j]
// << " 2: " << sparse_vec_list[i][j] << std::endl;
}
// must clear ^_^
sparse_indices_buffer.clear();
sparse_values_buffer.clear();
}
ASSERT_EQ(searcher->unload(), 0);
}
}
TEST_F(FlatSparseSearcherTest, TestSearch) {
IndexStreamer::Pointer streamer =
IndexFactory::CreateStreamer("FlatSparseStreamer");
ASSERT_TRUE(streamer != nullptr);
Params params;
Params stg_params;
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "TestLinearSearch.index", true));
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
ASSERT_EQ(0, streamer->open(storage));
size_t cnt = 5000UL;
auto ctx = streamer->create_context();
ASSERT_TRUE(!!ctx);
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
NumericalVector<float> sparse_velues(sparse_dim_count);
for (size_t i = 0; i < cnt; ++i) {
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = -1.0 * i - 1.0f;
}
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), qmeta, ctx));
}
// test dump
auto path = dir_ + "TestGeneral_dump";
auto dumper = IndexFactory::CreateDumper("FileDumper");
ASSERT_NE(dumper, nullptr);
ASSERT_EQ(0, dumper->create(path));
ASSERT_EQ(0, streamer->dump(dumper));
ASSERT_EQ(0, streamer->close());
ASSERT_EQ(0, dumper->close());
// do searcher get vector
IndexSearcher::Pointer searcher =
IndexFactory::CreateSearcher("FlatSparseSearcher");
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
ASSERT_TRUE(read_storage != nullptr);
ASSERT_TRUE(searcher != nullptr);
ASSERT_EQ(0, read_storage->open(path, false));
ASSERT_EQ(0, searcher->init(Params()));
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
size_t step = 50;
for (size_t i = 0; i < cnt; i += step) {
// std::cout << "search " << i << std::endl;
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = i + 1.0f;
}
ctx->set_topk(1U);
ASSERT_EQ(0,
searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(1UL, result1.size());
ASSERT_EQ(0, result1[0].key());
// std::cout << result1[0].key() << " " << result1[0].score() << std::endl;
ctx->set_topk(3U);
ASSERT_EQ(0,
searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(3UL, result2.size());
for (size_t i = 0; i < 3UL; ++i) {
// std::cout << result2[i].key() << " " << result2[i].score() <<
// std::endl;
ASSERT_EQ(i, result2[i].key());
}
}
ctx->set_topk(100U);
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = 10.1f;
}
ASSERT_EQ(0, searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), qmeta, ctx));
auto &result = ctx->result();
ASSERT_EQ(100U, result.size());
for (size_t i = 0; i < 100; ++i) {
ASSERT_EQ(i, result[i].key());
}
}
TEST_F(FlatSparseSearcherTest, TestSearchPKeys) {
IndexStreamer::Pointer streamer =
IndexFactory::CreateStreamer("FlatSparseStreamer");
ASSERT_TRUE(streamer != nullptr);
Params params;
Params stg_params;
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "TestLinearSearchByKeys.index", true));
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
ASSERT_EQ(0, streamer->open(storage));
size_t cnt = 5000UL;
auto ctx = streamer->create_context();
ASSERT_TRUE(!!ctx);
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
std::vector<std::vector<uint64_t>> p_keys;
p_keys.resize(1);
p_keys[0].resize(cnt);
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
NumericalVector<float> sparse_velues(sparse_dim_count);
for (size_t i = 0; i < cnt; ++i) {
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = -1.0 * i - 1.0f;
}
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), qmeta, ctx));
p_keys[0][i] = i;
}
// test dump
auto path = dir_ + "TestGeneral_dump";
auto dumper = IndexFactory::CreateDumper("FileDumper");
ASSERT_NE(dumper, nullptr);
ASSERT_EQ(0, dumper->create(path));
ASSERT_EQ(0, streamer->dump(dumper));
ASSERT_EQ(0, streamer->close());
ASSERT_EQ(0, dumper->close());
// do searcher get vector
IndexSearcher::Pointer searcher =
IndexFactory::CreateSearcher("FlatSparseSearcher");
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
ASSERT_TRUE(read_storage != nullptr);
ASSERT_TRUE(searcher != nullptr);
ASSERT_EQ(0, read_storage->open(path, false));
ASSERT_EQ(0, searcher->init(Params()));
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
size_t topk = 3;
for (size_t i = 0; i < cnt; i += 50) {
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = i + 1.0f;
}
ctx->set_topk(1U);
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), p_keys, qmeta, ctx));
auto &result1 = ctx->result();
ASSERT_EQ(1UL, result1.size());
ASSERT_EQ(0, result1[0].key());
ctx->set_topk(topk);
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), p_keys, qmeta, ctx));
auto &result2 = ctx->result();
ASSERT_EQ(topk, result2.size());
ASSERT_EQ(0, result2[0].key());
ASSERT_EQ(1, result2[1].key());
ASSERT_EQ(2, result2[2].key());
}
{
ctx->set_topk(100U);
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = 1.0f;
}
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), p_keys, qmeta, ctx));
auto &result = ctx->result();
ASSERT_EQ(100U, result.size());
ASSERT_EQ(0, result[0].key());
ASSERT_EQ(1, result[1].key());
ASSERT_EQ(10, result[10].key());
ASSERT_EQ(20, result[20].key());
ASSERT_EQ(30, result[30].key());
ASSERT_EQ(35, result[35].key());
ASSERT_EQ(99, result[99].key());
}
{
ctx->set_topk(100U);
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = 10.0f;
}
p_keys[0] = {{cnt + 1, 10, 1, 15, cnt + 2}};
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), p_keys, qmeta, ctx));
auto &result = ctx->result();
ASSERT_EQ(3U, result.size());
ASSERT_EQ(1, result[0].key());
ASSERT_EQ(10, result[1].key());
ASSERT_EQ(15, result[2].key());
}
{
ctx->set_topk(100U);
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = 9.0f;
}
p_keys[0].clear();
for (size_t j = 0; j < cnt; j += 10) {
p_keys[0].push_back((uint64_t)j);
}
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), p_keys, qmeta, ctx));
auto &result = ctx->result();
ASSERT_EQ(100U, result.size());
ASSERT_EQ(0, result[0].key());
ASSERT_EQ(10, result[1].key());
ASSERT_EQ(100, result[10].key());
ASSERT_EQ(200, result[20].key());
ASSERT_EQ(300, result[30].key());
ASSERT_EQ(350, result[35].key());
ASSERT_EQ(990, result[99].key());
}
}
TEST_F(FlatSparseSearcherTest, TestMultiThread) {
IndexStreamer::Pointer streamer =
IndexFactory::CreateStreamer("FlatSparseStreamer");
ASSERT_TRUE(streamer != nullptr);
Params params;
constexpr size_t static sparse_dim_count = 32;
IndexMeta meta(IndexMeta::MetaType::MT_SPARSE, IndexMeta::DataType::DT_FP32);
meta.set_metric("InnerProductSparse", 0, Params());
ASSERT_EQ(0, streamer->init(meta, params));
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
ASSERT_NE(nullptr, storage);
Params stg_params;
ASSERT_EQ(0, storage->init(stg_params));
ASSERT_EQ(0, storage->open(dir_ + "TessKnnMultiThread", true));
ASSERT_EQ(0, streamer->open(storage));
auto addVector = [&streamer](int baseKey, size_t addCnt) {
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
size_t succAdd = 0;
auto ctx = streamer->create_context();
for (size_t i = 0; i < addCnt; i++) {
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
NumericalVector<float> sparse_velues(sparse_dim_count);
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = (float)i + baseKey;
}
succAdd += !streamer->add_impl(baseKey + i, sparse_dim_count,
sparse_indices.data(),
sparse_velues.data(), qmeta, ctx);
}
streamer->flush(0UL);
return succAdd;
};
auto t2 = std::async(std::launch::async, addVector, 1000, 1000);
auto t3 = std::async(std::launch::async, addVector, 2000, 1000);
auto t1 = std::async(std::launch::async, addVector, 0, 1000);
ASSERT_EQ(1000U, t1.get());
ASSERT_EQ(1000U, t2.get());
ASSERT_EQ(1000U, t3.get());
streamer->close();
// checking data
ASSERT_EQ(0, streamer->open(storage));
auto provider = streamer->create_sparse_provider();
auto iter = provider->create_iterator();
ASSERT_TRUE(!!iter);
size_t total = 0;
uint64_t min = 1000;
uint64_t max = 0;
std::set<uint64_t> keys;
while (iter->is_valid()) {
const uint32_t sparse_count = iter->sparse_count();
ASSERT_EQ(sparse_count, sparse_dim_count);
const float *data = reinterpret_cast<const float *>(iter->sparse_data());
for (size_t j = 0; j < sparse_dim_count; ++j) {
ASSERT_EQ((float)iter->key(), data[j]);
}
total++;
min = std::min(min, iter->key());
max = std::max(max, iter->key());
keys.insert(iter->key());
iter->next();
}
ASSERT_EQ(3000, keys.size());
ASSERT_EQ(3000, total);
ASSERT_EQ(0, min);
ASSERT_EQ(2999, max);
// test dump
auto path = dir_ + "TestGeneral_dump";
auto dumper = IndexFactory::CreateDumper("FileDumper");
ASSERT_NE(dumper, nullptr);
ASSERT_EQ(0, dumper->create(path));
ASSERT_EQ(0, streamer->dump(dumper));
ASSERT_EQ(0, streamer->close());
ASSERT_EQ(0, dumper->close());
// do searcher get vector
IndexSearcher::Pointer searcher =
IndexFactory::CreateSearcher("FlatSparseSearcher");
auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage");
ASSERT_TRUE(read_storage != nullptr);
ASSERT_TRUE(searcher != nullptr);
ASSERT_EQ(0, read_storage->open(path, false));
ASSERT_EQ(0, searcher->init(Params()));
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
// ====== multi thread search
size_t topk = 10;
size_t cnt = 3000;
auto knnSearch = [&]() {
auto linearCtx = searcher->create_context();
auto linearByPkeysCtx = searcher->create_context();
auto ctx = searcher->create_context();
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
linearCtx->set_topk(topk);
linearByPkeysCtx->set_topk(topk);
ctx->set_topk(topk);
size_t totalCnts = 0;
size_t totalHits = 0;
for (size_t i = 0; i < cnt; i += 1) {
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
NumericalVector<float> sparse_velues(sparse_dim_count);
for (size_t j = 0; j < sparse_dim_count; ++j) {
sparse_indices[j] = j * 20;
sparse_velues[j] = ((float)i + 1.1f);
}
ASSERT_EQ(0,
searcher->search_impl(sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), qmeta, ctx));
ASSERT_EQ(
0, searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), qmeta, linearCtx));
std::vector<std::vector<uint64_t>> p_keys = {{cnt - 1, cnt - 2, cnt - 3}};
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
sparse_dim_count, sparse_indices.data(),
sparse_velues.data(), p_keys, qmeta, linearByPkeysCtx));
auto &r1 = ctx->result();
ASSERT_EQ(topk, r1.size());
// std::cout << "r1 top1: " << r1[0].key() << ", score: " << r1[0].score()
// << std::endl;
ASSERT_EQ(cnt - 1, r1[0].key());
auto &r2 = linearCtx->result();
ASSERT_EQ(topk, r2.size());
// std::cout << "r2 top1: " << r2[0].key() << ", score: " << r2[0].score()
// << std::endl;
ASSERT_EQ(cnt - 1, r2[0].key());
auto &r3 = linearByPkeysCtx->result();
ASSERT_EQ(std::min(topk, p_keys[0].size()), r3.size());
#if 0
printf("linear: %zd => %zd %zd %zd %zd %zd\n", i, r2[0].key,
r2[1].key, r2[2].key, r2[3].key, r2[4].key);
printf("knn: %zd => %zd %zd %zd %zd %zd\n", i, r1[0].key, r1[1].key,
r1[2].key, r1[3].key, r1[4].key);
#endif
for (size_t k = 0; k < topk; ++k) {
totalCnts++;
for (size_t j = 0; j < topk; ++j) {
if (r2[j].key() == r1[k].key()) {
totalHits++;
break;
}
}
}
}
printf("%f\n", totalHits * 1.0f / totalCnts);
ASSERT_FLOAT_EQ(1.0f, totalHits * 1.0f / totalCnts);
};
auto s1 = std::async(std::launch::async, knnSearch);
auto s2 = std::async(std::launch::async, knnSearch);
auto s3 = std::async(std::launch::async, knnSearch);
s1.wait();
s2.wait();
s3.wait();
}
#if defined(__GNUC__) || defined(__GNUG__)
#pragma GCC diagnostic pop
#endif
@@ -0,0 +1,230 @@
// 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 <string>
#include <vector>
#include <ailego/math/norm2_matrix.h>
#include <ailego/utility/math_helper.h>
#include <ailego/utility/memory_helper.h>
#include <algorithm/flat_sparse/flat_sparse_utility.h>
#include <gtest/gtest.h>
#include <zvec/core/framework/index_framework.h>
#include <zvec/core/framework/index_streamer.h>
#include "tests/test_util.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 FlatSparseStreamerTest : public testing::Test {
protected:
void SetUp(void) override;
void TearDown(void) override;
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 FlatSparseStreamerTest::dir_("FlatSparseStreamerTest/");
shared_ptr<IndexMeta> FlatSparseStreamerTest::index_meta_ptr_;
void FlatSparseStreamerTest::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 FlatSparseStreamerTest::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 FlatSparseStreamerTest::TearDown(void) {
zvec::test_util::RemoveTestPath(dir_);
}
TEST_F(FlatSparseStreamerTest, TestGeneral) {
IndexStreamer::Pointer write_streamer =
IndexFactory::CreateStreamer("FlatSparseStreamer");
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;
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/FlatSparseSearch", 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();
write_storage.reset();
IndexStreamer::Pointer read_streamer =
IndexFactory::CreateStreamer("FlatSparseStreamer");
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/FlatSparseSearch", 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;
std::cout << "knnTotalTime=" << knnTotalTime
<< " linearTotalTime=" << linearTotalTime << std::endl;
#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