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_flat
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,334 @@
|
||||
// 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/flat_builder.h"
|
||||
#include <future>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <gtest/gtest.h>
|
||||
#include "tests/test_util.h"
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
static inline size_t RandomDimension(void) {
|
||||
std::mt19937 gen((std::random_device())());
|
||||
return (std::uniform_int_distribution<size_t>(1, 129))(gen);
|
||||
}
|
||||
|
||||
static size_t DIMENSION = RandomDimension();
|
||||
class FlatBuilderTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
|
||||
public:
|
||||
static std::string dir_;
|
||||
static IndexMeta meta_;
|
||||
};
|
||||
|
||||
std::string FlatBuilderTest ::dir_("flat_builder_test/");
|
||||
IndexMeta FlatBuilderTest::meta_;
|
||||
|
||||
void FlatBuilderTest::SetUp(void) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
}
|
||||
|
||||
//! self-check column-major and row-major search.
|
||||
void FlatBuilderTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void build_process(IndexBuilder::Pointer &builder,
|
||||
IndexHolder::Pointer holder) {
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(FlatBuilderTest::meta_, params));
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
std::string path = FlatBuilderTest::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(0UL, stats.discarded_count());
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInitSuccess) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInitFailedWithInvalidMeasure) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("invalid", 0, Params());
|
||||
Params params;
|
||||
int ret = builder->init(meta_, params);
|
||||
EXPECT_EQ(IndexError_InvalidArgument, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8InvalidColumnMajor) {
|
||||
size_t dim = (DIMENSION + 3) / 4 * 4;
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, dim + 2);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
ASSERT_EQ(IndexMeta::MO_COLUMN, meta_.major_order());
|
||||
Params params;
|
||||
ASSERT_NE(0, builder->init(meta_, params));
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8WithRandomDimension) {
|
||||
size_t dim = DIMENSION;
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, dim);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_UNDEFINED);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestBuildWithRowMajor) {
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_ROW);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_FP32>>(DIMENSION);
|
||||
size_t doc_cnt = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
int ret = builder->train(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder->build(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8BuildWithRowMajor) {
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_meta(IndexMeta::DT_INT8, DIMENSION);
|
||||
meta_.set_major_order(IndexMeta::MO_ROW);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_INT8>>(DIMENSION);
|
||||
size_t doc_cnt = 128UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<int8_t> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = (int8_t)(i % 128);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
int ret = builder->train(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder->build(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestBuildWithColumnMajor) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_FP32>>(DIMENSION);
|
||||
size_t doc_cnt = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
int ret = builder->train(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder->build(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8BuildWithColumnMajor) {
|
||||
size_t dim = (DIMENSION + 3) / 4 * 4;
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, dim);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder = std::make_shared<OnePassIndexHolder<IndexMeta::DT_INT8>>(dim);
|
||||
size_t doc_cnt = 128UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<int8_t> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = (int8_t)(i % 128);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
int ret = builder->train(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder->build(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestWithRowMajor) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_ROW);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_FP32>>(DIMENSION);
|
||||
size_t doc_cnt = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
build_process(builder, holder);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8WithRowMajor) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_ROW);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_INT8>>(DIMENSION);
|
||||
size_t doc_cnt = 128UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<int8_t> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = (int8_t)(i % 128);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
build_process(builder, holder);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestWithColumnMajor) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_FP32>>(DIMENSION);
|
||||
size_t doc_cnt = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
build_process(builder, holder);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8WithColumnMajor) {
|
||||
size_t dim = (DIMENSION + 3) / 4 * 4;
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, dim);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder = std::make_shared<OnePassIndexHolder<IndexMeta::DT_INT8>>(dim);
|
||||
size_t doc_cnt = 128UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<int8_t> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = (int8_t)(i % 128);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
build_process(builder, holder);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
}
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,623 @@
|
||||
#include <future>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <ailego/utility/math_helper.h>
|
||||
#include <ailego/utility/memory_helper.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 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 dim = 16;
|
||||
|
||||
class FlatStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
void hybrid_scale(std::vector<float> &dense_value,
|
||||
std::vector<float> &sparse_value, float alpha_scale);
|
||||
|
||||
static std::string dir_;
|
||||
static std::shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string FlatStreamerTest::dir_("flat_streamer_buffer_test_dir/");
|
||||
std::shared_ptr<IndexMeta> FlatStreamerTest::index_meta_ptr_;
|
||||
|
||||
void FlatStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (std::nothrow)
|
||||
IndexMeta(IndexMeta::DataType::DT_FP32, dim));
|
||||
index_meta_ptr_->set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void FlatStreamerTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearch) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, 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_ + "Test/LinearSearch", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
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/LinearSearch", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchBuffer) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/LinearSearchBuffer", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
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/LinearSearchBuffer", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchBufferMMap) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/LinearSearchBuffer", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
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/LinearSearchBuffer", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchWithLRU) {
|
||||
MemoryLimitPool::get_instance().init(100 * 1024UL * 1024UL);
|
||||
#ifdef __ANDROID__
|
||||
GTEST_SKIP()
|
||||
<< "Skipped on Android: requires ~6GB memory/disk (emulator limit)";
|
||||
#endif
|
||||
constexpr size_t static dim = 1600;
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
IndexMeta meta = IndexMeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
meta.set_metric("SquaredEuclidean", 0, Params());
|
||||
ASSERT_EQ(0, write_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_ + "/Test/LinearSearchWithLRU", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 50000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(meta, 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/LinearSearchWithLRU", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < 10; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchMMap) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, 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_ + "Test/LinearSearchMMap", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
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/LinearSearchMMap", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
const float *data = (float *)provider->get_vector(result1[0].key());
|
||||
EXPECT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
}
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
@@ -0,0 +1,231 @@
|
||||
#include <future>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <ailego/utility/math_helper.h>
|
||||
#include <ailego/utility/memory_helper.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/utility/file_helper.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.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 dim = 128;
|
||||
|
||||
class FlatStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
void hybrid_scale(std::vector<float> &dense_value,
|
||||
std::vector<float> &sparse_value, float alpha_scale);
|
||||
|
||||
static std::string dir_;
|
||||
static std::shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string FlatStreamerTest::dir_("streamer_test/");
|
||||
std::shared_ptr<IndexMeta> FlatStreamerTest::index_meta_ptr_;
|
||||
|
||||
void FlatStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (std::nothrow)
|
||||
IndexMeta(IndexMeta::DataType::DT_FP32, dim));
|
||||
index_meta_ptr_->set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
zvec::ailego::FileHelper::RemovePath(dir_.c_str());
|
||||
}
|
||||
|
||||
void FlatStreamerTest::TearDown(void) {
|
||||
zvec::ailego::FileHelper::RemovePath(dir_.c_str());
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchMMap) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, 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_ + "/Test/LinearSearchMMap", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t data_cnt = 300000UL, cnt = 500UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < data_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
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/LinearSearchMMap", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 30;
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result1 = ctx->result();
|
||||
// ASSERT_EQ(topk, result1.size());
|
||||
// ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// vec[j] = i + 0.1f;
|
||||
// }
|
||||
// ctx->set_topk(topk);
|
||||
// ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result2 = ctx->result();
|
||||
// ASSERT_EQ(topk, result2.size());
|
||||
// ASSERT_EQ(i, result2[0].key());
|
||||
// ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
// ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.micro_seconds() << " us" << endl;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result1 = ctx->result();
|
||||
// ASSERT_EQ(topk, result1.size());
|
||||
// ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// vec[j] = i + 0.1f;
|
||||
// }
|
||||
// ctx->set_topk(topk);
|
||||
// ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result2 = ctx->result();
|
||||
// ASSERT_EQ(topk, result2.size());
|
||||
// ASSERT_EQ(i, result2[0].key());
|
||||
// ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
// ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.micro_seconds() << " us" << endl;
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchBuffer) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, 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_ + "/Test/LinearSearchBuffer", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t data_cnt = 300000UL, cnt = 500UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < data_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
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/LinearSearchBuffer", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 30;
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result1 = ctx->result();
|
||||
// ASSERT_EQ(topk, result1.size());
|
||||
// ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// vec[j] = i + 0.1f;
|
||||
// }
|
||||
// ctx->set_topk(topk);
|
||||
// ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result2 = ctx->result();
|
||||
// ASSERT_EQ(topk, result2.size());
|
||||
// ASSERT_EQ(i, result2[0].key());
|
||||
// ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
// ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.micro_seconds() << " us" << endl;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result1 = ctx->result();
|
||||
// ASSERT_EQ(topk, result1.size());
|
||||
// ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// vec[j] = i + 0.1f;
|
||||
// }
|
||||
// ctx->set_topk(topk);
|
||||
// ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result2 = ctx->result();
|
||||
// ASSERT_EQ(topk, result2.size());
|
||||
// ASSERT_EQ(i, result2[0].key());
|
||||
// ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
// ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.micro_seconds() << " us" << endl;
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
#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