Skip to content

Commit

Permalink
[fix][test] Fixup vector index unit test.
Browse files Browse the repository at this point in the history
  • Loading branch information
@rock-git @ketor
rock-git authored and ketor committed Mar 20, 2024
1 parent fbe9787 commit b713ef2
Show file tree
Hide file tree
Showing 15 changed files with 224 additions and 189 deletions.
32 changes: 19 additions & 13 deletions src/vector/vector_index.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -202,23 +202,29 @@ butil::Status VectorIndex::RangeSearchByParallel(
const std::vector<pb::common::VectorWithId>& vector_with_ids, float radius,
const std::vector<std::shared_ptr<VectorIndex::FilterFunctor>>& filters, bool reconstruct,
const pb::common::VectorSearchParameter& parameter, std::vector<pb::index::VectorWithDistanceResult>& results) {
results.resize(vector_with_ids.size());
if (VectorIndexType() == pb::common::VECTOR_INDEX_TYPE_HNSW) {
// parallel in inner
return RangeSearch(vector_with_ids, radius, filters, reconstruct, parameter, results);
} else {
// parallel in here
results.resize(vector_with_ids.size());

std::vector<std::vector<pb::common::VectorWithId>> vector_with_id_batchs;
SplitVectorWithId(vector_with_ids, FLAGS_vector_read_batch_size_per_task, vector_with_id_batchs);
std::vector<std::vector<pb::common::VectorWithId>> vector_with_id_batchs;
SplitVectorWithId(vector_with_ids, FLAGS_vector_read_batch_size_per_task, vector_with_id_batchs);

return ParallelRun(
thread_pool, vector_with_id_batchs, true,
[&](const std::vector<pb::common::VectorWithId>& vector_with_ids, uint32_t index) -> butil::Status {
std::vector<pb::index::VectorWithDistanceResult> part_results;
auto status = RangeSearch(vector_with_ids, radius, filters, reconstruct, parameter, part_results);
return ParallelRun(
thread_pool, vector_with_id_batchs, true,
[&](const std::vector<pb::common::VectorWithId>& vector_with_ids, uint32_t index) -> butil::Status {
std::vector<pb::index::VectorWithDistanceResult> part_results;
auto status = RangeSearch(vector_with_ids, radius, filters, reconstruct, parameter, part_results);

for (int i = 0; i < part_results.size(); ++i) {
results[index + i].Swap(&part_results[i]);
}
for (int i = 0; i < part_results.size(); ++i) {
results[index + i].Swap(&part_results[i]);
}

return status;
});
return status;
});
}
}

butil::Status VectorIndex::Save(const std::string& /*path*/) {
Expand Down
2 changes: 1 addition & 1 deletion src/vector/vector_index_flat.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -296,7 +296,7 @@ butil::Status VectorIndexFlat::Load(const std::string& path) {
// The outside has been locked. Remove the locking operation here.
faiss::Index* internal_raw_index = nullptr;
try {
faiss::Index* internal_raw_index = faiss::read_index(path.c_str(), 0);
internal_raw_index = faiss::read_index(path.c_str(), 0);
} catch (std::exception& e) {
std::string s = fmt::format("VectorIndexFlat::Load faiss::read_index failed. path : {} error : {}", path, e.what());
DINGO_LOG(ERROR) << s;
Expand Down
7 changes: 0 additions & 7 deletions test/unit_test/test_txn_gc.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,13 +113,6 @@ static void PrepareData(const std::vector<std::string> &prefix_key_array, int st
class TxnGcTest : public testing::Test {
protected:
static void SetUpTestSuite() {
DingoLogger::ChangeGlogLevelUsingDingoLevel(dingodb::pb::node::LogLevel::DEBUG, 0);

// Set whether log messages go to stderr in addition to logfiles.
FLAGS_alsologtostderr = true;

// If set this flag to true, the log will show in the terminal
FLAGS_logtostderr = true;
std::srand(std::time(nullptr));

Helper::CreateDirectories(kStorePath);
Expand Down
15 changes: 11 additions & 4 deletions test/unit_test/test_vector_index.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,15 +30,19 @@ namespace dingodb {

class VectorIndexWrapperTest : public testing::Test {
protected:
static void SetUpTestSuite() {}
static void SetUpTestSuite() { vector_index_thread_pool = std::make_shared<ThreadPool>("vector_index", 4); }

static void TearDownTestSuite() {}

void SetUp() override {}

void TearDown() override {}

static ThreadPoolPtr vector_index_thread_pool;
};

ThreadPoolPtr VectorIndexWrapperTest::vector_index_thread_pool = nullptr;

pb::common::RegionEpoch GenEpoch(int version) {
pb::common::RegionEpoch epoch;
epoch.set_conf_version(1);
Expand Down Expand Up @@ -89,7 +93,8 @@ TEST_F(VectorIndexWrapperTest, Add) {
index_parameter.mutable_hnsw_parameter()->set_max_elements(100000);
index_parameter.mutable_hnsw_parameter()->set_nlinks(2);

auto vector_index_hnsw = VectorIndexFactory::New(id, index_parameter, GenEpoch(10), GenRange(1, 1000));
auto vector_index_hnsw =
VectorIndexFactory::NewHnsw(id, index_parameter, GenEpoch(10), GenRange(1, 1000), vector_index_thread_pool);

auto vector_index_wrapper = VectorIndexWrapper::New(id, index_parameter);

Expand All @@ -111,8 +116,10 @@ TEST_F(VectorIndexWrapperTest, Add_Sibling) {
index_parameter.mutable_hnsw_parameter()->set_max_elements(100000);
index_parameter.mutable_hnsw_parameter()->set_nlinks(2);

auto vector_index = VectorIndexFactory::New(id, index_parameter, GenEpoch(10), GenRange(1, 1000));
auto sibling_vector_index = VectorIndexFactory::New(id, index_parameter, GenEpoch(10), GenRange(1000, 2000));
auto vector_index =
VectorIndexFactory::NewHnsw(id, index_parameter, GenEpoch(10), GenRange(1, 1000), vector_index_thread_pool);
auto sibling_vector_index =
VectorIndexFactory::NewHnsw(id, index_parameter, GenEpoch(10), GenRange(1000, 2000), vector_index_thread_pool);

auto vector_index_wrapper = VectorIndexWrapper::New(id, index_parameter);

Expand Down
61 changes: 15 additions & 46 deletions test/unit_test/test_vector_index_flat.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,10 @@ static const std::string kTempDataDirectory = "./unit_test/vector_index_flat";

class VectorIndexFlatTest : public testing::Test {
protected:
static void SetUpTestSuite() { Helper::CreateDirectories(kTempDataDirectory); }
static void SetUpTestSuite() {
Helper::CreateDirectories(kTempDataDirectory);
vector_index_thread_pool = std::make_shared<ThreadPool>("vector_index", 4);
}

static void TearDownTestSuite() {
vector_index_flat_l2.reset();
Expand All @@ -59,44 +62,24 @@ class VectorIndexFlatTest : public testing::Test {
inline static faiss::idx_t dimension = 8;
inline static int data_base_size = 10;
inline static std::vector<float> data_base;

static ThreadPoolPtr vector_index_thread_pool;
};

ThreadPoolPtr VectorIndexFlatTest::vector_index_thread_pool = nullptr;

TEST_F(VectorIndexFlatTest, Create) {
static const pb::common::Range kRange;
static pb::common::RegionEpoch kEpoch; // NOLINT
kEpoch.set_conf_version(1);
kEpoch.set_version(10);
// invalid param
{
int64_t id = 1;
pb::common::VectorIndexParameter index_parameter;
vector_index_flat_l2 = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
EXPECT_EQ(vector_index_flat_l2.get(), nullptr);
}

// invalid param
{
int64_t id = 1;
pb::common::VectorIndexParameter index_parameter;
vector_index_flat_l2 = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
EXPECT_EQ(vector_index_flat_l2.get(), nullptr);
}

// invalid param
{
int64_t id = 1;
pb::common::VectorIndexParameter index_parameter;
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_NONE);
vector_index_flat_l2 = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
EXPECT_EQ(vector_index_flat_l2.get(), nullptr);
}

// invalid param
{
int64_t id = 1;
pb::common::VectorIndexParameter index_parameter;
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
vector_index_flat_l2 = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_l2 = VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_EQ(vector_index_flat_l2.get(), nullptr);
}

Expand All @@ -106,7 +89,7 @@ TEST_F(VectorIndexFlatTest, Create) {
pb::common::VectorIndexParameter index_parameter;
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
index_parameter.mutable_flat_parameter()->set_dimension(64);
vector_index_flat_l2 = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_l2 = VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_EQ(vector_index_flat_l2.get(), nullptr);
}

Expand All @@ -117,7 +100,7 @@ TEST_F(VectorIndexFlatTest, Create) {
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
index_parameter.mutable_flat_parameter()->set_dimension(64);
index_parameter.mutable_flat_parameter()->set_metric_type(::dingodb::pb::common::MetricType::METRIC_TYPE_NONE);
vector_index_flat_l2 = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_l2 = VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_EQ(vector_index_flat_l2.get(), nullptr);
}

Expand All @@ -128,7 +111,7 @@ TEST_F(VectorIndexFlatTest, Create) {
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
index_parameter.mutable_flat_parameter()->set_dimension(dimension);
index_parameter.mutable_flat_parameter()->set_metric_type(::dingodb::pb::common::MetricType::METRIC_TYPE_L2);
vector_index_flat_l2 = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_l2 = VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_NE(vector_index_flat_l2.get(), nullptr);
}

Expand All @@ -140,7 +123,7 @@ TEST_F(VectorIndexFlatTest, Create) {
index_parameter.mutable_flat_parameter()->set_dimension(dimension);
index_parameter.mutable_flat_parameter()->set_metric_type(
::dingodb::pb::common::MetricType::METRIC_TYPE_INNER_PRODUCT);
vector_index_flat_ip = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_ip = VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_NE(vector_index_flat_ip.get(), nullptr);
}

Expand All @@ -151,7 +134,8 @@ TEST_F(VectorIndexFlatTest, Create) {
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
index_parameter.mutable_flat_parameter()->set_dimension(dimension);
index_parameter.mutable_flat_parameter()->set_metric_type(::dingodb::pb::common::MetricType::METRIC_TYPE_COSINE);
vector_index_flat_cosine = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_cosine =
VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_NE(vector_index_flat_cosine.get(), nullptr);
}
}
Expand Down Expand Up @@ -736,21 +720,6 @@ TEST_F(VectorIndexFlatTest, Upsert) {
EXPECT_EQ(ok.error_code(), pb::error::Errno::OK);
}

// // must be delete .................................................................
// {
// vector_index_flat_.reset();

// int64_t id = 1;
// pb::common::VectorIndexParameter index_parameter;
// index_parameter.set_vector_index_type(
// ::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
// index_parameter.mutable_flat_parameter()->set_dimension(dimension_);
// index_parameter.mutable_flat_parameter()->set_metric_type(
// ::dingodb::pb::common::MetricType::METRIC_TYPE_L2);
// vector_index_flat_ = VectorIndexFactory::New(id, index_parameter);
// EXPECT_NE(vector_index_flat_.get(), nullptr);
// }

// empty OK
{
std::vector<pb::common::VectorWithId> vector_with_ids;
Expand Down
15 changes: 11 additions & 4 deletions test/unit_test/test_vector_index_flat_search_limit.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ namespace dingodb {

class VectorIndexFlatSearchParamLimitTest : public testing::Test {
protected:
static void SetUpTestSuite() {}
static void SetUpTestSuite() { vector_index_thread_pool = std::make_shared<ThreadPool>("vector_index", 4); }

static void TearDownTestSuite() {
vector_index_flat_for_l2.reset();
Expand Down Expand Up @@ -67,8 +67,12 @@ class VectorIndexFlatSearchParamLimitTest : public testing::Test {
inline static int vector_ids_search_size = 10;

inline static int search_topk = 1;

static ThreadPoolPtr vector_index_thread_pool;
};

ThreadPoolPtr VectorIndexFlatSearchParamLimitTest::vector_index_thread_pool = nullptr;

TEST_F(VectorIndexFlatSearchParamLimitTest, Create) {
static const pb::common::Range kRange;
static pb::common::RegionEpoch kEpoch; // NOLINT
Expand All @@ -82,7 +86,8 @@ TEST_F(VectorIndexFlatSearchParamLimitTest, Create) {
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
index_parameter.mutable_flat_parameter()->set_dimension(kDimension);
index_parameter.mutable_flat_parameter()->set_metric_type(::dingodb::pb::common::MetricType::METRIC_TYPE_L2);
vector_index_flat_for_l2 = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_for_l2 =
VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_NE(vector_index_flat_for_l2.get(), nullptr);
}

Expand All @@ -94,7 +99,8 @@ TEST_F(VectorIndexFlatSearchParamLimitTest, Create) {
index_parameter.mutable_flat_parameter()->set_dimension(kDimension);
index_parameter.mutable_flat_parameter()->set_metric_type(
::dingodb::pb::common::MetricType::METRIC_TYPE_INNER_PRODUCT);
vector_index_flat_for_ip = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_for_ip =
VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_NE(vector_index_flat_for_ip.get(), nullptr);
}

Expand All @@ -105,7 +111,8 @@ TEST_F(VectorIndexFlatSearchParamLimitTest, Create) {
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
index_parameter.mutable_flat_parameter()->set_dimension(kDimension);
index_parameter.mutable_flat_parameter()->set_metric_type(::dingodb::pb::common::MetricType::METRIC_TYPE_COSINE);
vector_index_flat_for_cosine = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_for_cosine =
VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_NE(vector_index_flat_for_cosine.get(), nullptr);
}
}
Expand Down
15 changes: 11 additions & 4 deletions test/unit_test/test_vector_index_flat_search_param.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ namespace dingodb {

class VectorIndexFlatSearchParamTest : public testing::Test {
protected:
static void SetUpTestSuite() {}
static void SetUpTestSuite() { vector_index_thread_pool = std::make_shared<ThreadPool>("vector_index", 4); }

static void TearDownTestSuite() {
vector_index_flat_for_l2.reset();
Expand Down Expand Up @@ -66,8 +66,12 @@ class VectorIndexFlatSearchParamTest : public testing::Test {
inline static int vector_ids_search_size = 10;

inline static int search_topk = 3;

static ThreadPoolPtr vector_index_thread_pool;
};

ThreadPoolPtr VectorIndexFlatSearchParamTest::vector_index_thread_pool = nullptr;

TEST_F(VectorIndexFlatSearchParamTest, Create) {
static const pb::common::Range kRange;
static pb::common::RegionEpoch kEpoch; // NOLINT
Expand All @@ -81,7 +85,8 @@ TEST_F(VectorIndexFlatSearchParamTest, Create) {
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
index_parameter.mutable_flat_parameter()->set_dimension(dimension);
index_parameter.mutable_flat_parameter()->set_metric_type(::dingodb::pb::common::MetricType::METRIC_TYPE_L2);
vector_index_flat_for_l2 = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_for_l2 =
VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_NE(vector_index_flat_for_l2.get(), nullptr);
}

Expand All @@ -93,7 +98,8 @@ TEST_F(VectorIndexFlatSearchParamTest, Create) {
index_parameter.mutable_flat_parameter()->set_dimension(dimension);
index_parameter.mutable_flat_parameter()->set_metric_type(
::dingodb::pb::common::MetricType::METRIC_TYPE_INNER_PRODUCT);
vector_index_flat_for_ip = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_for_ip =
VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_NE(vector_index_flat_for_ip.get(), nullptr);
}

Expand All @@ -104,7 +110,8 @@ TEST_F(VectorIndexFlatSearchParamTest, Create) {
index_parameter.set_vector_index_type(::dingodb::pb::common::VectorIndexType::VECTOR_INDEX_TYPE_FLAT);
index_parameter.mutable_flat_parameter()->set_dimension(dimension);
index_parameter.mutable_flat_parameter()->set_metric_type(::dingodb::pb::common::MetricType::METRIC_TYPE_COSINE);
vector_index_flat_for_cosine = VectorIndexFactory::New(id, index_parameter, kEpoch, kRange);
vector_index_flat_for_cosine =
VectorIndexFactory::NewFlat(id, index_parameter, kEpoch, kRange, vector_index_thread_pool);
EXPECT_NE(vector_index_flat_for_cosine.get(), nullptr);
}
}
Expand Down
Loading

0 comments on commit b713ef2

Please sign in to comment.