mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
enhance: make DefaultValueChunk multi-cells to save memory usage (#46000)
issue: #45999 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> - Core invariant: DefaultValueChunkTranslator assumes a constant per-value byte size (value_size()) for a field and therefore models default-value storage as multiple deterministic, fixed-size cells (num_cells(), cell_id_of(uid) == uid). A single contiguous primary_buffer_ is built (optionally persisted to mmap) and sliced per-cell; this invariant enables safe sharing of the same underlying bytes across multiple Chunk views. - Removed/simplified logic: per-field direct file_path writes and duplicated make_chunk write paths were collapsed into a buffer-first flow (ChunkBuffer + create_chunk_buffer + make_chunk_from_buffer). This removes redundant per-field mmap/write code and allows multiple Chunk instances to reuse a shared ChunkBuffer/ChunkMmapGuard instead of performing independent writes. - Why no data loss or regression: byte layout, alignment, padding and per-row content are preserved because create_chunk_buffer implements the same alignment/padding and writes identical bytes previously produced by the direct make_chunk path; make_chunk_from_buffer constructs Chunk views over those same bytes. DefaultValueChunkTranslator uses deterministic math (total_rows_, num_rows_until_chunk_, primary_cell_rows_) to slice rows—no rows are dropped, reordered, or altered. Concrete code paths: callers that used create_group_chunk → per-field make_chunk now use create_group_chunk → create_chunk_buffer → make_chunk_from_buffer; DefaultValueChunkTranslator’s get_cells(), estimated_byte_size_of_cell(), and value_size() produce the same observable outputs as before for single-cell cases and correct per-cell outputs for multi-cell cases. - Enhancement / scope: adds multi-cell default-value chunks and optional mmap-backed persistence plus shared-memory ChunkBuffer support. Changes touch DefaultValueChunkTranslator (multi-cell behavior, value_size(), primary_buffer_, mmap_dir_path_), ChunkWriter (ChunkBuffer, create_chunk_buffer, make_chunk_from_buffer), and ChunkedSegmentSealedImpl (per-field mmap decision and mmap_dir_path propagation); comprehensive tests (mmap/non-mmap, nullable, multi/single-cell) validate correctness. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shawn Wang <shawn.wang@zilliz.com>
This commit is contained in:
sparknack
@@ -961,4 +961,403 @@ TEST(chunk, test_binary_search_methods_comparison) {
|
||||
EXPECT_EQ(ub, lb + 1) << "For existing target: " << target;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(chunk, test_create_group_chunk_basic) {
|
||||
// Prepare data for multiple fields
|
||||
size_t row_count = 5;
|
||||
|
||||
// Field 1: INT64
|
||||
FixedVector<int64_t> int_data = {1, 2, 3, 4, 5};
|
||||
auto int_field_data = milvus::storage::CreateFieldData(
|
||||
storage::DataType::INT64, DataType::NONE);
|
||||
int_field_data->FillFieldData(int_data.data(), int_data.size());
|
||||
|
||||
// Field 2: VARCHAR
|
||||
FixedVector<std::string> str_data = {
|
||||
"test1", "test2", "test3", "test4", "test5"};
|
||||
auto str_field_data = milvus::storage::CreateFieldData(
|
||||
storage::DataType::VARCHAR, DataType::NONE);
|
||||
str_field_data->FillFieldData(str_data.data(), str_data.size());
|
||||
|
||||
// Field 3: DOUBLE
|
||||
FixedVector<double> double_data = {1.1, 2.2, 3.3, 4.4, 5.5};
|
||||
auto double_field_data = milvus::storage::CreateFieldData(
|
||||
storage::DataType::DOUBLE, DataType::NONE);
|
||||
double_field_data->FillFieldData(double_data.data(), double_data.size());
|
||||
|
||||
// Create arrow arrays for each field
|
||||
std::vector<arrow::ArrayVector> array_vecs;
|
||||
|
||||
// Process INT64 field
|
||||
{
|
||||
storage::InsertEventData event_data;
|
||||
auto payload_reader =
|
||||
std::make_shared<milvus::storage::PayloadReader>(int_field_data);
|
||||
event_data.payload_reader = payload_reader;
|
||||
auto ser_data = event_data.Serialize();
|
||||
auto buffer = std::make_shared<arrow::io::BufferReader>(
|
||||
ser_data.data() + 2 * sizeof(milvus::Timestamp),
|
||||
ser_data.size() - 2 * sizeof(milvus::Timestamp));
|
||||
|
||||
parquet::arrow::FileReaderBuilder reader_builder;
|
||||
auto s = reader_builder.Open(buffer);
|
||||
EXPECT_TRUE(s.ok());
|
||||
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
|
||||
s = reader_builder.Build(&arrow_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
|
||||
s = arrow_reader->GetRecordBatchReader(&rb_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
array_vecs.push_back(read_single_column_batches(rb_reader));
|
||||
}
|
||||
|
||||
// Process VARCHAR field
|
||||
{
|
||||
storage::InsertEventData event_data;
|
||||
auto payload_reader =
|
||||
std::make_shared<milvus::storage::PayloadReader>(str_field_data);
|
||||
event_data.payload_reader = payload_reader;
|
||||
auto ser_data = event_data.Serialize();
|
||||
auto buffer = std::make_shared<arrow::io::BufferReader>(
|
||||
ser_data.data() + 2 * sizeof(milvus::Timestamp),
|
||||
ser_data.size() - 2 * sizeof(milvus::Timestamp));
|
||||
|
||||
parquet::arrow::FileReaderBuilder reader_builder;
|
||||
auto s = reader_builder.Open(buffer);
|
||||
EXPECT_TRUE(s.ok());
|
||||
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
|
||||
s = reader_builder.Build(&arrow_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
|
||||
s = arrow_reader->GetRecordBatchReader(&rb_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
array_vecs.push_back(read_single_column_batches(rb_reader));
|
||||
}
|
||||
|
||||
// Process DOUBLE field
|
||||
{
|
||||
storage::InsertEventData event_data;
|
||||
auto payload_reader =
|
||||
std::make_shared<milvus::storage::PayloadReader>(double_field_data);
|
||||
event_data.payload_reader = payload_reader;
|
||||
auto ser_data = event_data.Serialize();
|
||||
auto buffer = std::make_shared<arrow::io::BufferReader>(
|
||||
ser_data.data() + 2 * sizeof(milvus::Timestamp),
|
||||
ser_data.size() - 2 * sizeof(milvus::Timestamp));
|
||||
|
||||
parquet::arrow::FileReaderBuilder reader_builder;
|
||||
auto s = reader_builder.Open(buffer);
|
||||
EXPECT_TRUE(s.ok());
|
||||
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
|
||||
s = reader_builder.Build(&arrow_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
|
||||
s = arrow_reader->GetRecordBatchReader(&rb_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
array_vecs.push_back(read_single_column_batches(rb_reader));
|
||||
}
|
||||
|
||||
// Create field metadata
|
||||
std::vector<FieldId> field_ids = {FieldId(1), FieldId(2), FieldId(3)};
|
||||
std::vector<FieldMeta> field_metas = {FieldMeta(FieldName("int_field"),
|
||||
FieldId(1),
|
||||
DataType::INT64,
|
||||
false,
|
||||
std::nullopt),
|
||||
FieldMeta(FieldName("str_field"),
|
||||
FieldId(2),
|
||||
DataType::STRING,
|
||||
false,
|
||||
std::nullopt),
|
||||
FieldMeta(FieldName("double_field"),
|
||||
FieldId(3),
|
||||
DataType::DOUBLE,
|
||||
false,
|
||||
std::nullopt)};
|
||||
|
||||
// Create group chunk without mmap
|
||||
auto chunks = create_group_chunk(field_ids, field_metas, array_vecs);
|
||||
|
||||
// Verify all chunks were created
|
||||
EXPECT_EQ(chunks.size(), 3);
|
||||
EXPECT_NE(chunks.find(FieldId(1)), chunks.end());
|
||||
EXPECT_NE(chunks.find(FieldId(2)), chunks.end());
|
||||
EXPECT_NE(chunks.find(FieldId(3)), chunks.end());
|
||||
|
||||
// Verify INT64 chunk
|
||||
auto int_chunk = static_cast<FixedWidthChunk*>(chunks[FieldId(1)].get());
|
||||
auto int_span = int_chunk->Span();
|
||||
EXPECT_EQ(int_span.row_count(), row_count);
|
||||
for (size_t i = 0; i < row_count; ++i) {
|
||||
auto value =
|
||||
*(int64_t*)((char*)int_span.data() + i * int_span.element_sizeof());
|
||||
EXPECT_EQ(value, int_data[i]);
|
||||
}
|
||||
|
||||
// Verify STRING chunk
|
||||
auto str_chunk = static_cast<StringChunk*>(chunks[FieldId(2)].get());
|
||||
auto [str_views, str_valid] = str_chunk->StringViews(std::nullopt);
|
||||
EXPECT_EQ(str_views.size(), row_count);
|
||||
for (size_t i = 0; i < row_count; ++i) {
|
||||
EXPECT_EQ(str_views[i], str_data[i]);
|
||||
}
|
||||
|
||||
// Verify DOUBLE chunk
|
||||
auto double_chunk = static_cast<FixedWidthChunk*>(chunks[FieldId(3)].get());
|
||||
auto double_span = double_chunk->Span();
|
||||
EXPECT_EQ(double_span.row_count(), row_count);
|
||||
for (size_t i = 0; i < row_count; ++i) {
|
||||
auto value = *(double*)((char*)double_span.data() +
|
||||
i * double_span.element_sizeof());
|
||||
EXPECT_DOUBLE_EQ(value, double_data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(chunk, test_create_group_chunk_with_mmap) {
|
||||
// Prepare data for multiple fields
|
||||
size_t row_count = 5;
|
||||
|
||||
// Field 1: INT32
|
||||
FixedVector<int32_t> int32_data = {10, 20, 30, 40, 50};
|
||||
auto int32_field_data = milvus::storage::CreateFieldData(
|
||||
storage::DataType::INT32, DataType::NONE);
|
||||
int32_field_data->FillFieldData(int32_data.data(), int32_data.size());
|
||||
|
||||
// Field 2: FLOAT
|
||||
FixedVector<float> float_data = {1.5f, 2.5f, 3.5f, 4.5f, 5.5f};
|
||||
auto float_field_data = milvus::storage::CreateFieldData(
|
||||
storage::DataType::FLOAT, DataType::NONE);
|
||||
float_field_data->FillFieldData(float_data.data(), float_data.size());
|
||||
|
||||
// Create arrow arrays for each field
|
||||
std::vector<arrow::ArrayVector> array_vecs;
|
||||
|
||||
// Process INT32 field
|
||||
{
|
||||
storage::InsertEventData event_data;
|
||||
auto payload_reader =
|
||||
std::make_shared<milvus::storage::PayloadReader>(int32_field_data);
|
||||
event_data.payload_reader = payload_reader;
|
||||
auto ser_data = event_data.Serialize();
|
||||
auto buffer = std::make_shared<arrow::io::BufferReader>(
|
||||
ser_data.data() + 2 * sizeof(milvus::Timestamp),
|
||||
ser_data.size() - 2 * sizeof(milvus::Timestamp));
|
||||
|
||||
parquet::arrow::FileReaderBuilder reader_builder;
|
||||
auto s = reader_builder.Open(buffer);
|
||||
EXPECT_TRUE(s.ok());
|
||||
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
|
||||
s = reader_builder.Build(&arrow_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
|
||||
s = arrow_reader->GetRecordBatchReader(&rb_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
array_vecs.push_back(read_single_column_batches(rb_reader));
|
||||
}
|
||||
|
||||
// Process FLOAT field
|
||||
{
|
||||
storage::InsertEventData event_data;
|
||||
auto payload_reader =
|
||||
std::make_shared<milvus::storage::PayloadReader>(float_field_data);
|
||||
event_data.payload_reader = payload_reader;
|
||||
auto ser_data = event_data.Serialize();
|
||||
auto buffer = std::make_shared<arrow::io::BufferReader>(
|
||||
ser_data.data() + 2 * sizeof(milvus::Timestamp),
|
||||
ser_data.size() - 2 * sizeof(milvus::Timestamp));
|
||||
|
||||
parquet::arrow::FileReaderBuilder reader_builder;
|
||||
auto s = reader_builder.Open(buffer);
|
||||
EXPECT_TRUE(s.ok());
|
||||
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
|
||||
s = reader_builder.Build(&arrow_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
|
||||
s = arrow_reader->GetRecordBatchReader(&rb_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
array_vecs.push_back(read_single_column_batches(rb_reader));
|
||||
}
|
||||
|
||||
// Create field metadata
|
||||
std::vector<FieldId> field_ids = {FieldId(10), FieldId(11)};
|
||||
std::vector<FieldMeta> field_metas = {FieldMeta(FieldName("int32_field"),
|
||||
FieldId(10),
|
||||
DataType::INT32,
|
||||
false,
|
||||
std::nullopt),
|
||||
FieldMeta(FieldName("float_field"),
|
||||
FieldId(11),
|
||||
DataType::FLOAT,
|
||||
false,
|
||||
std::nullopt)};
|
||||
|
||||
// Create group chunk with mmap
|
||||
std::string mmap_file = "/tmp/test_group_chunk_mmap.bin";
|
||||
if (boost::filesystem::exists(mmap_file)) {
|
||||
boost::filesystem::remove(mmap_file);
|
||||
}
|
||||
auto chunks =
|
||||
create_group_chunk(field_ids, field_metas, array_vecs, true, mmap_file);
|
||||
|
||||
// Verify all chunks were created
|
||||
EXPECT_EQ(chunks.size(), 2);
|
||||
EXPECT_NE(chunks.find(FieldId(10)), chunks.end());
|
||||
EXPECT_NE(chunks.find(FieldId(11)), chunks.end());
|
||||
|
||||
// Verify INT32 chunk
|
||||
auto int32_chunk = static_cast<FixedWidthChunk*>(chunks[FieldId(10)].get());
|
||||
auto int32_span = int32_chunk->Span();
|
||||
EXPECT_EQ(int32_span.row_count(), row_count);
|
||||
for (size_t i = 0; i < row_count; ++i) {
|
||||
auto value = *(int32_t*)((char*)int32_span.data() +
|
||||
i * int32_span.element_sizeof());
|
||||
EXPECT_EQ(value, int32_data[i]);
|
||||
}
|
||||
|
||||
// Verify FLOAT chunk
|
||||
auto float_chunk = static_cast<FixedWidthChunk*>(chunks[FieldId(11)].get());
|
||||
auto float_span = float_chunk->Span();
|
||||
EXPECT_EQ(float_span.row_count(), row_count);
|
||||
for (size_t i = 0; i < row_count; ++i) {
|
||||
auto value = *(float*)((char*)float_span.data() +
|
||||
i * float_span.element_sizeof());
|
||||
EXPECT_FLOAT_EQ(value, float_data[i]);
|
||||
}
|
||||
|
||||
// Verify file exists
|
||||
EXPECT_TRUE(boost::filesystem::exists(mmap_file));
|
||||
|
||||
// Clean up mmap file
|
||||
chunks.clear();
|
||||
|
||||
// Verify file is removed by ChunkMmapGuard
|
||||
EXPECT_FALSE(boost::filesystem::exists(mmap_file));
|
||||
}
|
||||
|
||||
TEST(chunk, test_create_group_chunk_nullable_fields) {
|
||||
// Test group chunk with nullable fields
|
||||
size_t row_count = 5;
|
||||
|
||||
// Field 1: Nullable INT64
|
||||
FixedVector<int64_t> int_data = {1, 2, 3, 4, 5};
|
||||
auto int_field_data = milvus::storage::CreateFieldData(
|
||||
storage::DataType::INT64, DataType::NONE, true);
|
||||
uint8_t* int_valid_data = new uint8_t[1]{0x15}; // 10101 in binary
|
||||
int_field_data->FillFieldData(
|
||||
int_data.data(), int_valid_data, int_data.size(), 0);
|
||||
delete[] int_valid_data;
|
||||
|
||||
// Field 2: Nullable VARCHAR
|
||||
FixedVector<std::string> str_data = {
|
||||
"test1", "test2", "test3", "test4", "test5"};
|
||||
auto str_field_data = milvus::storage::CreateFieldData(
|
||||
storage::DataType::VARCHAR, DataType::NONE, true);
|
||||
uint8_t* str_valid_data = new uint8_t[1]{0x1A}; // 11010 in binary
|
||||
str_field_data->FillFieldData(
|
||||
str_data.data(), str_valid_data, str_data.size(), 0);
|
||||
delete[] str_valid_data;
|
||||
|
||||
// Create arrow arrays
|
||||
std::vector<arrow::ArrayVector> array_vecs;
|
||||
|
||||
// Process INT64 field
|
||||
{
|
||||
storage::InsertEventData event_data;
|
||||
auto payload_reader =
|
||||
std::make_shared<milvus::storage::PayloadReader>(int_field_data);
|
||||
event_data.payload_reader = payload_reader;
|
||||
auto ser_data = event_data.Serialize();
|
||||
auto buffer = std::make_shared<arrow::io::BufferReader>(
|
||||
ser_data.data() + 2 * sizeof(milvus::Timestamp),
|
||||
ser_data.size() - 2 * sizeof(milvus::Timestamp));
|
||||
|
||||
parquet::arrow::FileReaderBuilder reader_builder;
|
||||
auto s = reader_builder.Open(buffer);
|
||||
EXPECT_TRUE(s.ok());
|
||||
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
|
||||
s = reader_builder.Build(&arrow_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
|
||||
s = arrow_reader->GetRecordBatchReader(&rb_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
array_vecs.push_back(read_single_column_batches(rb_reader));
|
||||
}
|
||||
|
||||
// Process VARCHAR field
|
||||
{
|
||||
storage::InsertEventData event_data;
|
||||
auto payload_reader =
|
||||
std::make_shared<milvus::storage::PayloadReader>(str_field_data);
|
||||
event_data.payload_reader = payload_reader;
|
||||
auto ser_data = event_data.Serialize();
|
||||
auto buffer = std::make_shared<arrow::io::BufferReader>(
|
||||
ser_data.data() + 2 * sizeof(milvus::Timestamp),
|
||||
ser_data.size() - 2 * sizeof(milvus::Timestamp));
|
||||
|
||||
parquet::arrow::FileReaderBuilder reader_builder;
|
||||
auto s = reader_builder.Open(buffer);
|
||||
EXPECT_TRUE(s.ok());
|
||||
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
|
||||
s = reader_builder.Build(&arrow_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
|
||||
s = arrow_reader->GetRecordBatchReader(&rb_reader);
|
||||
EXPECT_TRUE(s.ok());
|
||||
|
||||
array_vecs.push_back(read_single_column_batches(rb_reader));
|
||||
}
|
||||
|
||||
std::vector<FieldId> field_ids = {FieldId(30), FieldId(31)};
|
||||
std::vector<FieldMeta> field_metas = {FieldMeta(FieldName("nullable_int"),
|
||||
FieldId(30),
|
||||
DataType::INT64,
|
||||
true,
|
||||
std::nullopt),
|
||||
FieldMeta(FieldName("nullable_str"),
|
||||
FieldId(31),
|
||||
DataType::STRING,
|
||||
true,
|
||||
std::nullopt)};
|
||||
|
||||
auto chunks = create_group_chunk(field_ids, field_metas, array_vecs);
|
||||
|
||||
EXPECT_EQ(chunks.size(), 2);
|
||||
|
||||
// Verify nullable INT64 chunk (validity: 10101)
|
||||
auto int_chunk = static_cast<FixedWidthChunk*>(chunks[FieldId(30)].get());
|
||||
EXPECT_TRUE(int_chunk->isValid(0));
|
||||
EXPECT_FALSE(int_chunk->isValid(1));
|
||||
EXPECT_TRUE(int_chunk->isValid(2));
|
||||
EXPECT_FALSE(int_chunk->isValid(3));
|
||||
EXPECT_TRUE(int_chunk->isValid(4));
|
||||
|
||||
// Verify nullable VARCHAR chunk (validity: 11010)
|
||||
auto str_chunk = static_cast<StringChunk*>(chunks[FieldId(31)].get());
|
||||
auto [str_views, str_valid] = str_chunk->StringViews(std::nullopt);
|
||||
EXPECT_FALSE(str_valid[0]);
|
||||
EXPECT_TRUE(str_valid[1]);
|
||||
EXPECT_FALSE(str_valid[2]);
|
||||
EXPECT_TRUE(str_valid[3]);
|
||||
EXPECT_TRUE(str_valid[4]);
|
||||
|
||||
// Verify data for valid entries
|
||||
for (size_t i = 0; i < row_count; ++i) {
|
||||
if (str_valid[i]) {
|
||||
EXPECT_EQ(str_views[i], str_data[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -609,7 +609,6 @@ make_chunk(const FieldMeta& field_meta,
|
||||
size_t row_nums,
|
||||
char* data,
|
||||
size_t size,
|
||||
const std::string& file_path,
|
||||
std::shared_ptr<ChunkMmapGuard> chunk_mmap_guard) {
|
||||
int dim = IsVectorDataType(field_meta.get_data_type()) &&
|
||||
!IsSparseFloatVectorDataType(field_meta.get_data_type())
|
||||
@@ -756,12 +755,12 @@ make_chunk(const FieldMeta& field_meta,
|
||||
}
|
||||
}
|
||||
|
||||
std::unique_ptr<Chunk>
|
||||
create_chunk(const FieldMeta& field_meta,
|
||||
const arrow::ArrayVector& array_vec,
|
||||
const std::string& file_path,
|
||||
bool mmap_populate,
|
||||
proto::common::LoadPriority load_priority) {
|
||||
ChunkBuffer
|
||||
create_chunk_buffer(const FieldMeta& field_meta,
|
||||
const arrow::ArrayVector& array_vec,
|
||||
bool mmap_populate,
|
||||
const std::string& file_path,
|
||||
proto::common::LoadPriority load_priority) {
|
||||
auto cw = create_chunk_writer(field_meta);
|
||||
auto [size, row_nums] = cw->calculate_size(array_vec);
|
||||
size_t aligned_size = (size + ChunkTarget::ALIGNED_SIZE - 1) &
|
||||
@@ -783,22 +782,47 @@ create_chunk(const FieldMeta& field_meta,
|
||||
} else {
|
||||
chunk_mmap_guard = std::make_shared<ChunkMmapGuard>(data, size, "");
|
||||
}
|
||||
ChunkBuffer buffer;
|
||||
buffer.data = data;
|
||||
buffer.size = size;
|
||||
buffer.row_nums = row_nums;
|
||||
buffer.guard = std::move(chunk_mmap_guard);
|
||||
return buffer;
|
||||
}
|
||||
|
||||
std::unique_ptr<Chunk>
|
||||
make_chunk_from_buffer(const FieldMeta& field_meta,
|
||||
const ChunkBuffer& buffer,
|
||||
size_t row_nums_override) {
|
||||
auto row_nums =
|
||||
row_nums_override == 0 ? buffer.row_nums : row_nums_override;
|
||||
return make_chunk(
|
||||
field_meta, row_nums, data, size, file_path, chunk_mmap_guard);
|
||||
field_meta, row_nums, buffer.data, buffer.size, buffer.guard);
|
||||
}
|
||||
|
||||
std::unique_ptr<Chunk>
|
||||
create_chunk(const FieldMeta& field_meta,
|
||||
const arrow::ArrayVector& array_vec,
|
||||
bool mmap_populate,
|
||||
const std::string& file_path,
|
||||
proto::common::LoadPriority load_priority) {
|
||||
auto buffer = create_chunk_buffer(
|
||||
field_meta, array_vec, mmap_populate, file_path, load_priority);
|
||||
return make_chunk_from_buffer(field_meta, buffer, 0);
|
||||
}
|
||||
|
||||
std::unordered_map<FieldId, std::shared_ptr<Chunk>>
|
||||
create_group_chunk(const std::vector<FieldId>& field_ids,
|
||||
const std::vector<FieldMeta>& field_metas,
|
||||
const std::vector<arrow::ArrayVector>& array_vec,
|
||||
const std::string& file_path,
|
||||
bool mmap_populate,
|
||||
const std::string& file_path,
|
||||
proto::common::LoadPriority load_priority) {
|
||||
std::vector<std::shared_ptr<ChunkWriterBase>> cws;
|
||||
cws.reserve(field_ids.size());
|
||||
size_t total_aligned_size = 0, final_row_nums = 0;
|
||||
for (size_t i = 0; i < field_ids.size(); i++) {
|
||||
auto field_meta = field_metas[i];
|
||||
const auto& field_meta = field_metas[i];
|
||||
cws.push_back(create_chunk_writer(field_meta));
|
||||
}
|
||||
std::vector<size_t> chunk_sizes;
|
||||
@@ -876,7 +900,6 @@ create_group_chunk(const std::vector<FieldId>& field_ids,
|
||||
final_row_nums,
|
||||
data + chunk_offsets[i],
|
||||
chunk_sizes[i],
|
||||
file_path,
|
||||
chunk_mmap_guard));
|
||||
LOG_INFO(
|
||||
"created chunk for field {} with chunk offset: {}, chunk "
|
||||
@@ -900,4 +923,4 @@ read_single_column_batches(std::shared_ptr<arrow::RecordBatchReader> reader) {
|
||||
return array_vec;
|
||||
}
|
||||
|
||||
} // namespace milvus
|
||||
} // namespace milvus
|
||||
|
||||
@@ -20,6 +20,7 @@
|
||||
#include "common/ChunkTarget.h"
|
||||
#include "arrow/record_batch.h"
|
||||
#include "common/Chunk.h"
|
||||
#include "pb/common.pb.h"
|
||||
|
||||
namespace milvus {
|
||||
class ChunkWriterBase {
|
||||
@@ -290,11 +291,40 @@ class SparseFloatVectorChunkWriter : public ChunkWriterBase {
|
||||
const std::shared_ptr<ChunkTarget>& target) override;
|
||||
};
|
||||
|
||||
// A reusable buffer that holds the raw chunk memory and its mmap guard.
|
||||
// This can be used to create multiple Chunk instances that share the same
|
||||
// underlying memory.
|
||||
struct ChunkBuffer {
|
||||
char* data{nullptr};
|
||||
size_t size{0};
|
||||
size_t row_nums{0};
|
||||
std::shared_ptr<ChunkMmapGuard> guard;
|
||||
};
|
||||
|
||||
// Build a chunk buffer from Arrow arrays, but do not materialize the Chunk
|
||||
// object yet. This is useful when multiple Chunk instances need to share
|
||||
// the same underlying memory.
|
||||
ChunkBuffer
|
||||
create_chunk_buffer(const FieldMeta& field_meta,
|
||||
const arrow::ArrayVector& array_vec,
|
||||
bool mmap_populate = true,
|
||||
const std::string& file_path = "",
|
||||
proto::common::LoadPriority load_priority =
|
||||
proto::common::LoadPriority::HIGH);
|
||||
|
||||
// Create a Chunk view from an existing ChunkBuffer. Multiple Chunk instances
|
||||
// created from the same buffer will share the same underlying memory via
|
||||
// the shared ChunkMmapGuard in the buffer.
|
||||
std::unique_ptr<Chunk>
|
||||
make_chunk_from_buffer(const FieldMeta& field_meta,
|
||||
const ChunkBuffer& buffer,
|
||||
size_t row_nums_override = 0);
|
||||
|
||||
std::unique_ptr<Chunk>
|
||||
create_chunk(const FieldMeta& field_meta,
|
||||
const arrow::ArrayVector& array_vec,
|
||||
bool mmap_populate = true,
|
||||
const std::string& file_path = "",
|
||||
bool mmap_populate_ = true,
|
||||
proto::common::LoadPriority load_priority =
|
||||
proto::common::LoadPriority::HIGH);
|
||||
|
||||
@@ -302,8 +332,8 @@ std::unordered_map<FieldId, std::shared_ptr<Chunk>>
|
||||
create_group_chunk(const std::vector<FieldId>& field_ids,
|
||||
const std::vector<FieldMeta>& field_metas,
|
||||
const std::vector<arrow::ArrayVector>& array_vec,
|
||||
bool mmap_populate = true,
|
||||
const std::string& file_path = "",
|
||||
bool mmap_populate_ = true,
|
||||
proto::common::LoadPriority load_priority =
|
||||
proto::common::LoadPriority::HIGH);
|
||||
|
||||
|
||||
@@ -308,7 +308,7 @@ TEST_F(VectorArrayChunkTest, TestWriteWithMmap) {
|
||||
DataType::VECTOR_FLOAT,
|
||||
dim,
|
||||
std::nullopt);
|
||||
auto chunk = create_chunk(field_meta, array_vec, temp_file);
|
||||
auto chunk = create_chunk(field_meta, array_vec, true, temp_file);
|
||||
auto vector_array_chunk = static_cast<VectorArrayChunk*>(chunk.get());
|
||||
|
||||
// Verify mmap write
|
||||
|
||||
@@ -38,7 +38,7 @@ using namespace milvus;
|
||||
using namespace milvus::storage;
|
||||
|
||||
std::shared_ptr<Chunk>
|
||||
create_chunk(const FixedVector<int64_t>& data) {
|
||||
create_chunk_int64(const FixedVector<int64_t>& data) {
|
||||
auto field_data = milvus::storage::CreateFieldData(storage::DataType::INT64,
|
||||
DataType::NONE);
|
||||
field_data->FillFieldData(data.data(), data.size());
|
||||
@@ -73,7 +73,7 @@ create_chunk(const FixedVector<int64_t>& data) {
|
||||
|
||||
// Helper function to create chunks for string data
|
||||
std::shared_ptr<Chunk>
|
||||
create_chunk(const FixedVector<std::string>& data) {
|
||||
create_chunk_string(const FixedVector<std::string>& data) {
|
||||
auto field_data = milvus::storage::CreateFieldData(
|
||||
storage::DataType::VARCHAR, DataType::NONE);
|
||||
field_data->FillFieldData(data.data(), data.size());
|
||||
@@ -139,8 +139,8 @@ class ChunkedColumnGroupTest : public ::testing::Test {
|
||||
std::nullopt);
|
||||
|
||||
// Create chunks
|
||||
int64_chunk = std::move(create_chunk(int64_data));
|
||||
string_chunk = std::move(create_chunk(string_data));
|
||||
int64_chunk = std::move(create_chunk_int64(int64_data));
|
||||
string_chunk = std::move(create_chunk_string(string_data));
|
||||
}
|
||||
|
||||
FixedVector<int64_t> int64_data;
|
||||
@@ -177,12 +177,13 @@ TEST_F(ChunkedColumnGroupTest, GroupChunk) {
|
||||
EXPECT_EQ(all_chunks[FieldId(2)], string_chunk);
|
||||
|
||||
// Add chunk
|
||||
auto new_int64_chunk = create_chunk(FixedVector<int64_t>{6, 7, 8, 9, 10});
|
||||
auto new_int64_chunk =
|
||||
create_chunk_int64(FixedVector<int64_t>{6, 7, 8, 9, 10});
|
||||
EXPECT_NO_THROW(group_chunk->AddChunk(FieldId(3), new_int64_chunk));
|
||||
EXPECT_TRUE(group_chunk->HasChunk(FieldId(3)));
|
||||
EXPECT_EQ(group_chunk->GetChunk(FieldId(3))->RowNums(), 5);
|
||||
auto another_int64_chunk =
|
||||
create_chunk(FixedVector<int64_t>{11, 12, 13, 14, 15});
|
||||
create_chunk_int64(FixedVector<int64_t>{11, 12, 13, 14, 15});
|
||||
EXPECT_THROW(group_chunk->AddChunk(FieldId(3), another_int64_chunk),
|
||||
std::exception);
|
||||
|
||||
|
||||
@@ -2883,12 +2883,27 @@ ChunkedSegmentSealedImpl::fill_empty_field(const FieldMeta& field_meta) {
|
||||
data_type,
|
||||
field_id.get(),
|
||||
id_);
|
||||
auto [field_has_setting, field_mmap_enabled] =
|
||||
schema_->MmapEnabled(field_id);
|
||||
auto is_vector = IsVectorDataType(field_meta.get_data_type());
|
||||
auto& mmap_config = storage::MmapManager::GetInstance().GetMmapConfig();
|
||||
bool global_use_mmap = is_vector ? mmap_config.GetVectorFieldEnableMmap()
|
||||
: mmap_config.GetScalarFieldEnableMmap();
|
||||
bool use_mmap = field_has_setting ? field_mmap_enabled : global_use_mmap;
|
||||
auto mmap_dir_path =
|
||||
milvus::storage::LocalChunkManagerSingleton::GetInstance()
|
||||
.GetChunkManager()
|
||||
->GetRootPath();
|
||||
int64_t size = num_rows_.value();
|
||||
AssertInfo(size > 0, "Chunked Sealed segment must have more than 0 row");
|
||||
auto field_data_info = FieldDataInfo(field_id.get(), size, "");
|
||||
auto field_data_info = FieldDataInfo(field_id.get(), size, mmap_dir_path);
|
||||
std::unique_ptr<Translator<milvus::Chunk>> translator =
|
||||
std::make_unique<storagev1translator::DefaultValueChunkTranslator>(
|
||||
get_segment_id(), field_meta, field_data_info, false);
|
||||
get_segment_id(),
|
||||
field_meta,
|
||||
field_data_info,
|
||||
use_mmap,
|
||||
mmap_config.GetMmapPopulate());
|
||||
auto column =
|
||||
MakeChunkedColumnBase(data_type, std::move(translator), field_meta);
|
||||
|
||||
|
||||
@@ -205,8 +205,8 @@ ChunkTranslator::get_cells(
|
||||
read_single_column_batches(r->reader);
|
||||
chunk = create_chunk(field_meta_,
|
||||
array_vec,
|
||||
filepath.string(),
|
||||
mmap_populate_,
|
||||
filepath.string(),
|
||||
load_priority_);
|
||||
}
|
||||
cells.emplace_back(cid, std::move(chunk));
|
||||
|
||||
@@ -22,14 +22,21 @@ DefaultValueChunkTranslator::DefaultValueChunkTranslator(
|
||||
int64_t segment_id,
|
||||
FieldMeta field_meta,
|
||||
FieldDataInfo field_data_info,
|
||||
bool use_mmap)
|
||||
: segment_id_(segment_id),
|
||||
key_(fmt::format("seg_{}_f_{}", segment_id, field_data_info.field_id)),
|
||||
bool use_mmap,
|
||||
bool mmap_populate)
|
||||
: total_rows_(field_data_info.row_count),
|
||||
segment_id_(segment_id),
|
||||
key_(
|
||||
fmt::format("seg_{}_f_{}_def", segment_id, field_data_info.field_id)),
|
||||
use_mmap_(use_mmap),
|
||||
mmap_populate_(mmap_populate),
|
||||
mmap_dir_path_(field_data_info.mmap_dir_path),
|
||||
field_meta_(field_meta),
|
||||
meta_(use_mmap ? milvus::cachinglayer::StorageType::DISK
|
||||
: milvus::cachinglayer::StorageType::MEMORY,
|
||||
milvus::cachinglayer::CellIdMappingMode::ALWAYS_ZERO,
|
||||
// For default-value fields, one logical chunk per caching cell.
|
||||
// Cell IDs are identical to chunk IDs.
|
||||
milvus::cachinglayer::CellIdMappingMode::IDENTICAL,
|
||||
milvus::segcore::getCellDataType(
|
||||
IsVectorDataType(field_meta.get_data_type()),
|
||||
/* is_index */ false),
|
||||
@@ -38,13 +45,84 @@ DefaultValueChunkTranslator::DefaultValueChunkTranslator(
|
||||
/* is_index */ false,
|
||||
/* in_load_list, set to false to reduce memory usage */ false),
|
||||
/* support_eviction */ false) {
|
||||
// Split rows into ~64KB cells according to value_size().
|
||||
// Fallback to single-cell if value_size() is not well-defined.
|
||||
auto vsize = this->value_size();
|
||||
int64_t rows_per_cell = total_rows_;
|
||||
if (vsize > 0) {
|
||||
rows_per_cell = std::max<int64_t>(
|
||||
1, kTargetCellBytes / static_cast<int64_t>(vsize));
|
||||
}
|
||||
// primary_cell_rows_ is the standard cell row count, but should not exceed
|
||||
// total_rows_.
|
||||
primary_cell_rows_ = std::min(rows_per_cell, total_rows_);
|
||||
|
||||
meta_.num_rows_until_chunk_.clear();
|
||||
meta_.num_rows_until_chunk_.reserve(
|
||||
static_cast<size_t>(total_rows_ / rows_per_cell) + 2);
|
||||
meta_.num_rows_until_chunk_.push_back(0);
|
||||
meta_.num_rows_until_chunk_.push_back(field_data_info.row_count);
|
||||
virtual_chunk_config(field_data_info.row_count,
|
||||
1,
|
||||
while (meta_.num_rows_until_chunk_.back() < total_rows_) {
|
||||
auto prev = meta_.num_rows_until_chunk_.back();
|
||||
auto remain = total_rows_ - prev;
|
||||
auto this_rows = std::min(remain, rows_per_cell);
|
||||
meta_.num_rows_until_chunk_.push_back(prev + this_rows);
|
||||
}
|
||||
|
||||
auto nr_chunks =
|
||||
static_cast<int64_t>(meta_.num_rows_until_chunk_.size() - 1);
|
||||
virtual_chunk_config(total_rows_,
|
||||
nr_chunks,
|
||||
meta_.num_rows_until_chunk_,
|
||||
meta_.virt_chunk_order_,
|
||||
meta_.vcid_to_cid_arr_);
|
||||
|
||||
// Pre-build shared buffer for default-value cells: all cells, including
|
||||
// the tail one, will share this buffer. Tail cells will use a smaller
|
||||
// logical row count while reusing the same underlying memory.
|
||||
auto build_buffer_for_rows = [&](int64_t num_rows) -> milvus::ChunkBuffer {
|
||||
auto data_type = field_meta_.get_data_type();
|
||||
std::shared_ptr<arrow::ArrayBuilder> builder;
|
||||
if (IsVectorDataType(data_type)) {
|
||||
AssertInfo(field_meta_.is_nullable(),
|
||||
"only nullable vector fields can be dynamically added");
|
||||
builder = std::make_shared<arrow::BinaryBuilder>();
|
||||
} else {
|
||||
builder = milvus::storage::CreateArrowBuilder(data_type);
|
||||
}
|
||||
arrow::Status ast;
|
||||
if (field_meta_.default_value().has_value()) {
|
||||
ast = builder->Reserve(num_rows);
|
||||
AssertInfo(
|
||||
ast.ok(), "reserve arrow builder failed: {}", ast.ToString());
|
||||
auto default_scalar =
|
||||
storage::CreateArrowScalarFromDefaultValue(field_meta_);
|
||||
ast = builder->AppendScalar(*default_scalar, num_rows);
|
||||
} else {
|
||||
ast = builder->AppendNulls(num_rows);
|
||||
}
|
||||
AssertInfo(ast.ok(),
|
||||
"append null/default values to arrow builder failed: {}",
|
||||
ast.ToString());
|
||||
arrow::ArrayVector array_vec;
|
||||
array_vec.emplace_back(builder->Finish().ValueOrDie());
|
||||
if (!use_mmap_ || mmap_dir_path_.empty()) {
|
||||
return milvus::create_chunk_buffer(
|
||||
field_meta_, array_vec, mmap_populate_);
|
||||
} else {
|
||||
auto filepath =
|
||||
std::filesystem::path(mmap_dir_path_) /
|
||||
fmt::format(
|
||||
"seg_{}_f_{}_def", segment_id_, field_data_info.field_id);
|
||||
std::filesystem::create_directories(filepath.parent_path());
|
||||
// just use default load priority: proto::common::LoadPriority::HIGH
|
||||
return milvus::create_chunk_buffer(
|
||||
field_meta_, array_vec, mmap_populate_, filepath.string());
|
||||
}
|
||||
};
|
||||
|
||||
if (primary_cell_rows_ > 0) {
|
||||
primary_buffer_ = build_buffer_for_rows(primary_cell_rows_);
|
||||
}
|
||||
}
|
||||
|
||||
DefaultValueChunkTranslator::~DefaultValueChunkTranslator() {
|
||||
@@ -52,18 +130,18 @@ DefaultValueChunkTranslator::~DefaultValueChunkTranslator() {
|
||||
|
||||
size_t
|
||||
DefaultValueChunkTranslator::num_cells() const {
|
||||
return 1;
|
||||
return meta_.num_rows_until_chunk_.size() > 0
|
||||
? meta_.num_rows_until_chunk_.size() - 1
|
||||
: 0;
|
||||
}
|
||||
|
||||
milvus::cachinglayer::cid_t
|
||||
DefaultValueChunkTranslator::cell_id_of(milvus::cachinglayer::uid_t uid) const {
|
||||
return 0;
|
||||
return uid;
|
||||
}
|
||||
|
||||
std::pair<milvus::cachinglayer::ResourceUsage,
|
||||
milvus::cachinglayer::ResourceUsage>
|
||||
DefaultValueChunkTranslator::estimated_byte_size_of_cell(
|
||||
milvus::cachinglayer::cid_t cid) const {
|
||||
int64_t
|
||||
DefaultValueChunkTranslator::value_size() const {
|
||||
int64_t value_size = 0;
|
||||
switch (field_meta_.get_data_type()) {
|
||||
case milvus::DataType::BOOL:
|
||||
@@ -122,8 +200,25 @@ DefaultValueChunkTranslator::estimated_byte_size_of_cell(
|
||||
"unsupported default value data type {}",
|
||||
field_meta_.get_data_type());
|
||||
}
|
||||
return {{value_size * meta_.num_rows_until_chunk_[1], 0},
|
||||
{2 * value_size * meta_.num_rows_until_chunk_[1], 0}};
|
||||
return value_size;
|
||||
}
|
||||
|
||||
std::pair<milvus::cachinglayer::ResourceUsage,
|
||||
milvus::cachinglayer::ResourceUsage>
|
||||
DefaultValueChunkTranslator::estimated_byte_size_of_cell(
|
||||
milvus::cachinglayer::cid_t cid) const {
|
||||
// TODO: actually only the first cell is used, other cells share the same buffer,
|
||||
// but for now we estimate the same size for all cells
|
||||
auto value_size = this->value_size();
|
||||
auto rows_begin = meta_.num_rows_until_chunk_[cid];
|
||||
auto rows_end = meta_.num_rows_until_chunk_[cid + 1];
|
||||
auto rows = rows_end - rows_begin;
|
||||
auto cell_bytes = value_size * rows;
|
||||
if (use_mmap_) {
|
||||
return {{0, cell_bytes}, {0, cell_bytes}};
|
||||
} else {
|
||||
return {{cell_bytes, 0}, {cell_bytes, 0}};
|
||||
}
|
||||
}
|
||||
|
||||
const std::string&
|
||||
@@ -135,39 +230,28 @@ std::vector<
|
||||
std::pair<milvus::cachinglayer::cid_t, std::unique_ptr<milvus::Chunk>>>
|
||||
DefaultValueChunkTranslator::get_cells(
|
||||
const std::vector<milvus::cachinglayer::cid_t>& cids) {
|
||||
AssertInfo(cids.size() == 1 && cids[0] == 0,
|
||||
"DefaultValueChunkTranslator only supports one cell");
|
||||
auto num_rows = meta_.num_rows_until_chunk_[1];
|
||||
auto data_type = field_meta_.get_data_type();
|
||||
std::shared_ptr<arrow::ArrayBuilder> builder;
|
||||
if (IsVectorDataType(data_type)) {
|
||||
AssertInfo(field_meta_.is_nullable(),
|
||||
"only nullable vector fields can be dynamically added");
|
||||
builder = std::make_shared<arrow::BinaryBuilder>();
|
||||
} else {
|
||||
builder = milvus::storage::CreateArrowBuilder(data_type);
|
||||
}
|
||||
arrow::Status ast;
|
||||
if (field_meta_.default_value().has_value()) {
|
||||
ast = builder->Reserve(num_rows);
|
||||
AssertInfo(ast.ok(), "reserve arrow build failed: {}", ast.ToString());
|
||||
auto scalar = storage::CreateArrowScalarFromDefaultValue(field_meta_);
|
||||
ast = builder->AppendScalar(*scalar, num_rows);
|
||||
} else {
|
||||
ast = builder->AppendNulls(num_rows);
|
||||
}
|
||||
AssertInfo(ast.ok(),
|
||||
"append null/default values to arrow builder failed: {}",
|
||||
ast.ToString());
|
||||
arrow::ArrayVector array_vec;
|
||||
array_vec.emplace_back(builder->Finish().ValueOrDie());
|
||||
auto chunk = create_chunk(field_meta_, array_vec);
|
||||
AssertInfo(primary_buffer_.has_value(),
|
||||
"primary buffer is not initialized");
|
||||
|
||||
std::vector<
|
||||
std::pair<milvus::cachinglayer::cid_t, std::unique_ptr<milvus::Chunk>>>
|
||||
res;
|
||||
res.reserve(1);
|
||||
res.emplace_back(0, std::move(chunk));
|
||||
res.reserve(cids.size());
|
||||
|
||||
for (auto cid : cids) {
|
||||
assert(cid + 1 < meta_.num_rows_until_chunk_.size());
|
||||
|
||||
auto rows_begin = meta_.num_rows_until_chunk_[cid];
|
||||
auto rows_end = meta_.num_rows_until_chunk_[cid + 1];
|
||||
auto num_rows = rows_end - rows_begin;
|
||||
|
||||
const milvus::ChunkBuffer& buffer = primary_buffer_.value();
|
||||
|
||||
auto chunk =
|
||||
milvus::make_chunk_from_buffer(field_meta_, buffer, num_rows);
|
||||
res.emplace_back(cid, std::move(chunk));
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
|
||||
@@ -13,11 +13,13 @@
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <optional>
|
||||
|
||||
#include "cachinglayer/Translator.h"
|
||||
#include "cachinglayer/Utils.h"
|
||||
#include "common/Chunk.h"
|
||||
#include "common/FieldMeta.h"
|
||||
#include "common/ChunkWriter.h"
|
||||
#include "mmap/Types.h"
|
||||
#include "segcore/storagev1translator/ChunkTranslator.h"
|
||||
|
||||
@@ -30,7 +32,8 @@ class DefaultValueChunkTranslator
|
||||
DefaultValueChunkTranslator(int64_t segment_id,
|
||||
FieldMeta field_meta,
|
||||
FieldDataInfo field_data_info,
|
||||
bool use_mmap);
|
||||
bool use_mmap,
|
||||
bool mmap_populate);
|
||||
|
||||
~DefaultValueChunkTranslator() override;
|
||||
|
||||
@@ -58,10 +61,29 @@ class DefaultValueChunkTranslator
|
||||
return 0;
|
||||
}
|
||||
|
||||
int64_t
|
||||
value_size() const;
|
||||
|
||||
// preferred bytes per cell when splitting default-value column
|
||||
static constexpr int64_t kTargetCellBytes = 64 * 1024; // 64KB
|
||||
|
||||
private:
|
||||
// total rows of this field in the segment
|
||||
int64_t total_rows_{0};
|
||||
|
||||
// Number of rows in the primary cell (all full-sized cells).
|
||||
// The last cell may contain fewer rows.
|
||||
int64_t primary_cell_rows_{0};
|
||||
|
||||
// Shared chunk buffers for default-value cells. All cells with the same
|
||||
// row count will share the same underlying memory via these buffers.
|
||||
std::optional<milvus::ChunkBuffer> primary_buffer_;
|
||||
|
||||
int64_t segment_id_;
|
||||
std::string key_;
|
||||
bool use_mmap_;
|
||||
bool mmap_populate_;
|
||||
std::string mmap_dir_path_;
|
||||
CTMeta meta_;
|
||||
milvus::FieldMeta field_meta_;
|
||||
};
|
||||
|
||||
@@ -0,0 +1,856 @@
|
||||
// Copyright (C) 2019-2025 Zilliz. All rights reserved.
|
||||
//
|
||||
// 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 <gtest/gtest.h>
|
||||
#include <fmt/core.h>
|
||||
#include <filesystem>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#include "common/Chunk.h"
|
||||
#include "common/FieldMeta.h"
|
||||
#include "common/Types.h"
|
||||
#include "mmap/Types.h"
|
||||
#include "segcore/storagev1translator/DefaultValueChunkTranslator.h"
|
||||
#include "pb/schema.pb.h"
|
||||
|
||||
using namespace milvus;
|
||||
using namespace milvus::segcore::storagev1translator;
|
||||
|
||||
class DefaultValueChunkTranslatorTest : public ::testing::TestWithParam<bool> {
|
||||
protected:
|
||||
void
|
||||
SetUp() override {
|
||||
// Create a unique temp directory for mmap tests
|
||||
temp_dir_ = std::filesystem::temp_directory_path() /
|
||||
("milvus_param_test_" + std::to_string(segment_id_) + "_" +
|
||||
std::to_string(reinterpret_cast<uintptr_t>(this)));
|
||||
std::filesystem::create_directories(temp_dir_);
|
||||
}
|
||||
|
||||
void
|
||||
TearDown() override {
|
||||
// Clean up temp directory
|
||||
if (std::filesystem::exists(temp_dir_)) {
|
||||
std::filesystem::remove_all(temp_dir_);
|
||||
}
|
||||
}
|
||||
|
||||
// Helper to get mmap_dir_path based on use_mmap parameter
|
||||
std::string
|
||||
getMmapDirPath() const {
|
||||
return GetParam() ? temp_dir_.string() : "";
|
||||
}
|
||||
|
||||
std::filesystem::path temp_dir_;
|
||||
int64_t segment_id_ = 12345;
|
||||
};
|
||||
|
||||
// Test basic int64 field with default value
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestInt64WithDefaultValue) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 1000;
|
||||
int64_t default_value = 42;
|
||||
|
||||
// Create field meta with default value
|
||||
DefaultValueType value_field;
|
||||
value_field.set_long_data(default_value);
|
||||
FieldMeta field_meta(FieldName("test_int64"),
|
||||
FieldId(101),
|
||||
DataType::INT64,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(101, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
// Test num_cells
|
||||
EXPECT_GT(translator->num_cells(), 0);
|
||||
|
||||
// Test cell_id_of - should be identical mapping
|
||||
for (size_t i = 0; i < translator->num_cells(); ++i) {
|
||||
EXPECT_EQ(translator->cell_id_of(i), i);
|
||||
}
|
||||
|
||||
// Test estimated_byte_size_of_cell
|
||||
auto [usage, peak_usage] = translator->estimated_byte_size_of_cell(0);
|
||||
if (use_mmap) {
|
||||
EXPECT_GT(usage.file_bytes, 0);
|
||||
} else {
|
||||
EXPECT_GT(usage.memory_bytes, 0);
|
||||
}
|
||||
|
||||
// Test get_cells
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
EXPECT_EQ(cells.size(), 1);
|
||||
|
||||
auto& [cid, chunk] = cells[0];
|
||||
EXPECT_EQ(cid, 0);
|
||||
ASSERT_NE(chunk, nullptr);
|
||||
|
||||
// Verify the chunk contains default values
|
||||
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
EXPECT_GT(span.row_count(), 0);
|
||||
|
||||
for (size_t i = 0; i < span.row_count(); ++i) {
|
||||
auto value = *reinterpret_cast<int64_t*>((char*)span.data() +
|
||||
i * span.element_sizeof());
|
||||
EXPECT_EQ(value, default_value);
|
||||
}
|
||||
}
|
||||
|
||||
// Test int64 field without default value (nulls)
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestInt64WithoutDefaultValue) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 500;
|
||||
|
||||
FieldMeta field_meta(FieldName("test_int64_null"),
|
||||
FieldId(102),
|
||||
DataType::INT64,
|
||||
true,
|
||||
std::nullopt);
|
||||
|
||||
FieldDataInfo field_data_info(102, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_GT(translator->num_cells(), 0);
|
||||
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
EXPECT_EQ(cells.size(), 1);
|
||||
|
||||
auto& [cid, chunk] = cells[0];
|
||||
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
EXPECT_GT(span.row_count(), 0);
|
||||
|
||||
// All values should be marked as invalid (null)
|
||||
for (size_t i = 0; i < span.row_count(); ++i) {
|
||||
EXPECT_FALSE(fixed_chunk->isValid(i));
|
||||
}
|
||||
}
|
||||
|
||||
// Test various fixed-width data types with default values
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestVariousFixedWidthTypes) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 100;
|
||||
|
||||
// Test BOOL
|
||||
{
|
||||
DefaultValueType value_field;
|
||||
value_field.set_bool_data(true);
|
||||
FieldMeta field_meta(FieldName("test_bool"),
|
||||
FieldId(201),
|
||||
DataType::BOOL,
|
||||
false,
|
||||
value_field);
|
||||
FieldDataInfo field_data_info(201, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_GT(translator->num_cells(), 0);
|
||||
EXPECT_GT(translator->value_size(), 0);
|
||||
}
|
||||
|
||||
// Test INT32
|
||||
{
|
||||
DefaultValueType value_field;
|
||||
value_field.set_int_data(100);
|
||||
FieldMeta field_meta(FieldName("test_int32"),
|
||||
FieldId(202),
|
||||
DataType::INT32,
|
||||
false,
|
||||
value_field);
|
||||
FieldDataInfo field_data_info(202, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_EQ(translator->value_size(), sizeof(int32_t));
|
||||
}
|
||||
|
||||
// Test FLOAT
|
||||
{
|
||||
DefaultValueType value_field;
|
||||
value_field.set_float_data(3.14f);
|
||||
FieldMeta field_meta(FieldName("test_float"),
|
||||
FieldId(203),
|
||||
DataType::FLOAT,
|
||||
false,
|
||||
value_field);
|
||||
FieldDataInfo field_data_info(203, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_EQ(translator->value_size(), sizeof(float));
|
||||
}
|
||||
|
||||
// Test DOUBLE
|
||||
{
|
||||
DefaultValueType value_field;
|
||||
value_field.set_double_data(2.718281828);
|
||||
FieldMeta field_meta(FieldName("test_double"),
|
||||
FieldId(204),
|
||||
DataType::DOUBLE,
|
||||
false,
|
||||
value_field);
|
||||
FieldDataInfo field_data_info(204, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_EQ(translator->value_size(), sizeof(double));
|
||||
}
|
||||
}
|
||||
|
||||
// Test VARCHAR/STRING with default value
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestStringWithDefaultValue) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 200;
|
||||
std::string default_string = "default_value";
|
||||
|
||||
DefaultValueType value_field;
|
||||
value_field.set_string_data(default_string);
|
||||
FieldMeta field_meta(FieldName("test_string"),
|
||||
FieldId(301),
|
||||
DataType::STRING,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(301, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_GT(translator->num_cells(), 0);
|
||||
EXPECT_GT(translator->value_size(), 0);
|
||||
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
EXPECT_EQ(cells.size(), 1);
|
||||
|
||||
auto& [cid, chunk] = cells[0];
|
||||
auto string_chunk = static_cast<StringChunk*>(chunk.get());
|
||||
auto [views, valid] = string_chunk->StringViews(std::nullopt);
|
||||
|
||||
EXPECT_GT(views.size(), 0);
|
||||
for (const auto& view : views) {
|
||||
EXPECT_EQ(view, default_string);
|
||||
}
|
||||
}
|
||||
|
||||
// Test VARCHAR without default value
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestStringWithoutDefaultValue) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 150;
|
||||
|
||||
FieldMeta field_meta(FieldName("test_string_null"),
|
||||
FieldId(302),
|
||||
DataType::VARCHAR,
|
||||
true,
|
||||
std::nullopt);
|
||||
|
||||
FieldDataInfo field_data_info(302, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_GT(translator->num_cells(), 0);
|
||||
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
EXPECT_EQ(cells.size(), 1);
|
||||
|
||||
auto& [cid, chunk] = cells[0];
|
||||
auto string_chunk = static_cast<StringChunk*>(chunk.get());
|
||||
auto [views, valid] = string_chunk->StringViews(std::nullopt);
|
||||
|
||||
EXPECT_GT(views.size(), 0);
|
||||
// All should be marked as invalid (null)
|
||||
for (const auto& v : valid) {
|
||||
EXPECT_FALSE(v);
|
||||
}
|
||||
}
|
||||
|
||||
// Test large row count that requires multiple cells
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestMultipleCells) {
|
||||
bool use_mmap = GetParam();
|
||||
// Use a large row count to ensure multiple cells are created
|
||||
int64_t row_count = 10000000; // 10 million rows
|
||||
|
||||
DefaultValueType value_field;
|
||||
value_field.set_long_data(999);
|
||||
FieldMeta field_meta(FieldName("test_large"),
|
||||
FieldId(401),
|
||||
DataType::INT64,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(401, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
// With 10M rows and int64 (8 bytes), we expect multiple cells
|
||||
// since target cell size is 64KB
|
||||
EXPECT_GT(translator->num_cells(), 1);
|
||||
|
||||
// Test multiple cells
|
||||
std::vector<cachinglayer::cid_t> cids;
|
||||
for (size_t i = 0; i < std::min<size_t>(3, translator->num_cells()); ++i) {
|
||||
cids.push_back(i);
|
||||
}
|
||||
|
||||
auto cells = translator->get_cells(cids);
|
||||
EXPECT_EQ(cells.size(), cids.size());
|
||||
|
||||
int64_t total_rows = 0;
|
||||
for (const auto& [cid, chunk] : cells) {
|
||||
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
total_rows += span.row_count();
|
||||
|
||||
// Verify all values are the default
|
||||
for (size_t i = 0; i < span.row_count(); ++i) {
|
||||
auto value =
|
||||
*(int64_t*)((char*)span.data() + i * span.element_sizeof());
|
||||
EXPECT_EQ(value, 999);
|
||||
}
|
||||
}
|
||||
|
||||
EXPECT_GT(total_rows, 0);
|
||||
}
|
||||
|
||||
// Test small row count (single cell)
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestSmallRowCount) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 10;
|
||||
|
||||
DefaultValueType value_field;
|
||||
value_field.set_int_data(55);
|
||||
FieldMeta field_meta(FieldName("test_small"),
|
||||
FieldId(501),
|
||||
DataType::INT32,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(501, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
// Small row count should result in a single cell
|
||||
EXPECT_EQ(translator->num_cells(), 1);
|
||||
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
EXPECT_EQ(cells.size(), 1);
|
||||
|
||||
auto& [cid, chunk] = cells[0];
|
||||
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
EXPECT_EQ(span.row_count(), row_count);
|
||||
}
|
||||
|
||||
// Test cells_storage_bytes
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestCellsStorageBytes) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 500;
|
||||
|
||||
proto::schema::ValueField value_field;
|
||||
value_field.set_long_data(777);
|
||||
FieldMeta field_meta(FieldName("test_storage_bytes"),
|
||||
FieldId(701),
|
||||
DataType::INT64,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(701, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
// cells_storage_bytes always returns 0 for default-value chunks
|
||||
EXPECT_EQ(translator->cells_storage_bytes(cids), 0);
|
||||
}
|
||||
|
||||
// Test get_cells with multiple cell IDs
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestGetMultipleCells) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count =
|
||||
DefaultValueChunkTranslator::kTargetCellBytes / sizeof(float) * 3;
|
||||
|
||||
DefaultValueType value_field;
|
||||
value_field.set_float_data(1.5f);
|
||||
FieldMeta field_meta(FieldName("test_multi_cells"),
|
||||
FieldId(801),
|
||||
DataType::FLOAT,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(801, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
size_t num_cells = translator->num_cells();
|
||||
if (num_cells > 1) {
|
||||
std::vector<cachinglayer::cid_t> cids;
|
||||
for (size_t i = 0; i < num_cells; ++i) {
|
||||
cids.push_back(i);
|
||||
}
|
||||
|
||||
auto cells = translator->get_cells(cids);
|
||||
EXPECT_EQ(cells.size(), cids.size());
|
||||
|
||||
for (size_t i = 0; i < cells.size(); ++i) {
|
||||
EXPECT_EQ(cells[i].first, cids[i]);
|
||||
ASSERT_NE(cells[i].second, nullptr);
|
||||
auto fixed_chunk =
|
||||
static_cast<FixedWidthChunk*>(cells[i].second.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
EXPECT_GT(span.row_count(), 0);
|
||||
for (size_t j = 0; j < span.row_count(); ++j) {
|
||||
auto value = *reinterpret_cast<float*>(
|
||||
(char*)span.data() + j * span.element_sizeof());
|
||||
EXPECT_EQ(value, 1.5f);
|
||||
}
|
||||
if (i > 0) {
|
||||
EXPECT_EQ(
|
||||
cells[i].second->RawData() == cells[0].second->RawData(),
|
||||
true);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Test JSON type
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestJsonType) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 100;
|
||||
|
||||
FieldMeta field_meta(FieldName("test_json"),
|
||||
FieldId(901),
|
||||
DataType::JSON,
|
||||
false,
|
||||
std::nullopt);
|
||||
|
||||
FieldDataInfo field_data_info(901, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_GT(translator->num_cells(), 0);
|
||||
EXPECT_EQ(translator->value_size(), sizeof(Json));
|
||||
}
|
||||
|
||||
// Test ARRAY type
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestArrayType) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 100;
|
||||
|
||||
FieldMeta field_meta(FieldName("test_array"),
|
||||
FieldId(902),
|
||||
DataType::ARRAY,
|
||||
DataType::INT32,
|
||||
false,
|
||||
std::nullopt);
|
||||
|
||||
FieldDataInfo field_data_info(902, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_GT(translator->num_cells(), 0);
|
||||
EXPECT_EQ(translator->value_size(), sizeof(Array));
|
||||
}
|
||||
|
||||
// Test TIMESTAMPTZ type
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestTimestamptzType) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 200;
|
||||
|
||||
DefaultValueType value_field;
|
||||
value_field.set_timestamptz_data(1234567890); // Timestamp value
|
||||
FieldMeta field_meta(FieldName("test_timestamptz"),
|
||||
FieldId(903),
|
||||
DataType::TIMESTAMPTZ,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(903, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
EXPECT_GT(translator->num_cells(), 0);
|
||||
EXPECT_EQ(translator->value_size(), sizeof(int64_t));
|
||||
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
EXPECT_EQ(cells.size(), 1);
|
||||
|
||||
auto& [cid, chunk] = cells[0];
|
||||
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
EXPECT_GT(span.row_count(), 0);
|
||||
|
||||
for (size_t i = 0; i < span.row_count(); ++i) {
|
||||
auto value = *reinterpret_cast<int64_t*>((char*)span.data() +
|
||||
i * span.element_sizeof());
|
||||
EXPECT_EQ(value, 1234567890);
|
||||
}
|
||||
}
|
||||
|
||||
// Test with zero rows (edge case)
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestZeroRows) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 0;
|
||||
|
||||
DefaultValueType value_field;
|
||||
value_field.set_long_data(0);
|
||||
FieldMeta field_meta(FieldName("test_zero"),
|
||||
FieldId(1001),
|
||||
DataType::INT64,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(1001, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
// Zero rows should result in zero cells
|
||||
EXPECT_EQ(translator->num_cells(), 0);
|
||||
}
|
||||
|
||||
// Test estimated_byte_size_of_cell for different cell indices
|
||||
TEST_P(DefaultValueChunkTranslatorTest, TestEstimatedByteSizeMultipleCells) {
|
||||
bool use_mmap = GetParam();
|
||||
int64_t row_count = 5000000; // Ensure multiple cells
|
||||
|
||||
DefaultValueType value_field;
|
||||
value_field.set_long_data(42);
|
||||
FieldMeta field_meta(FieldName("test_byte_size"),
|
||||
FieldId(1101),
|
||||
DataType::INT64,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(1101, row_count, getMmapDirPath());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, use_mmap, true);
|
||||
|
||||
size_t num_cells = translator->num_cells();
|
||||
if (num_cells > 1) {
|
||||
for (size_t i = 0; i < num_cells; ++i) {
|
||||
auto [usage, peak_usage] =
|
||||
translator->estimated_byte_size_of_cell(i);
|
||||
if (use_mmap) {
|
||||
EXPECT_GT(usage.file_bytes, 0);
|
||||
} else {
|
||||
EXPECT_GT(usage.memory_bytes, 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Parameterized test with both mmap modes
|
||||
INSTANTIATE_TEST_SUITE_P(MmapModes,
|
||||
DefaultValueChunkTranslatorTest,
|
||||
testing::Bool());
|
||||
|
||||
// Non-parameterized test class for mmap file verification tests
|
||||
class DefaultValueChunkTranslatorMmapTest : public ::testing::Test {
|
||||
protected:
|
||||
void
|
||||
SetUp() override {
|
||||
// Create a unique temp directory for each test
|
||||
temp_dir_ =
|
||||
std::filesystem::temp_directory_path() /
|
||||
("milvus_test_" + std::to_string(::testing::UnitTest::GetInstance()
|
||||
->current_test_info()
|
||||
->line()));
|
||||
std::filesystem::create_directories(temp_dir_);
|
||||
}
|
||||
|
||||
void
|
||||
TearDown() override {
|
||||
// Clean up temp directory
|
||||
if (std::filesystem::exists(temp_dir_)) {
|
||||
std::filesystem::remove_all(temp_dir_);
|
||||
}
|
||||
}
|
||||
|
||||
std::filesystem::path temp_dir_;
|
||||
int64_t segment_id_ = 99999;
|
||||
};
|
||||
|
||||
// Test that mmap creates file on disk
|
||||
TEST_F(DefaultValueChunkTranslatorMmapTest, TestMmapCreatesFile) {
|
||||
int64_t row_count = 1000;
|
||||
int64_t field_id = 101;
|
||||
int64_t default_value = 42;
|
||||
|
||||
proto::schema::ValueField value_field;
|
||||
value_field.set_long_data(default_value);
|
||||
FieldMeta field_meta(FieldName("test_mmap_file"),
|
||||
FieldId(field_id),
|
||||
DataType::INT64,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
// Pass mmap_dir_path to FieldDataInfo
|
||||
FieldDataInfo field_data_info(field_id, row_count, temp_dir_.string());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, true /* use_mmap */, true);
|
||||
|
||||
// Trigger cell creation
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
ASSERT_EQ(cells.size(), 1);
|
||||
|
||||
// Verify file was created
|
||||
auto expected_file =
|
||||
temp_dir_ / fmt::format("seg_{}_f_{}_def", segment_id_, field_id);
|
||||
EXPECT_TRUE(std::filesystem::exists(expected_file))
|
||||
<< "Expected mmap file to be created at: " << expected_file;
|
||||
|
||||
// Verify the chunk data is correct
|
||||
auto& [cid, chunk] = cells[0];
|
||||
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
EXPECT_GT(span.row_count(), 0);
|
||||
|
||||
for (size_t i = 0; i < span.row_count(); ++i) {
|
||||
auto value =
|
||||
*(int64_t*)((char*)span.data() + i * span.element_sizeof());
|
||||
EXPECT_EQ(value, default_value);
|
||||
}
|
||||
}
|
||||
|
||||
// Test that non-mmap mode does not create file
|
||||
TEST_F(DefaultValueChunkTranslatorMmapTest, TestNoMmapNoFile) {
|
||||
int64_t row_count = 1000;
|
||||
int64_t field_id = 102;
|
||||
int64_t default_value = 99;
|
||||
|
||||
proto::schema::ValueField value_field;
|
||||
value_field.set_long_data(default_value);
|
||||
FieldMeta field_meta(FieldName("test_no_mmap_file"),
|
||||
FieldId(field_id),
|
||||
DataType::INT64,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
// Pass mmap_dir_path even though mmap is disabled
|
||||
FieldDataInfo field_data_info(field_id, row_count, temp_dir_.string());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, false /* use_mmap */, true);
|
||||
|
||||
// Trigger cell creation
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
ASSERT_EQ(cells.size(), 1);
|
||||
|
||||
// Verify no file was created (memory-only mode)
|
||||
auto unexpected_file =
|
||||
temp_dir_ / fmt::format("seg_{}_f_{}_def", segment_id_, field_id);
|
||||
EXPECT_FALSE(std::filesystem::exists(unexpected_file))
|
||||
<< "Expected no mmap file when use_mmap=false, but found: "
|
||||
<< unexpected_file;
|
||||
|
||||
// Verify the chunk data is still correct
|
||||
auto& [cid, chunk] = cells[0];
|
||||
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
EXPECT_GT(span.row_count(), 0);
|
||||
|
||||
for (size_t i = 0; i < span.row_count(); ++i) {
|
||||
auto value =
|
||||
*(int64_t*)((char*)span.data() + i * span.element_sizeof());
|
||||
EXPECT_EQ(value, default_value);
|
||||
}
|
||||
}
|
||||
|
||||
// Test mmap with string type
|
||||
TEST_F(DefaultValueChunkTranslatorMmapTest, TestMmapWithString) {
|
||||
int64_t row_count = 500;
|
||||
int64_t field_id = 103;
|
||||
std::string default_string = "mmap_default_value";
|
||||
|
||||
proto::schema::ValueField value_field;
|
||||
value_field.set_string_data(default_string);
|
||||
FieldMeta field_meta(FieldName("test_mmap_string"),
|
||||
FieldId(field_id),
|
||||
DataType::VARCHAR,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(field_id, row_count, temp_dir_.string());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, true /* use_mmap */, true);
|
||||
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
ASSERT_EQ(cells.size(), 1);
|
||||
|
||||
// Verify file was created
|
||||
auto expected_file =
|
||||
temp_dir_ / fmt::format("seg_{}_f_{}_def", segment_id_, field_id);
|
||||
EXPECT_TRUE(std::filesystem::exists(expected_file))
|
||||
<< "Expected mmap file for string type at: " << expected_file;
|
||||
|
||||
// Verify string data
|
||||
auto& [cid, chunk] = cells[0];
|
||||
auto string_chunk = static_cast<StringChunk*>(chunk.get());
|
||||
auto [views, valid] = string_chunk->StringViews(std::nullopt);
|
||||
|
||||
EXPECT_GT(views.size(), 0);
|
||||
for (size_t i = 0; i < views.size(); ++i) {
|
||||
EXPECT_EQ(views[i], default_string);
|
||||
}
|
||||
}
|
||||
|
||||
// Test mmap with nullable field (nulls)
|
||||
TEST_F(DefaultValueChunkTranslatorMmapTest, TestMmapWithNullableField) {
|
||||
int64_t row_count = 300;
|
||||
int64_t field_id = 104;
|
||||
|
||||
FieldMeta field_meta(FieldName("test_mmap_nullable"),
|
||||
FieldId(field_id),
|
||||
DataType::INT64,
|
||||
true, // nullable
|
||||
std::nullopt);
|
||||
|
||||
FieldDataInfo field_data_info(field_id, row_count, temp_dir_.string());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, true /* use_mmap */, true);
|
||||
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
ASSERT_EQ(cells.size(), 1);
|
||||
|
||||
// Verify file was created
|
||||
auto expected_file =
|
||||
temp_dir_ / fmt::format("seg_{}_f_{}_def", segment_id_, field_id);
|
||||
EXPECT_TRUE(std::filesystem::exists(expected_file))
|
||||
<< "Expected mmap file for nullable field at: " << expected_file;
|
||||
|
||||
// Verify all values are null
|
||||
auto& [cid, chunk] = cells[0];
|
||||
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
EXPECT_GT(span.row_count(), 0);
|
||||
|
||||
for (size_t i = 0; i < span.row_count(); ++i) {
|
||||
EXPECT_FALSE(fixed_chunk->isValid(i));
|
||||
}
|
||||
}
|
||||
|
||||
// Test mmap with multiple cells
|
||||
TEST_F(DefaultValueChunkTranslatorMmapTest, TestMmapMultipleCells) {
|
||||
int64_t row_count = 10000000; // Large enough for multiple cells
|
||||
int64_t field_id = 105;
|
||||
int64_t default_value = 12345;
|
||||
|
||||
proto::schema::ValueField value_field;
|
||||
value_field.set_long_data(default_value);
|
||||
FieldMeta field_meta(FieldName("test_mmap_multi_cells"),
|
||||
FieldId(field_id),
|
||||
DataType::INT64,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(field_id, row_count, temp_dir_.string());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, true /* use_mmap */, true);
|
||||
|
||||
// Ensure we have multiple cells
|
||||
size_t num_cells = translator->num_cells();
|
||||
ASSERT_GT(num_cells, 1) << "Expected multiple cells for large row count";
|
||||
|
||||
// Request multiple cells
|
||||
std::vector<cachinglayer::cid_t> cids;
|
||||
for (size_t i = 0; i < std::min<size_t>(3, num_cells); ++i) {
|
||||
cids.push_back(i);
|
||||
}
|
||||
auto cells = translator->get_cells(cids);
|
||||
ASSERT_EQ(cells.size(), cids.size());
|
||||
|
||||
// Verify file was created (all cells share the same buffer/file)
|
||||
auto expected_file =
|
||||
temp_dir_ / fmt::format("seg_{}_f_{}_def", segment_id_, field_id);
|
||||
EXPECT_TRUE(std::filesystem::exists(expected_file))
|
||||
<< "Expected mmap file at: " << expected_file;
|
||||
|
||||
// Verify data in all cells
|
||||
for (const auto& [cid, chunk] : cells) {
|
||||
auto fixed_chunk = static_cast<FixedWidthChunk*>(chunk.get());
|
||||
auto span = fixed_chunk->Span();
|
||||
EXPECT_GT(span.row_count(), 0);
|
||||
|
||||
for (size_t i = 0; i < span.row_count(); ++i) {
|
||||
auto value =
|
||||
*(int64_t*)((char*)span.data() + i * span.element_sizeof());
|
||||
EXPECT_EQ(value, default_value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Test mmap file is properly sized
|
||||
TEST_F(DefaultValueChunkTranslatorMmapTest, TestMmapFileSize) {
|
||||
int64_t row_count = 1000;
|
||||
int64_t field_id = 106;
|
||||
int64_t default_value = 77;
|
||||
|
||||
proto::schema::ValueField value_field;
|
||||
value_field.set_long_data(default_value);
|
||||
FieldMeta field_meta(FieldName("test_mmap_file_size"),
|
||||
FieldId(field_id),
|
||||
DataType::INT64,
|
||||
false,
|
||||
value_field);
|
||||
|
||||
FieldDataInfo field_data_info(field_id, row_count, temp_dir_.string());
|
||||
|
||||
auto translator = std::make_unique<DefaultValueChunkTranslator>(
|
||||
segment_id_, field_meta, field_data_info, true /* use_mmap */, true);
|
||||
|
||||
std::vector<cachinglayer::cid_t> cids = {0};
|
||||
auto cells = translator->get_cells(cids);
|
||||
ASSERT_EQ(cells.size(), 1);
|
||||
|
||||
auto expected_file =
|
||||
temp_dir_ / fmt::format("seg_{}_f_{}_def", segment_id_, field_id);
|
||||
ASSERT_TRUE(std::filesystem::exists(expected_file));
|
||||
|
||||
// File size should be at least row_count * sizeof(int64_t)
|
||||
auto file_size = std::filesystem::file_size(expected_file);
|
||||
EXPECT_GE(file_size, row_count * sizeof(int64_t))
|
||||
<< "Mmap file size should be at least " << row_count * sizeof(int64_t)
|
||||
<< " bytes";
|
||||
}
|
||||
@@ -437,7 +437,8 @@ GroupChunkTranslator::load_group_chunk(
|
||||
|
||||
std::unordered_map<FieldId, std::shared_ptr<Chunk>> chunks;
|
||||
if (!use_mmap_) {
|
||||
chunks = create_group_chunk(field_ids, field_metas, array_vecs);
|
||||
chunks = create_group_chunk(
|
||||
field_ids, field_metas, array_vecs, mmap_populate_);
|
||||
} else {
|
||||
std::filesystem::path filepath;
|
||||
switch (group_chunk_type_) {
|
||||
@@ -467,8 +468,9 @@ GroupChunkTranslator::load_group_chunk(
|
||||
chunks = create_group_chunk(field_ids,
|
||||
field_metas,
|
||||
array_vecs,
|
||||
mmap_populate_,
|
||||
filepath.string(),
|
||||
mmap_populate_);
|
||||
load_priority_);
|
||||
}
|
||||
return std::make_unique<milvus::GroupChunk>(chunks);
|
||||
}
|
||||
|
||||
@@ -291,7 +291,8 @@ ManifestGroupTranslator::load_group_chunk(
|
||||
std::unordered_map<FieldId, std::shared_ptr<Chunk>> chunks;
|
||||
if (!use_mmap_) {
|
||||
// Memory mode
|
||||
chunks = create_group_chunk(field_ids, field_metas, array_vecs);
|
||||
chunks = create_group_chunk(
|
||||
field_ids, field_metas, array_vecs, mmap_populate_);
|
||||
} else {
|
||||
// Mmap mode
|
||||
std::filesystem::path filepath;
|
||||
@@ -323,8 +324,9 @@ ManifestGroupTranslator::load_group_chunk(
|
||||
chunks = create_group_chunk(field_ids,
|
||||
field_metas,
|
||||
array_vecs,
|
||||
mmap_populate_,
|
||||
filepath.string(),
|
||||
mmap_populate_);
|
||||
load_priority_);
|
||||
}
|
||||
|
||||
return std::make_unique<milvus::GroupChunk>(chunks);
|
||||
|
||||
Reference in New Issue
Block a user