mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-08 18:18:30 +08:00
328 lines
11 KiB
C++
328 lines
11 KiB
C++
#include "gtest/gtest.h"
|
|
|
|
#include "db/VectorSource.h"
|
|
#include "db/MemTableFile.h"
|
|
#include "db/MemTable.h"
|
|
#include "utils.h"
|
|
#include "db/Factories.h"
|
|
#include "db/Constants.h"
|
|
#include "db/EngineFactory.h"
|
|
#include "metrics/Metrics.h"
|
|
|
|
#include <thread>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
|
|
using namespace zilliz::milvus;
|
|
|
|
namespace {
|
|
|
|
static const std::string TABLE_NAME = "test_group";
|
|
static constexpr int64_t TABLE_DIM = 256;
|
|
static constexpr int64_t VECTOR_COUNT = 250000;
|
|
static constexpr int64_t INSERT_LOOP = 10000;
|
|
|
|
engine::meta::TableSchema BuildTableSchema() {
|
|
engine::meta::TableSchema table_info;
|
|
table_info.dimension_ = TABLE_DIM;
|
|
table_info.table_id_ = TABLE_NAME;
|
|
table_info.engine_type_ = (int)engine::EngineType::FAISS_IDMAP;
|
|
return table_info;
|
|
}
|
|
|
|
void BuildVectors(int64_t n, std::vector<float>& vectors) {
|
|
vectors.clear();
|
|
vectors.resize(n*TABLE_DIM);
|
|
float* data = vectors.data();
|
|
// std::random_device rd;
|
|
// std::mt19937 gen(rd());
|
|
// std::uniform_real_distribution<> dis(0.0, 1.0);
|
|
for(int i = 0; i < n; i++) {
|
|
for(int j = 0; j < TABLE_DIM; j++) data[TABLE_DIM * i + j] = drand48();
|
|
data[TABLE_DIM * i] += i / 2000.;
|
|
}
|
|
}
|
|
}
|
|
|
|
TEST(MEM_TEST, VECTOR_SOURCE_TEST) {
|
|
|
|
std::shared_ptr<engine::meta::DBMetaImpl> impl_ = engine::DBMetaImplFactory::Build();
|
|
|
|
engine::meta::TableSchema table_schema = BuildTableSchema();
|
|
auto status = impl_->CreateTable(table_schema);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
engine::meta::TableFileSchema table_file_schema;
|
|
table_file_schema.table_id_ = TABLE_NAME;
|
|
status = impl_->CreateTableFile(table_file_schema);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
int64_t n = 100;
|
|
std::vector<float> vectors;
|
|
BuildVectors(n, vectors);
|
|
|
|
engine::VectorSource source(n, vectors.data());
|
|
|
|
size_t num_vectors_added;
|
|
engine::ExecutionEnginePtr execution_engine_ = engine::EngineFactory::Build(table_file_schema.dimension_,
|
|
table_file_schema.location_,
|
|
(engine::EngineType)table_file_schema.engine_type_);
|
|
status = source.Add(execution_engine_, table_file_schema, 50, num_vectors_added);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
ASSERT_EQ(num_vectors_added, 50);
|
|
|
|
engine::IDNumbers vector_ids = source.GetVectorIds();
|
|
ASSERT_EQ(vector_ids.size(), 50);
|
|
|
|
status = source.Add(execution_engine_, table_file_schema, 60, num_vectors_added);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
ASSERT_EQ(num_vectors_added, 50);
|
|
|
|
vector_ids = source.GetVectorIds();
|
|
ASSERT_EQ(vector_ids.size(), 100);
|
|
|
|
// for (auto& id : vector_ids) {
|
|
// std::cout << id << std::endl;
|
|
// }
|
|
|
|
status = impl_->DropAll();
|
|
ASSERT_TRUE(status.ok());
|
|
}
|
|
|
|
TEST(MEM_TEST, MEM_TABLE_FILE_TEST) {
|
|
|
|
std::shared_ptr<engine::meta::DBMetaImpl> impl_ = engine::DBMetaImplFactory::Build();
|
|
auto options = engine::OptionsFactory::Build();
|
|
|
|
engine::meta::TableSchema table_schema = BuildTableSchema();
|
|
auto status = impl_->CreateTable(table_schema);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
engine::MemTableFile memTableFile(TABLE_NAME, impl_, options);
|
|
|
|
int64_t n_100 = 100;
|
|
std::vector<float> vectors_100;
|
|
BuildVectors(n_100, vectors_100);
|
|
|
|
engine::VectorSource::Ptr source = std::make_shared<engine::VectorSource>(n_100, vectors_100.data());
|
|
|
|
status = memTableFile.Add(source);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
// std::cout << memTableFile.GetCurrentMem() << " " << memTableFile.GetMemLeft() << std::endl;
|
|
|
|
engine::IDNumbers vector_ids = source->GetVectorIds();
|
|
ASSERT_EQ(vector_ids.size(), 100);
|
|
|
|
size_t singleVectorMem = sizeof(float) * TABLE_DIM;
|
|
ASSERT_EQ(memTableFile.GetCurrentMem(), n_100 * singleVectorMem);
|
|
|
|
int64_t n_max = engine::MAX_TABLE_FILE_MEM / singleVectorMem;
|
|
std::vector<float> vectors_128M;
|
|
BuildVectors(n_max, vectors_128M);
|
|
|
|
engine::VectorSource::Ptr source_128M = std::make_shared<engine::VectorSource>(n_max, vectors_128M.data());
|
|
status = memTableFile.Add(source_128M);
|
|
|
|
vector_ids = source_128M->GetVectorIds();
|
|
ASSERT_EQ(vector_ids.size(), n_max - n_100);
|
|
|
|
ASSERT_TRUE(memTableFile.IsFull());
|
|
|
|
status = impl_->DropAll();
|
|
ASSERT_TRUE(status.ok());
|
|
}
|
|
|
|
TEST(MEM_TEST, MEM_TABLE_TEST) {
|
|
|
|
std::shared_ptr<engine::meta::DBMetaImpl> impl_ = engine::DBMetaImplFactory::Build();
|
|
auto options = engine::OptionsFactory::Build();
|
|
|
|
engine::meta::TableSchema table_schema = BuildTableSchema();
|
|
auto status = impl_->CreateTable(table_schema);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
int64_t n_100 = 100;
|
|
std::vector<float> vectors_100;
|
|
BuildVectors(n_100, vectors_100);
|
|
|
|
engine::VectorSource::Ptr source_100 = std::make_shared<engine::VectorSource>(n_100, vectors_100.data());
|
|
|
|
engine::MemTable memTable(TABLE_NAME, impl_, options);
|
|
|
|
status = memTable.Add(source_100);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
engine::IDNumbers vector_ids = source_100->GetVectorIds();
|
|
ASSERT_EQ(vector_ids.size(), 100);
|
|
|
|
engine::MemTableFile::Ptr memTableFile;
|
|
memTable.GetCurrentMemTableFile(memTableFile);
|
|
size_t singleVectorMem = sizeof(float) * TABLE_DIM;
|
|
ASSERT_EQ(memTableFile->GetCurrentMem(), n_100 * singleVectorMem);
|
|
|
|
int64_t n_max = engine::MAX_TABLE_FILE_MEM / singleVectorMem;
|
|
std::vector<float> vectors_128M;
|
|
BuildVectors(n_max, vectors_128M);
|
|
|
|
engine::VectorSource::Ptr source_128M = std::make_shared<engine::VectorSource>(n_max, vectors_128M.data());
|
|
status = memTable.Add(source_128M);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
vector_ids = source_128M->GetVectorIds();
|
|
ASSERT_EQ(vector_ids.size(), n_max);
|
|
|
|
memTable.GetCurrentMemTableFile(memTableFile);
|
|
ASSERT_EQ(memTableFile->GetCurrentMem(), n_100 * singleVectorMem);
|
|
|
|
ASSERT_EQ(memTable.GetTableFileCount(), 2);
|
|
|
|
int64_t n_1G = 1024000;
|
|
std::vector<float> vectors_1G;
|
|
BuildVectors(n_1G, vectors_1G);
|
|
|
|
engine::VectorSource::Ptr source_1G = std::make_shared<engine::VectorSource>(n_1G, vectors_1G.data());
|
|
|
|
status = memTable.Add(source_1G);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
vector_ids = source_1G->GetVectorIds();
|
|
ASSERT_EQ(vector_ids.size(), n_1G);
|
|
|
|
int expectedTableFileCount = 2 + std::ceil((n_1G - n_100) * singleVectorMem / engine::MAX_TABLE_FILE_MEM);
|
|
ASSERT_EQ(memTable.GetTableFileCount(), expectedTableFileCount);
|
|
|
|
status = memTable.Serialize();
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
status = impl_->DropAll();
|
|
ASSERT_TRUE(status.ok());
|
|
}
|
|
|
|
TEST(MEM_TEST, MEM_MANAGER_TEST) {
|
|
|
|
auto options = engine::OptionsFactory::Build();
|
|
options.meta.path = "/tmp/milvus_test";
|
|
options.meta.backend_uri = "sqlite://:@:/";
|
|
auto db_ = engine::DBFactory::Build(options);
|
|
|
|
engine::meta::TableSchema table_info = BuildTableSchema();
|
|
engine::Status stat = db_->CreateTable(table_info);
|
|
|
|
engine::meta::TableSchema table_info_get;
|
|
table_info_get.table_id_ = TABLE_NAME;
|
|
stat = db_->DescribeTable(table_info_get);
|
|
ASSERT_STATS(stat);
|
|
ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);
|
|
|
|
std::map<int64_t , std::vector<float>> search_vectors;
|
|
// std::map<int64_t , std::vector<float>> vectors_ids_map;
|
|
{
|
|
engine::IDNumbers vector_ids;
|
|
int64_t nb = 1024000;
|
|
std::vector<float> xb;
|
|
BuildVectors(nb, xb);
|
|
engine::Status status = db_->InsertVectors(TABLE_NAME, nb, xb.data(), vector_ids);
|
|
ASSERT_TRUE(status.ok());
|
|
|
|
// std::ofstream myfile("mem_test.txt");
|
|
// for (int64_t i = 0; i < nb; ++i) {
|
|
// int64_t vector_id = vector_ids[i];
|
|
// std::vector<float> vectors;
|
|
// for (int64_t j = 0; j < TABLE_DIM; j++) {
|
|
// vectors.emplace_back(xb[i*TABLE_DIM + j]);
|
|
//// std::cout << xb[i*TABLE_DIM + j] << std::endl;
|
|
// }
|
|
// vectors_ids_map[vector_id] = vectors;
|
|
// }
|
|
|
|
std::this_thread::sleep_for(std::chrono::seconds(3));
|
|
|
|
std::random_device rd;
|
|
std::mt19937 gen(rd());
|
|
std::uniform_int_distribution<int64_t> dis(0, nb - 1);
|
|
|
|
int64_t numQuery = 1000;
|
|
for (int64_t i = 0; i < numQuery; ++i) {
|
|
int64_t index = dis(gen);
|
|
std::vector<float> search;
|
|
for (int64_t j = 0; j < TABLE_DIM; j++) {
|
|
search.push_back(xb[index * TABLE_DIM + j]);
|
|
}
|
|
search_vectors.insert(std::make_pair(vector_ids[index], search));
|
|
// std::cout << "index: " << index << " vector_ids[index]: " << vector_ids[index] << std::endl;
|
|
}
|
|
|
|
// for (int64_t i = 0; i < nb; i += 100000) {
|
|
// std::vector<float> search;
|
|
// for (int64_t j = 0; j < TABLE_DIM; j++) {
|
|
// search.push_back(xb[i * TABLE_DIM + j]);
|
|
// }
|
|
// search_vectors.insert(std::make_pair(vector_ids[i], search));
|
|
// }
|
|
|
|
}
|
|
|
|
int k = 10;
|
|
for(auto& pair : search_vectors) {
|
|
auto& search = pair.second;
|
|
engine::QueryResults results;
|
|
stat = db_->Query(TABLE_NAME, k, 1, search.data(), results);
|
|
for(int t = 0; t < k; t++) {
|
|
// std::cout << "ID=" << results[0][t].first << " DISTANCE=" << results[0][t].second << std::endl;
|
|
|
|
// std::cout << vectors_ids_map[results[0][t].first].size() << std::endl;
|
|
// for (auto& data : vectors_ids_map[results[0][t].first]) {
|
|
// std::cout << data << " ";
|
|
// }
|
|
// std::cout << std::endl;
|
|
}
|
|
// std::cout << "results[0][0].first: " << results[0][0].first << " pair.first: " << pair.first << " results[0][0].second: " << results[0][0].second << std::endl;
|
|
ASSERT_EQ(results[0][0].first, pair.first);
|
|
ASSERT_LT(results[0][0].second, 0.00001);
|
|
}
|
|
|
|
stat = db_->DropAll();
|
|
ASSERT_TRUE(stat.ok());
|
|
|
|
}
|
|
|
|
TEST(MEM_TEST, INSERT_TEST) {
|
|
|
|
auto options = engine::OptionsFactory::Build();
|
|
options.meta.path = "/tmp/milvus_test";
|
|
options.meta.backend_uri = "sqlite://:@:/";
|
|
auto db_ = engine::DBFactory::Build(options);
|
|
|
|
engine::meta::TableSchema table_info = BuildTableSchema();
|
|
engine::Status stat = db_->CreateTable(table_info);
|
|
|
|
engine::meta::TableSchema table_info_get;
|
|
table_info_get.table_id_ = TABLE_NAME;
|
|
stat = db_->DescribeTable(table_info_get);
|
|
ASSERT_STATS(stat);
|
|
ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);
|
|
|
|
auto start_time = METRICS_NOW_TIME;
|
|
|
|
int insert_loop = 1000;
|
|
for (int i = 0; i < insert_loop; ++i) {
|
|
int64_t nb = 204800;
|
|
std::vector<float> xb;
|
|
BuildVectors(nb, xb);
|
|
engine::IDNumbers vector_ids;
|
|
engine::Status status = db_->InsertVectors(TABLE_NAME, nb, xb.data(), vector_ids);
|
|
ASSERT_TRUE(status.ok());
|
|
}
|
|
auto end_time = METRICS_NOW_TIME;
|
|
auto total_time = METRICS_MICROSECONDS(start_time, end_time);
|
|
std::cout << "total_time(ms) : " << total_time << std::endl;
|
|
|
|
stat = db_->DropAll();
|
|
ASSERT_TRUE(stat.ok());
|
|
|
|
}
|
|
|