furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 10:15:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: [StorageV2] handle correct cid with multiple files and add storage v2 prefix logs (#43539)

Browse Source 
related: #43372

---------

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
This commit is contained in:
sthuang
2025-07-25 11:22:54 +08:00
committed by GitHub
parent 87326a5a64
commit 5cebc9f7f6
13 changed files with 383 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/core/src/common/ArrowDataWrapper.h
+8 -1
View File
@@ -19,6 +19,13 @@
namespace milvus {
struct ArrowTableInfo {
size_t file_index; // file index in the file list
size_t row_group_index; // row group index in the file, start from 0
std::shared_ptr<arrow::Table>
table; // the arrow table of the row group at row group index in the file index file
};
struct ArrowDataWrapper {
ArrowDataWrapper() = default;
ArrowDataWrapper(std::shared_ptr<arrow::RecordBatchReader> reader,
@@ -33,7 +40,7 @@ struct ArrowDataWrapper {
std::shared_ptr<parquet::arrow::FileReader> arrow_reader;
// underlying file data memory, must outlive the arrow reader
std::shared_ptr<uint8_t[]> file_data;
std::vector<std::pair<int64_t, std::shared_ptr<arrow::Table>>> arrow_tables;
std::vector<ArrowTableInfo> arrow_tables;
};
using ArrowReaderChannel = Channel<std::shared_ptr<milvus::ArrowDataWrapper>>;
internal/core/src/mmap/ChunkedColumnGroup.h
+18 -16
View File
@@ -47,8 +47,7 @@ class ChunkedColumnGroup {
public:
explicit ChunkedColumnGroup(
std::unique_ptr<Translator<GroupChunk>> translator)
: slot_(Manager::GetInstance().CreateCacheSlot(
std::move(translator))) {
: slot_(Manager::GetInstance().CreateCacheSlot(std::move(translator))) {
num_chunks_ = slot_->num_cells();
num_rows_ = GetNumRowsUntilChunk().back();
}
@@ -86,7 +85,7 @@ class ChunkedColumnGroup {
int64_t
GetNumRowsUntilChunk(int64_t chunk_id) const {
AssertInfo(chunk_id >= 0 && chunk_id <= num_chunks_,
"chunk_id out of range: " + std::to_string(chunk_id));
"[StorageV2] chunk_id out of range: " + std::to_string(chunk_id));
return GetNumRowsUntilChunk()[chunk_id];
}
@@ -104,7 +103,8 @@ class ChunkedColumnGroup {
auto iter = std::lower_bound(num_rows_until_chunk.begin(),
num_rows_until_chunk.end(),
offset + 1);
size_t chunk_idx = std::distance(num_rows_until_chunk.begin(), iter) - 1;
size_t chunk_idx =
std::distance(num_rows_until_chunk.begin(), iter) - 1;
size_t offset_in_chunk = offset - num_rows_until_chunk[chunk_idx];
return {chunk_idx, offset_in_chunk};
}
@@ -250,7 +250,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
Span(int64_t chunk_id) const override {
if (!IsChunkedColumnDataType(data_type_)) {
ThrowInfo(ErrorCode::Unsupported,
"Span only supported for ChunkedColumn");
"[StorageV2] Span only supported for ChunkedColumn");
}
auto chunk_wrapper = group_->GetGroupChunk(chunk_id);
auto chunk = chunk_wrapper.get()->GetChunk(field_id_);
@@ -264,7 +264,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
std::nullopt) const override {
if (!IsChunkedVariableColumnDataType(data_type_)) {
ThrowInfo(ErrorCode::Unsupported,
"StringViews only supported for ChunkedVariableColumn");
"[StorageV2] StringViews only supported for ChunkedVariableColumn");
}
auto chunk_wrapper = group_->GetGroupChunk(chunk_id);
auto chunk = chunk_wrapper.get()->GetChunk(field_id_);
@@ -280,7 +280,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
std::nullopt) const override {
if (!IsChunkedArrayColumnDataType(data_type_)) {
ThrowInfo(ErrorCode::Unsupported,
"ArrayViews only supported for ChunkedArrayColumn");
"[StorageV2] ArrayViews only supported for ChunkedArrayColumn");
}
auto chunk_wrapper = group_->GetGroupChunk(chunk_id);
auto chunk = chunk_wrapper.get()->GetChunk(field_id_);
@@ -294,7 +294,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
if (!IsChunkedVectorArrayColumnDataType(data_type_)) {
ThrowInfo(
ErrorCode::Unsupported,
"VectorArrayViews only supported for ChunkedVectorArrayColumn");
"[StorageV2] VectorArrayViews only supported for ChunkedVectorArrayColumn");
}
auto chunk_wrapper = group_->GetGroupChunk(chunk_id);
auto chunk = chunk_wrapper.get()->GetChunk(field_id_);
@@ -309,7 +309,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
if (!IsChunkedVariableColumnDataType(data_type_)) {
ThrowInfo(
ErrorCode::Unsupported,
"ViewsByOffsets only supported for ChunkedVariableColumn");
"[StorageV2] ViewsByOffsets only supported for ChunkedVariableColumn");
}
auto chunk_wrapper = group_->GetGroupChunk(chunk_id);
auto chunk = chunk_wrapper.get()->GetChunk(field_id_);
@@ -411,7 +411,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
default: {
ThrowInfo(
ErrorCode::Unsupported,
"BulkScalarValueAt is not supported for unknown scalar "
"[StorageV2] BulkScalarValueAt is not supported for unknown scalar "
"data type: {}",
data_type_);
}
@@ -441,7 +441,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
if (!IsChunkedVariableColumnDataType(data_type_) ||
data_type_ == DataType::JSON) {
ThrowInfo(ErrorCode::Unsupported,
"BulkRawStringAt only supported for ProxyChunkColumn of "
"[StorageV2] BulkRawStringAt only supported for ProxyChunkColumn of "
"variable length type(except Json)");
}
if (offsets == nullptr) {
@@ -466,7 +466,8 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
auto chunk = group_chunk->GetChunk(field_id_);
auto valid = chunk->isValid(offsets_in_chunk[i]);
auto value = static_cast<StringChunk*>(chunk.get())
->operator[](offsets_in_chunk[i]);
->
operator[](offsets_in_chunk[i]);
fn(value, i, valid);
}
}
@@ -480,7 +481,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
if (data_type_ != DataType::JSON) {
ThrowInfo(
ErrorCode::Unsupported,
"RawJsonAt only supported for ProxyChunkColumn of Json type");
"[StorageV2] RawJsonAt only supported for ProxyChunkColumn of Json type");
}
if (count == 0) {
return;
@@ -493,7 +494,8 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
auto chunk = group_chunk->GetChunk(field_id_);
auto valid = chunk->isValid(offsets_in_chunk[i]);
auto str_view = static_cast<StringChunk*>(chunk.get())
->operator[](offsets_in_chunk[i]);
->
operator[](offsets_in_chunk[i]);
fn(Json(str_view.data(), str_view.size()), i, valid);
}
}
@@ -504,7 +506,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
int64_t count) const override {
if (!IsChunkedArrayColumnDataType(data_type_)) {
ThrowInfo(ErrorCode::Unsupported,
"BulkArrayAt only supported for ChunkedArrayColumn");
"[StorageV2] BulkArrayAt only supported for ChunkedArrayColumn");
}
auto [cids, offsets_in_chunk] = ToChunkIdAndOffset(offsets, count);
auto ca = group_->GetGroupChunks(cids);
@@ -524,7 +526,7 @@ class ProxyChunkColumn : public ChunkedColumnInterface {
int64_t count) const override {
if (!IsChunkedVectorArrayColumnDataType(data_type_)) {
ThrowInfo(ErrorCode::Unsupported,
"BulkVectorArrayAt only supported for "
"[StorageV2] BulkVectorArrayAt only supported for "
"ChunkedVectorArrayColumn");
}
auto [cids, offsets_in_chunk] = ToChunkIdAndOffset(offsets, count);
internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp
+10 -9
View File
@@ -252,11 +252,11 @@ ChunkedSegmentSealedImpl::LoadFieldData(const LoadFieldDataInfo& load_info) {
}
init_timestamp_index(timestamps, num_rows);
system_ready_count_++;
AssertInfo(
offset == num_rows,
"timestamp total row count {} not equal to expected {}",
offset,
num_rows);
AssertInfo(offset == num_rows,
"[StorageV2] timestamp total row count {} not equal "
"to expected {}",
offset,
num_rows);
}
break;
default:
@@ -272,7 +272,8 @@ ChunkedSegmentSealedImpl::load_column_group_data_internal(
ArrowSchemaPtr arrow_schema = schema_->ConvertToArrowSchema();
for (auto& [id, info] : load_info.field_infos) {
AssertInfo(info.row_count > 0, "The row count of field data is 0");
AssertInfo(info.row_count > 0,
"[StorageV2] The row count of field data is 0");
auto column_group_id = FieldId(id);
auto insert_files = info.insert_files;
@@ -299,7 +300,8 @@ ChunkedSegmentSealedImpl::load_column_group_data_internal(
load_info.mmap_dir_path,
merged_in_load_list);
LOG_INFO(
"segment {} loads column group {} with field ids {} with "
"[StorageV2] segment {} loads column group {} with field ids {} "
"with "
"num_rows "
"{}",
this->get_segment_id(),
@@ -1082,8 +1084,7 @@ ChunkedSegmentSealedImpl::bulk_subscript_impl(ChunkedColumnInterface* field,
T* dst) {
static_assert(std::is_fundamental_v<S> && std::is_fundamental_v<T>);
// use field->data_type_ to determine the type of dst
field->BulkPrimitiveValueAt(
static_cast<void*>(dst), seg_offsets, count);
field->BulkPrimitiveValueAt(static_cast<void*>(dst), seg_offsets, count);
}
// for dense vector
internal/core/src/segcore/SegmentGrowingImpl.cpp
+30 -20
View File
@@ -449,10 +449,11 @@ SegmentGrowingImpl::load_column_group_data_internal(
FieldDataInfo(column_group_id.get(), num_rows, infos.mmap_dir_path);
column_group_info.arrow_reader_channel->set_capacity(parallel_degree);
LOG_INFO("segment {} loads column group {} with num_rows {}",
this->get_segment_id(),
column_group_id.get(),
num_rows);
LOG_INFO(
"[StorageV2] segment {} loads column group {} with num_rows {}",
this->get_segment_id(),
column_group_id.get(),
num_rows);
auto& pool =
ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::MIDDLE);
@@ -469,10 +470,12 @@ SegmentGrowingImpl::load_column_group_data_internal(
std::iota(all_row_groups.begin(), all_row_groups.end(), 0);
row_group_lists.push_back(all_row_groups);
auto status = reader->Close();
AssertInfo(status.ok(),
"failed to close file reader when get row group "
"metadata from file: " +
file + " with error: " + status.ToString());
AssertInfo(
status.ok(),
"[StorageV2] failed to close file reader when get row group "
"metadata from file {} with error {}",
file,
status.ToString());
}
// create parallel degree split strategy
@@ -489,23 +492,30 @@ SegmentGrowingImpl::load_column_group_data_internal(
infos.load_priority);
});
LOG_INFO("segment {} submits load column group {} task to thread pool",
this->get_segment_id(),
info.field_id);
LOG_INFO(
"[StorageV2] segment {} submits load column group {} task to "
"thread pool",
this->get_segment_id(),
column_group_id.get());
std::shared_ptr<milvus::ArrowDataWrapper> r;
std::unordered_map<FieldId, std::vector<FieldDataPtr>> field_data_map;
while (column_group_info.arrow_reader_channel->pop(r)) {
for (const auto& [row_group_id, table] : r->arrow_tables) {
size_t batch_num_rows = table->num_rows();
for (int i = 0; i < table->schema()->num_fields(); ++i) {
AssertInfo(table->schema()->field(i)->metadata()->Contains(
milvus_storage::ARROW_FIELD_ID_KEY),
"field id not found in metadata for field {}",
table->schema()->field(i)->name());
for (const auto& table_info : r->arrow_tables) {
size_t batch_num_rows = table_info.table->num_rows();
for (int i = 0; i < table_info.table->schema()->num_fields();
++i) {
AssertInfo(
table_info.table->schema()
->field(i)
->metadata()
->Contains(milvus_storage::ARROW_FIELD_ID_KEY),
"[StorageV2] field id not found in metadata for field "
"{}",
table_info.table->schema()->field(i)->name());
auto field_id =
std::stoll(table->schema()
std::stoll(table_info.table->schema()
->field(i)
->metadata()
->Get(milvus_storage::ARROW_FIELD_ID_KEY)
@@ -523,7 +533,7 @@ SegmentGrowingImpl::load_column_group_data_internal(
? field.second.get_dim()
: 1,
batch_num_rows);
field_data->FillFieldData(table->column(i));
field_data->FillFieldData(table_info.table->column(i));
field_data_map[FieldId(field_id)].push_back(field_data);
}
}
internal/core/src/segcore/memory_planner.cpp
+25 -13
View File
@@ -151,9 +151,9 @@ LoadWithStrategy(const std::vector<std::string>& remote_files,
const std::shared_ptr<arrow::Schema> schema,
milvus::proto::common::LoadPriority priority) {
try {
AssertInfo(
remote_files.size() == row_group_lists.size(),
"Number of remote files must match number of row group lists");
AssertInfo(remote_files.size() == row_group_lists.size(),
"[StorageV2] Number of remote files must match number of "
"row group lists");
auto fs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
auto& pool =
@@ -170,8 +170,9 @@ LoadWithStrategy(const std::vector<std::string>& remote_files,
// Use provided strategy to split row groups
auto blocks = strategy->split(row_groups);
LOG_INFO("split row groups into blocks: {} for file {}",
blocks.size(), file);
LOG_INFO("[StorageV2] split row groups into blocks: {} for file {}",
blocks.size(),
file);
// Create and submit tasks for each block
std::vector<std::future<std::shared_ptr<milvus::ArrowDataWrapper>>>
@@ -185,33 +186,44 @@ LoadWithStrategy(const std::vector<std::string>& remote_files,
futures.emplace_back(pool.Submit([block,
fs,
file,
file_idx,
schema,
reader_memory_limit]() {
AssertInfo(fs != nullptr, "file system is nullptr");
AssertInfo(fs != nullptr,
"[StorageV2] file system is nullptr");
auto row_group_reader =
std::make_shared<milvus_storage::FileRowGroupReader>(
fs, file, schema, reader_memory_limit);
AssertInfo(row_group_reader != nullptr,
"row group reader is nullptr");
"[StorageV2] row group reader is nullptr");
row_group_reader->SetRowGroupOffsetAndCount(block.offset,
block.count);
LOG_INFO("read row groups from file {} with offset {} and count {}",
file, block.offset, block.count);
LOG_INFO(
"[StorageV2] read row groups from file {} with offset "
"{} and count "
"{}",
file,
block.offset,
block.count);
auto ret = std::make_shared<ArrowDataWrapper>();
for (int64_t i = 0; i < block.count; ++i) {
std::shared_ptr<arrow::Table> table;
auto status =
row_group_reader->ReadNextRowGroup(&table);
AssertInfo(status.ok(),
"Failed to read row group " +
"[StorageV2] Failed to read row group " +
std::to_string(block.offset + i) +
" from file " + file + " with error " +
status.ToString());
ret->arrow_tables.push_back(std::make_pair(block.offset + i, table));
ret->arrow_tables.push_back(
{file_idx,
static_cast<size_t>(block.offset + i),
table});
}
auto close_status = row_group_reader->Close();
AssertInfo(close_status.ok(),
"Failed to close row group reader for file " +
"[StorageV2] Failed to close row group reader "
"for file " +
file + " with error " +
close_status.ToString());
return ret;
@@ -226,7 +238,7 @@ LoadWithStrategy(const std::vector<std::string>& remote_files,
channel->close();
} catch (std::exception& e) {
LOG_INFO("failed to load data from remote: {}", e.what());
LOG_INFO("[StorageV2] failed to load data from remote: {}", e.what());
channel->close();
}
}
internal/core/src/segcore/packed_reader_c.cpp
+6 -4
View File
@@ -61,8 +61,9 @@ NewPackedReaderWithStorageConfig(char** paths,
auto trueFs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
if (!trueFs) {
return milvus::FailureCStatus(milvus::ErrorCode::FileReadFailed,
"Failed to get filesystem");
return milvus::FailureCStatus(
milvus::ErrorCode::FileReadFailed,
"[StorageV2] Failed to get filesystem");
}
auto trueSchema = arrow::ImportSchema(schema).ValueOrDie();
auto reader = std::make_unique<milvus_storage::PackedRecordBatchReader>(
@@ -87,8 +88,9 @@ NewPackedReader(char** paths,
auto trueFs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
if (!trueFs) {
return milvus::FailureCStatus(milvus::ErrorCode::FileReadFailed,
"Failed to get filesystem");
return milvus::FailureCStatus(
milvus::ErrorCode::FileReadFailed,
"[StorageV2] Failed to get filesystem");
}
auto trueSchema = arrow::ImportSchema(schema).ValueOrDie();
auto reader = std::make_unique<milvus_storage::PackedRecordBatchReader>(
internal/core/src/segcore/packed_writer_c.cpp
+8 -6
View File
@@ -65,8 +65,9 @@ NewPackedWriterWithStorageConfig(struct ArrowSchema* schema,
auto trueFs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
if (!trueFs) {
return milvus::FailureCStatus(milvus::ErrorCode::FileReadFailed,
"Failed to get filesystem");
return milvus::FailureCStatus(
milvus::ErrorCode::FileReadFailed,
"[StorageV2] Failed to get filesystem");
}
auto trueSchema = arrow::ImportSchema(schema).ValueOrDie();
@@ -81,7 +82,7 @@ NewPackedWriterWithStorageConfig(struct ArrowSchema* schema,
storage_config,
columnGroups,
buffer_size);
AssertInfo(writer, "write must not be null");
AssertInfo(writer, "[StorageV2] writer pointer is null");
*c_packed_writer = writer.release();
return milvus::SuccessCStatus();
@@ -109,8 +110,9 @@ NewPackedWriter(struct ArrowSchema* schema,
auto trueFs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
if (!trueFs) {
return milvus::FailureCStatus(milvus::ErrorCode::FileWriteFailed,
"Failed to get filesystem");
return milvus::FailureCStatus(
milvus::ErrorCode::FileWriteFailed,
"[StorageV2] Failed to get filesystem");
}
auto trueSchema = arrow::ImportSchema(schema).ValueOrDie();
@@ -120,7 +122,7 @@ NewPackedWriter(struct ArrowSchema* schema,
auto writer = std::make_unique<milvus_storage::PackedRecordBatchWriter>(
trueFs, truePaths, trueSchema, conf, columnGroups, buffer_size);
AssertInfo(writer, "write must not be null");
AssertInfo(writer, "[StorageV2] writer pointer is null");
*c_packed_writer = writer.release();
return milvus::SuccessCStatus();
internal/core/src/segcore/storagev2translator/GroupChunkTranslator.cpp
+88 -49
View File
@@ -67,7 +67,7 @@ GroupChunkTranslator::GroupChunkTranslator(
/* support_eviction */ true) {
auto fs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
// Get row group metadata from files
for (const auto& file : insert_files_) {
auto reader =
@@ -76,9 +76,20 @@ GroupChunkTranslator::GroupChunkTranslator(
reader->file_metadata()->GetRowGroupMetadataVector());
auto status = reader->Close();
AssertInfo(status.ok(),
"failed to close file reader when get row group "
"metadata from file: " +
file + " with error: " + status.ToString());
"[StorageV2] translator {} failed to close file reader when "
"get row group "
"metadata from file {} with error {}",
key_,
file + " with error: " + status.ToString());
}
// Build prefix sum for O(1) lookup in get_cid_from_file_and_row_group_index
file_row_group_prefix_sum_.reserve(row_group_meta_list_.size() + 1);
file_row_group_prefix_sum_.push_back(
0); // Base case: 0 row groups before first file
for (const auto& file_metas : row_group_meta_list_) {
file_row_group_prefix_sum_.push_back(file_row_group_prefix_sum_.back() +
file_metas.size());
}
meta_.num_rows_until_chunk_.push_back(0);
@@ -93,11 +104,12 @@ GroupChunkTranslator::GroupChunkTranslator(
}
AssertInfo(
meta_.num_rows_until_chunk_.back() == column_group_info_.row_count,
fmt::format("data lost while loading column group {}: found "
"num rows {} but expected {}",
column_group_info_.field_id,
meta_.num_rows_until_chunk_.back(),
column_group_info_.row_count));
fmt::format(
"[StorageV2] data lost while loading column group {}: found "
"num rows {} but expected {}",
column_group_info_.field_id,
meta_.num_rows_until_chunk_.back(),
column_group_info_.row_count));
}
GroupChunkTranslator::~GroupChunkTranslator() {
@@ -130,30 +142,46 @@ GroupChunkTranslator::key() const {
std::pair<size_t, size_t>
GroupChunkTranslator::get_file_and_row_group_index(
milvus::cachinglayer::cid_t cid) const {
size_t file_idx = 0;
size_t remaining_cid = cid;
for (; file_idx < row_group_meta_list_.size(); ++file_idx) {
const auto& file_metas = row_group_meta_list_[file_idx];
if (remaining_cid < file_metas.size()) {
return {file_idx, remaining_cid};
for (size_t file_idx = 0; file_idx < file_row_group_prefix_sum_.size() - 1;
++file_idx) {
if (cid < file_row_group_prefix_sum_[file_idx + 1]) {
return {file_idx, cid - file_row_group_prefix_sum_[file_idx]};
}
remaining_cid -= file_metas.size();
}
// cid is out of range, this should not happen
AssertInfo(false,
fmt::format("cid {} is out of range. Total row groups across all files: {}",
cid,
[this]() {
size_t total = 0;
for (const auto& file_metas : row_group_meta_list_) {
total += file_metas.size();
}
return total;
}()));
AssertInfo(false,
fmt::format("[StorageV2] translator {} cid {} is out of range. "
"Total row groups across all files: {}",
key_,
cid,
file_row_group_prefix_sum_.back()));
}
milvus::cachinglayer::cid_t
GroupChunkTranslator::get_cid_from_file_and_row_group_index(
size_t file_idx, size_t row_group_idx) const {
AssertInfo(file_idx < file_row_group_prefix_sum_.size() - 1,
fmt::format("[StorageV2] translator {} file_idx {} is out of "
"range. Total files: {}",
key_,
file_idx,
file_row_group_prefix_sum_.size() - 1));
size_t file_start = file_row_group_prefix_sum_[file_idx];
size_t file_end = file_row_group_prefix_sum_[file_idx + 1];
AssertInfo(row_group_idx < file_end - file_start,
fmt::format("[StorageV2] translator {} row_group_idx {} is out "
"of range for file {}. "
"Total row groups in file: {}",
key_,
row_group_idx,
file_idx,
file_end - file_start));
return file_start + row_group_idx;
}
// the returned cids are sorted. It may not follow the order of cids.
std::vector<std::pair<cachinglayer::cid_t, std::unique_ptr<milvus::GroupChunk>>>
GroupChunkTranslator::get_cells(const std::vector<cachinglayer::cid_t>& cids) {
std::vector<std::pair<milvus::cachinglayer::cid_t,
@@ -185,26 +213,32 @@ GroupChunkTranslator::get_cells(const std::vector<cachinglayer::cid_t>& cids) {
nullptr,
load_priority_);
});
LOG_INFO("segment {} submits load column group {} task to thread pool",
segment_id_,
column_group_info_.field_id);
LOG_INFO(
"[StorageV2] translator {} submits load column group {} task to thread "
"pool",
key_,
column_group_info_.field_id);
std::shared_ptr<milvus::ArrowDataWrapper> r;
std::unordered_set<cachinglayer::cid_t> filled_cids;
filled_cids.reserve(cids.size());
while (column_group_info_.arrow_reader_channel->pop(r)) {
for (const auto& [row_group_id, table] : r->arrow_tables) {
auto cid = static_cast<cachinglayer::cid_t>(row_group_id);
cells.emplace_back(cid, load_group_chunk(table, cid));
for (const auto& table_info : r->arrow_tables) {
// Convert file_index and row_group_index to global cid
auto cid = get_cid_from_file_and_row_group_index(
table_info.file_index, table_info.row_group_index);
cells.emplace_back(cid, load_group_chunk(table_info.table, cid));
filled_cids.insert(cid);
}
}
// Verify all requested cids have been filled
for (auto cid : cids) {
AssertInfo(filled_cids.find(cid) != filled_cids.end(),
"Cid {} was not filled, missing row group id {}",
cid, cid);
"[StorageV2] translator {} cid {} was not filled, missing "
"row group id {}",
key_,
cid,
cid);
}
return cells;
}
@@ -219,7 +253,9 @@ GroupChunkTranslator::load_group_chunk(
for (int i = 0; i < table->schema()->num_fields(); ++i) {
AssertInfo(table->schema()->field(i)->metadata()->Contains(
milvus_storage::ARROW_FIELD_ID_KEY),
"field id not found in metadata for field {}",
"[StorageV2] translator {} field id not found in metadata "
"for field {}",
key_,
table->schema()->field(i)->name());
auto field_id = std::stoll(table->schema()
->field(i)
@@ -233,8 +269,11 @@ GroupChunkTranslator::load_group_chunk(
continue;
}
auto it = field_metas_.find(fid);
AssertInfo(it != field_metas_.end(),
"Field id not found in field_metas");
AssertInfo(
it != field_metas_.end(),
"[StorageV2] translator {} field id {} not found in field_metas",
key_,
fid.get());
const auto& field_meta = it->second;
const arrow::ArrayVector& array_vec = table->column(i)->chunks();
std::unique_ptr<Chunk> chunk;
@@ -249,8 +288,9 @@ GroupChunkTranslator::load_group_chunk(
std::to_string(cid);
LOG_INFO(
"storage v2 segment {} mmaping field {} chunk {} to path {}",
segment_id_,
"[StorageV2] translator {} mmaping field {} chunk {} to path "
"{}",
key_,
field_id,
cid,
filepath.string());
@@ -259,12 +299,12 @@ GroupChunkTranslator::load_group_chunk(
chunk = create_chunk(field_meta, array_vec, filepath.string());
auto ok = unlink(filepath.c_str());
AssertInfo(
ok == 0,
fmt::format(
"storage v2 failed to unlink mmap data file {}, err: {}",
filepath.c_str(),
strerror(errno)));
AssertInfo(ok == 0,
fmt::format("[StorageV2] translator {} failed to unlink "
"mmap data file {}, err: {}",
key_,
filepath.c_str(),
strerror(errno)));
}
chunks[fid] = std::move(chunk);
@@ -273,4 +313,3 @@ GroupChunkTranslator::load_group_chunk(
}
} // namespace milvus::segcore::storagev2translator
internal/core/src/segcore/storagev2translator/GroupChunkTranslator.h
+5
View File
@@ -65,6 +65,10 @@ class GroupChunkTranslator
std::pair<size_t, size_t>
get_file_and_row_group_index(milvus::cachinglayer::cid_t cid) const;
milvus::cachinglayer::cid_t
get_cid_from_file_and_row_group_index(size_t file_idx,
size_t row_group_idx) const;
milvus::cachinglayer::Meta*
meta() override {
return &meta_;
@@ -81,6 +85,7 @@ class GroupChunkTranslator
FieldDataInfo column_group_info_;
std::vector<std::string> insert_files_;
std::vector<milvus_storage::RowGroupMetadataVector> row_group_meta_list_;
std::vector<size_t> file_row_group_prefix_sum_;
SchemaPtr schema_;
bool is_sorted_by_pk_;
ChunkedSegmentSealedImpl* chunked_segment_;
internal/core/src/storage/MemFileManagerImpl.cpp
+4 -2
View File
@@ -186,13 +186,15 @@ MemFileManagerImpl::cache_raw_data_to_memory_internal(const Config& config) {
std::vector<FieldDataPtr>
MemFileManagerImpl::cache_raw_data_to_memory_storage_v2(const Config& config) {
auto data_type = index::GetValueFromConfig<DataType>(config, DATA_TYPE_KEY);
AssertInfo(data_type.has_value(), "data type is empty when build index");
AssertInfo(data_type.has_value(),
"[StorageV2] data type is empty when build index");
auto dim = index::GetValueFromConfig<int64_t>(config, DIM_KEY).value_or(0);
auto segment_insert_files =
index::GetValueFromConfig<std::vector<std::vector<std::string>>>(
config, SEGMENT_INSERT_FILES_KEY);
AssertInfo(segment_insert_files.has_value(),
"insert file paths for storage v2 is empty when build index");
"[StorageV2] insert file paths for storage v2 is empty when "
"build index");
auto remote_files = segment_insert_files.value();
for (auto& files : remote_files) {
SortByPath(files);
internal/core/src/storage/Util.cpp
+15 -12
View File
@@ -1068,14 +1068,15 @@ GetFieldDatasFromStorageV2(std::vector<std::vector<std::string>>& remote_files,
DataType data_type,
int64_t dim,
milvus_storage::ArrowFileSystemPtr fs) {
AssertInfo(remote_files.size() > 0, "remote files size is 0");
AssertInfo(remote_files.size() > 0, "[StorageV2] remote files size is 0");
std::vector<FieldDataPtr> field_data_list;
// remote files might not followed the sequence of column group id,
// so we need to put into map<column_group_id, remote_chunk_files>
std::unordered_map<int64_t, std::vector<std::string>> column_group_files;
for (auto& remote_chunk_files : remote_files) {
AssertInfo(remote_chunk_files.size() > 0, "remote files size is 0");
AssertInfo(remote_chunk_files.size() > 0,
"[StorageV2] remote files size is 0");
int64_t group_id = ExtractGroupIdFromPath(remote_chunk_files[0]);
column_group_files[group_id] = remote_chunk_files;
}
@@ -1090,7 +1091,8 @@ GetFieldDatasFromStorageV2(std::vector<std::vector<std::string>>& remote_files,
column_group_id = field_id;
}
AssertInfo(remote_chunk_files.size() > 0, "remote files size is 0");
AssertInfo(remote_chunk_files.size() > 0,
"[StorageV2] remote files size is 0");
// find column offset
milvus_storage::FieldIDList field_id_list = storage::GetFieldIDList(
@@ -1105,14 +1107,15 @@ GetFieldDatasFromStorageV2(std::vector<std::vector<std::string>>& remote_files,
// field not found, must be newly added field, return empty resultset
if (col_offset == -1) {
LOG_INFO(
"field {} not found in column group {}, return empty result set",
"[StorageV2] field {} not found in column group {}, return empty "
"result set",
field_id,
column_group_id);
return field_data_list;
}
AssertInfo(fs != nullptr,
"storage v2 arrow file system is not initialized");
"[StorageV2] storage v2 arrow file system is not initialized");
// set up channel for arrow reader
auto field_data_info = FieldDataInfo();
@@ -1138,10 +1141,10 @@ GetFieldDatasFromStorageV2(std::vector<std::vector<std::string>>& remote_files,
auto field_schema = reader->schema()->field(col_offset)->Copy();
auto arrow_schema = arrow::schema({field_schema});
auto status = reader->Close();
AssertInfo(
status.ok(),
"failed to close file reader when get arrow schema from file: " +
column_group_file + " with error: " + status.ToString());
AssertInfo(status.ok(),
"[StorageV2] failed to close file reader when get arrow "
"schema from file: " +
column_group_file + " with error: " + status.ToString());
// split row groups for parallel reading
auto strategy = std::make_unique<segcore::ParallelDegreeSplitStrategy>(
@@ -1159,9 +1162,9 @@ GetFieldDatasFromStorageV2(std::vector<std::vector<std::string>>& remote_files,
while (field_data_info.arrow_reader_channel->pop(r)) {
size_t num_rows = 0;
std::vector<std::shared_ptr<arrow::ChunkedArray>> chunked_arrays;
for (const auto& [row_group_id, table] : r->arrow_tables) {
num_rows += table->num_rows();
chunked_arrays.push_back(table->column(col_offset));
for (const auto& table_info : r->arrow_tables) {
num_rows += table_info.table->num_rows();
chunked_arrays.push_back(table_info.table->column(col_offset));
}
auto field_data = storage::CreateFieldData(
data_type, field_schema->nullable(), dim, num_rows);
internal/core/unittest/test_group_chunk_translator.cpp
+146 -13
View File
@@ -101,18 +101,22 @@ TEST_P(GroupChunkTranslatorTest, TestWithMmap) {
std::unordered_map<FieldId, FieldMeta> field_metas = schema_->get_fields();
auto column_group_info = FieldDataInfo(0, 3000, "");
auto translator = std::make_unique<GroupChunkTranslator>(segment_id_,
field_metas,
column_group_info,
paths_,
use_mmap,
schema_->get_field_ids().size(),
milvus::proto::common::LoadPriority::LOW);
auto translator = std::make_unique<GroupChunkTranslator>(
segment_id_,
field_metas,
column_group_info,
paths_,
use_mmap,
schema_->get_field_ids().size(),
milvus::proto::common::LoadPriority::LOW);
// num cells - get the expected number from the file directly
auto fr = std::make_shared<milvus_storage::FileRowGroupReader>(fs_, paths_[0]);
auto expected_num_cells = fr->file_metadata()->GetRowGroupMetadataVector().size();
auto row_group_metadata_vector = fr->file_metadata()->GetRowGroupMetadataVector();
auto fr =
std::make_shared<milvus_storage::FileRowGroupReader>(fs_, paths_[0]);
auto expected_num_cells =
fr->file_metadata()->GetRowGroupMetadataVector().size();
auto row_group_metadata_vector =
fr->file_metadata()->GetRowGroupMetadataVector();
auto status = fr->Close();
AssertInfo(status.ok(), "failed to close file reader");
EXPECT_EQ(translator->num_cells(), expected_num_cells);
@@ -125,14 +129,13 @@ TEST_P(GroupChunkTranslatorTest, TestWithMmap) {
// key
EXPECT_EQ(translator->key(), "seg_0_cg_0");
// estimated byte size
for (size_t i = 0; i < translator->num_cells(); ++i) {
auto [file_idx, row_group_idx] =
translator->get_file_and_row_group_index(i);
// Get the expected memory size from the file directly
auto expected_memory_size = static_cast<int64_t>(row_group_metadata_vector.Get(row_group_idx).memory_size());
auto expected_memory_size = static_cast<int64_t>(
row_group_metadata_vector.Get(row_group_idx).memory_size());
auto usage = translator->estimated_byte_size_of_cell(i);
EXPECT_EQ(usage.memory_bytes, expected_memory_size);
}
@@ -171,6 +174,136 @@ TEST_P(GroupChunkTranslatorTest, TestWithMmap) {
}
}
TEST_P(GroupChunkTranslatorTest, TestMultipleFiles) {
auto use_mmap = GetParam();
std::unordered_map<FieldId, FieldMeta> field_metas = schema_->get_fields();
// Create multiple files for testing
std::vector<std::string> multi_file_paths;
std::vector<int64_t> expected_row_groups_per_file;
int64_t total_rows = 0;
// Create 3 files with different numbers of row groups
for (int file_idx = 0; file_idx < 3; ++file_idx) {
std::string file_path =
path_ + "/multi_file_" + std::to_string(file_idx) + ".parquet";
multi_file_paths.push_back(file_path);
int64_t per_batch = 1000;
int64_t n_batch =
2 + file_idx; // Different number of batches per file: 2, 3, 4
int64_t dim = 128;
auto column_groups = std::vector<std::vector<int>>{
{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}};
auto writer_memory = 16 * 1024 * 1024;
auto storage_config = milvus_storage::StorageConfig();
std::vector<std::string> single_file_paths{file_path};
milvus_storage::PackedRecordBatchWriter writer(fs_,
single_file_paths,
arrow_schema_,
storage_config,
column_groups,
writer_memory);
for (int64_t i = 0; i < n_batch; i++) {
auto dataset = DataGen(schema_, per_batch);
auto record_batch =
ConvertToArrowRecordBatch(dataset, dim, arrow_schema_);
total_rows += record_batch->num_rows();
EXPECT_TRUE(writer.Write(record_batch).ok());
}
EXPECT_TRUE(writer.Close().ok());
// Get the number of row groups in this file
auto fr = std::make_shared<milvus_storage::FileRowGroupReader>(
fs_, file_path);
expected_row_groups_per_file.push_back(
fr->file_metadata()->GetRowGroupMetadataVector().size());
auto status = fr->Close();
AssertInfo(status.ok(), "failed to close file reader");
}
auto column_group_info = FieldDataInfo(0, total_rows, "");
auto translator = std::make_unique<GroupChunkTranslator>(
segment_id_,
field_metas,
column_group_info,
multi_file_paths,
use_mmap,
schema_->get_field_ids().size(),
milvus::proto::common::LoadPriority::LOW);
// Test total number of cells across all files
int64_t expected_total_cells = 0;
for (auto row_groups : expected_row_groups_per_file) {
expected_total_cells += row_groups;
}
EXPECT_EQ(translator->num_cells(), expected_total_cells);
// Test get_file_and_row_group_index for cids across different files
int64_t cid_offset = 0;
for (size_t file_idx = 0; file_idx < expected_row_groups_per_file.size();
++file_idx) {
for (int64_t row_group_idx = 0;
row_group_idx < expected_row_groups_per_file[file_idx];
++row_group_idx) {
auto cid = cid_offset + row_group_idx;
auto [actual_file_idx, actual_row_group_idx] =
translator->get_file_and_row_group_index(cid);
EXPECT_EQ(actual_file_idx, file_idx);
EXPECT_EQ(actual_row_group_idx, row_group_idx);
}
cid_offset += expected_row_groups_per_file[file_idx];
}
// Test get_cells with cids from the same file
std::vector<cachinglayer::cid_t> same_file_cids = {0,
1}; // Both from file 0
auto same_file_cells = translator->get_cells(same_file_cids);
EXPECT_EQ(same_file_cells.size(), same_file_cids.size());
int i = 0;
for (const auto& [cid, chunk] : same_file_cells) {
EXPECT_EQ(cid, same_file_cids[i++]);
}
// Test get_cells with cids in reverse order to test sorting
std::vector<cachinglayer::cid_t> cross_file_cids = {4, 7, 0};
auto cells = translator->get_cells(cross_file_cids);
std::vector<cachinglayer::cid_t> returned_cids = {0, 4, 7};
i = 0;
for (const auto& [cid, chunk] : cells) {
EXPECT_EQ(cid, returned_cids[i++]);
}
// Test estimated byte size for cids across different files
for (size_t i = 0; i < translator->num_cells(); ++i) {
auto [file_idx, row_group_idx] =
translator->get_file_and_row_group_index(i);
auto usage = translator->estimated_byte_size_of_cell(i);
// Get the expected memory size from the corresponding file
auto fr = std::make_shared<milvus_storage::FileRowGroupReader>(
fs_, multi_file_paths[file_idx]);
auto row_group_metadata_vector =
fr->file_metadata()->GetRowGroupMetadataVector();
auto expected_memory_size = static_cast<int64_t>(
row_group_metadata_vector.Get(row_group_idx).memory_size());
auto status = fr->Close();
AssertInfo(status.ok(), "failed to close file reader");
EXPECT_EQ(usage.memory_bytes, expected_memory_size);
}
// Clean up test files
for (const auto& file_path : multi_file_paths) {
if (std::filesystem::exists(file_path)) {
std::filesystem::remove(file_path);
}
}
}
INSTANTIATE_TEST_SUITE_P(GroupChunkTranslatorTest,
GroupChunkTranslatorTest,
testing::Bool());
internal/core/unittest/test_growing_storage_v2.cpp
+20 -14
View File
@@ -227,11 +227,11 @@ TEST_F(TestGrowingStorageV2, LoadWithStrategy) {
int64_t current_row_group = 0;
while (channel->pop(wrapper)) {
for (const auto& [row_group_id, table] : wrapper->arrow_tables) {
for (const auto& table_info : wrapper->arrow_tables) {
// Verify batch size matches row group metadata
EXPECT_EQ(table->num_rows(),
EXPECT_EQ(table_info.table->num_rows(),
row_group_metadata.Get(current_row_group).row_num());
total_rows += table->num_rows();
total_rows += table_info.table->num_rows();
current_row_group++;
}
}
@@ -259,11 +259,12 @@ TEST_F(TestGrowingStorageV2, LoadWithStrategy) {
int64_t total_rows = 0;
while (channel->pop(wrapper)) {
for (const auto& [row_group_id, table] : wrapper->arrow_tables) {
for (const auto& table_info : wrapper->arrow_tables) {
// Verify batch size matches row group metadata
EXPECT_EQ(table->num_rows(),
row_group_metadata.Get(row_group_id).row_num());
total_rows += table->num_rows();
EXPECT_EQ(table_info.table->num_rows(),
row_group_metadata.Get(table_info.row_group_index)
.row_num());
total_rows += table_info.table->num_rows();
}
}
@@ -289,14 +290,19 @@ TEST_F(TestGrowingStorageV2, LoadWithStrategy) {
std::vector<int64_t> selected_row_groups = {0, 2};
while (channel->pop(wrapper)) {
for (const auto& [row_group_id, table] : wrapper->arrow_tables) {
// row_group_id is the actual row group ID (0 or 2), not an index
for (const auto& table_info : wrapper->arrow_tables) {
// row_group_index is the actual row group ID (0 or 2), not an index
// We need to find its position in selected_row_groups
auto it = std::find(selected_row_groups.begin(), selected_row_groups.end(), row_group_id);
ASSERT_NE(it, selected_row_groups.end()) << "Row group " << row_group_id << " not found in selected_row_groups";
EXPECT_EQ(table->num_rows(),
row_group_metadata.Get(row_group_id).row_num());
total_rows += table->num_rows();
auto it = std::find(selected_row_groups.begin(),
selected_row_groups.end(),
table_info.row_group_index);
ASSERT_NE(it, selected_row_groups.end())
<< "Row group " << table_info.row_group_index
<< " not found in selected_row_groups";
EXPECT_EQ(table_info.table->num_rows(),
row_group_metadata.Get(table_info.row_group_index)
.row_num());
total_rows += table_info.table->num_rows();
}
}