mirror of https://github.com/milvus-io/milvus.git
278 lines
11 KiB
C++
278 lines
11 KiB
C++
// Licensed to the LF AI & Data foundation under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you 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 <gtest/gtest.h>
|
|
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "fmt/format.h"
|
|
#include "common/Schema.h"
|
|
#include "test_utils/DataGen.h"
|
|
#include "test_utils/storage_test_utils.h"
|
|
#include "storage/ChunkCache.h"
|
|
#include "storage/LocalChunkManagerSingleton.h"
|
|
|
|
#define DEFAULT_READ_AHEAD_POLICY "willneed"
|
|
class ChunkCacheTest : public testing::TestWithParam</*mmap enabled*/ bool> {
|
|
public:
|
|
void
|
|
SetUp() override {
|
|
mcm = milvus::storage::MmapManager::GetInstance().GetMmapChunkManager();
|
|
mcm->Register(descriptor);
|
|
}
|
|
void
|
|
TearDown() override {
|
|
mcm->UnRegister(descriptor);
|
|
}
|
|
const char* dense_file_name = "chunk_cache_test/insert_log/2/101/1000000";
|
|
const char* sparse_file_name = "chunk_cache_test/insert_log/2/102/1000000";
|
|
milvus::storage::MmapChunkManagerPtr mcm;
|
|
milvus::segcore::SegcoreConfig config;
|
|
milvus::storage::MmapChunkDescriptorPtr descriptor =
|
|
std::shared_ptr<milvus::storage::MmapChunkDescriptor>(
|
|
new milvus::storage::MmapChunkDescriptor(
|
|
{101, SegmentType::Sealed}));
|
|
};
|
|
|
|
INSTANTIATE_TEST_SUITE_P(ChunkCacheTestSuite,
|
|
ChunkCacheTest,
|
|
testing::Values(true, false));
|
|
|
|
TEST_P(ChunkCacheTest, Read) {
|
|
auto N = 10000;
|
|
auto dim = 128;
|
|
auto dense_metric_type = knowhere::metric::L2;
|
|
auto sparse_metric_type = knowhere::metric::IP;
|
|
|
|
auto schema = std::make_shared<milvus::Schema>();
|
|
auto fake_dense_vec_id = schema->AddDebugField(
|
|
"fakevec", milvus::DataType::VECTOR_FLOAT, dim, dense_metric_type);
|
|
auto i64_fid = schema->AddDebugField("counter", milvus::DataType::INT64);
|
|
auto fake_sparse_vec_id =
|
|
schema->AddDebugField("fakevec_sparse",
|
|
milvus::DataType::VECTOR_SPARSE_FLOAT,
|
|
dim,
|
|
sparse_metric_type);
|
|
schema->set_primary_field_id(i64_fid);
|
|
|
|
auto dataset = milvus::segcore::DataGen(schema, N);
|
|
|
|
auto dense_field_data_meta =
|
|
milvus::storage::FieldDataMeta{1, 2, 3, fake_dense_vec_id.get()};
|
|
auto sparse_field_data_meta =
|
|
milvus::storage::FieldDataMeta{1, 2, 3, fake_sparse_vec_id.get()};
|
|
auto dense_field_meta = milvus::FieldMeta(milvus::FieldName("fakevec"),
|
|
fake_dense_vec_id,
|
|
milvus::DataType::VECTOR_FLOAT,
|
|
dim,
|
|
dense_metric_type,
|
|
false);
|
|
auto sparse_field_meta =
|
|
milvus::FieldMeta(milvus::FieldName("fakevec_sparse"),
|
|
fake_sparse_vec_id,
|
|
milvus::DataType::VECTOR_SPARSE_FLOAT,
|
|
dim,
|
|
sparse_metric_type,
|
|
false);
|
|
|
|
auto lcm = milvus::storage::LocalChunkManagerSingleton::GetInstance()
|
|
.GetChunkManager();
|
|
auto dense_data = dataset.get_col<float>(fake_dense_vec_id);
|
|
auto sparse_data =
|
|
dataset.get_col<knowhere::sparse::SparseRow<float>>(fake_sparse_vec_id);
|
|
|
|
auto data_slices = std::vector<void*>{dense_data.data()};
|
|
auto slice_sizes = std::vector<int64_t>{static_cast<int64_t>(N)};
|
|
auto slice_names = std::vector<std::string>{dense_file_name};
|
|
PutFieldData(lcm.get(),
|
|
data_slices,
|
|
slice_sizes,
|
|
slice_names,
|
|
dense_field_data_meta,
|
|
dense_field_meta);
|
|
|
|
data_slices = std::vector<void*>{sparse_data.data()};
|
|
slice_sizes = std::vector<int64_t>{static_cast<int64_t>(N)};
|
|
slice_names = std::vector<std::string>{sparse_file_name};
|
|
PutFieldData(lcm.get(),
|
|
data_slices,
|
|
slice_sizes,
|
|
slice_names,
|
|
sparse_field_data_meta,
|
|
sparse_field_meta);
|
|
|
|
auto cc = milvus::storage::MmapManager::GetInstance().GetChunkCache();
|
|
|
|
// validate dense data
|
|
const auto& dense_column =
|
|
cc->Read(dense_file_name, descriptor, dense_field_meta, GetParam());
|
|
Assert(dense_column->DataByteSize() == dim * N * 4);
|
|
auto actual_dense = (const float*)(dense_column->Data());
|
|
for (auto i = 0; i < N * dim; i++) {
|
|
AssertInfo(dense_data[i] == actual_dense[i],
|
|
fmt::format(
|
|
"expect {}, actual {}", dense_data[i], actual_dense[i]));
|
|
}
|
|
|
|
// validate sparse data
|
|
const auto& sparse_column =
|
|
cc->Read(sparse_file_name, descriptor, sparse_field_meta, GetParam());
|
|
auto expected_sparse_size = 0;
|
|
auto actual_sparse =
|
|
(const knowhere::sparse::SparseRow<float>*)(sparse_column->Data());
|
|
for (auto i = 0; i < N; i++) {
|
|
const auto& actual_sparse_row = actual_sparse[i];
|
|
const auto& expect_sparse_row = sparse_data[i];
|
|
AssertInfo(
|
|
actual_sparse_row.size() == expect_sparse_row.size(),
|
|
fmt::format("Incorrect size of sparse row: expect {}, actual {}",
|
|
expect_sparse_row.size(),
|
|
actual_sparse_row.size()));
|
|
auto bytes = actual_sparse_row.data_byte_size();
|
|
AssertInfo(
|
|
memcmp(actual_sparse_row.data(), expect_sparse_row.data(), bytes) ==
|
|
0,
|
|
fmt::format("Incorrect data of sparse row: expect {}, actual {}",
|
|
expect_sparse_row.data(),
|
|
actual_sparse_row.data()));
|
|
expected_sparse_size += bytes;
|
|
}
|
|
|
|
ASSERT_EQ(sparse_column->DataByteSize(), expected_sparse_size);
|
|
|
|
cc->Remove(dense_file_name);
|
|
cc->Remove(sparse_file_name);
|
|
lcm->Remove(dense_file_name);
|
|
lcm->Remove(sparse_file_name);
|
|
}
|
|
|
|
TEST_P(ChunkCacheTest, TestMultithreads) {
|
|
auto N = 1000;
|
|
auto dim = 128;
|
|
auto dense_metric_type = knowhere::metric::L2;
|
|
auto sparse_metric_type = knowhere::metric::IP;
|
|
|
|
auto schema = std::make_shared<milvus::Schema>();
|
|
auto fake_dense_vec_id = schema->AddDebugField(
|
|
"fakevec", milvus::DataType::VECTOR_FLOAT, dim, dense_metric_type);
|
|
auto fake_sparse_vec_id =
|
|
schema->AddDebugField("fakevec_sparse",
|
|
milvus::DataType::VECTOR_SPARSE_FLOAT,
|
|
dim,
|
|
sparse_metric_type);
|
|
auto i64_fid = schema->AddDebugField("counter", milvus::DataType::INT64);
|
|
schema->set_primary_field_id(i64_fid);
|
|
|
|
auto dataset = milvus::segcore::DataGen(schema, N);
|
|
|
|
auto dense_field_data_meta =
|
|
milvus::storage::FieldDataMeta{1, 2, 3, fake_dense_vec_id.get()};
|
|
auto sparse_field_data_meta =
|
|
milvus::storage::FieldDataMeta{1, 2, 3, fake_sparse_vec_id.get()};
|
|
auto dense_field_meta = milvus::FieldMeta(milvus::FieldName("fakevec"),
|
|
fake_dense_vec_id,
|
|
milvus::DataType::VECTOR_FLOAT,
|
|
dim,
|
|
dense_metric_type,
|
|
false);
|
|
auto sparse_field_meta =
|
|
milvus::FieldMeta(milvus::FieldName("fakevec_sparse"),
|
|
fake_sparse_vec_id,
|
|
milvus::DataType::VECTOR_SPARSE_FLOAT,
|
|
dim,
|
|
sparse_metric_type,
|
|
false);
|
|
|
|
auto lcm = milvus::storage::LocalChunkManagerSingleton::GetInstance()
|
|
.GetChunkManager();
|
|
auto dense_data = dataset.get_col<float>(fake_dense_vec_id);
|
|
auto sparse_data =
|
|
dataset.get_col<knowhere::sparse::SparseRow<float>>(fake_sparse_vec_id);
|
|
|
|
auto dense_data_slices = std::vector<void*>{dense_data.data()};
|
|
auto sparse_data_slices = std::vector<void*>{sparse_data.data()};
|
|
auto slice_sizes = std::vector<int64_t>{static_cast<int64_t>(N)};
|
|
auto dense_slice_names = std::vector<std::string>{dense_file_name};
|
|
auto sparse_slice_names = std::vector<std::string>{sparse_file_name};
|
|
|
|
PutFieldData(lcm.get(),
|
|
dense_data_slices,
|
|
slice_sizes,
|
|
dense_slice_names,
|
|
dense_field_data_meta,
|
|
dense_field_meta);
|
|
|
|
PutFieldData(lcm.get(),
|
|
sparse_data_slices,
|
|
slice_sizes,
|
|
sparse_slice_names,
|
|
sparse_field_data_meta,
|
|
sparse_field_meta);
|
|
|
|
auto cc = milvus::storage::MmapManager::GetInstance().GetChunkCache();
|
|
|
|
constexpr int threads = 16;
|
|
std::vector<int64_t> total_counts(threads);
|
|
auto executor = [&](int thread_id) {
|
|
const auto& dense_column =
|
|
cc->Read(dense_file_name, descriptor, dense_field_meta, GetParam());
|
|
Assert(dense_column->DataByteSize() == dim * N * 4);
|
|
|
|
auto actual_dense = (const float*)dense_column->Data();
|
|
for (auto i = 0; i < N * dim; i++) {
|
|
AssertInfo(
|
|
dense_data[i] == actual_dense[i],
|
|
fmt::format(
|
|
"expect {}, actual {}", dense_data[i], actual_dense[i]));
|
|
}
|
|
|
|
const auto& sparse_column = cc->Read(
|
|
sparse_file_name, descriptor, sparse_field_meta, GetParam());
|
|
auto actual_sparse =
|
|
(const knowhere::sparse::SparseRow<float>*)sparse_column->Data();
|
|
for (auto i = 0; i < N; i++) {
|
|
const auto& actual_sparse_row = actual_sparse[i];
|
|
const auto& expect_sparse_row = sparse_data[i];
|
|
AssertInfo(actual_sparse_row.size() == expect_sparse_row.size(),
|
|
fmt::format(
|
|
"Incorrect size of sparse row: expect {}, actual {}",
|
|
expect_sparse_row.size(),
|
|
actual_sparse_row.size()));
|
|
auto bytes = actual_sparse_row.data_byte_size();
|
|
AssertInfo(memcmp(actual_sparse_row.data(),
|
|
expect_sparse_row.data(),
|
|
bytes) == 0,
|
|
fmt::format(
|
|
"Incorrect data of sparse row: expect {}, actual {}",
|
|
expect_sparse_row.data(),
|
|
actual_sparse_row.data()));
|
|
}
|
|
};
|
|
std::vector<std::thread> pool;
|
|
for (int i = 0; i < threads; ++i) {
|
|
pool.emplace_back(executor, i);
|
|
}
|
|
for (auto& thread : pool) {
|
|
thread.join();
|
|
}
|
|
|
|
cc->Remove(dense_file_name);
|
|
cc->Remove(sparse_file_name);
|
|
lcm->Remove(dense_file_name);
|
|
lcm->Remove(sparse_file_name);
|
|
}
|