enhance: map multi row groups into one cache cell

Signed-off-by: Shawn Wang <[email protected]>

Author: sparknack

internal/core/src/segcore/storagev2translator/GroupCTMeta.h

@@ -19,10 +19,21 @@
 
 namespace milvus::segcore::storagev2translator {
 
+// Number of row groups (parquet row groups) merged into one cache cell,
+// for now it is a constant.
+// herierchy: 1 group chunk <-> 1 cache cell <-> kRowGroupsPerCell row groups
+constexpr size_t kRowGroupsPerCell = 4;
+
 struct GroupCTMeta : public milvus::cachinglayer::Meta {
+    // num_rows_until_chunk_[i] = total rows(prefix sum) in cells [0, i-1]
+    // the size of num_rows_until_chunk_ is num_cells + 1
     std::vector<int64_t> num_rows_until_chunk_;
+    // memory size for each group chunk(cache cell)
     std::vector<int64_t> chunk_memory_size_;
     size_t num_fields_;
+    // total number of row groups
+    size_t total_row_groups_;
+
     GroupCTMeta(size_t num_fields,
                 milvus::cachinglayer::StorageType storage_type,
                 milvus::cachinglayer::CellIdMappingMode cell_id_mapping_mode,
@@ -34,7 +45,16 @@ struct GroupCTMeta : public milvus::cachinglayer::Meta {
                                      cell_data_type,
                                      cache_warmup_policy,
                                      support_eviction),
-          num_fields_(num_fields) {
+          num_fields_(num_fields),
+          total_row_groups_(0) {
+    }
+
+    // Get the range of row groups for a cell [start, end)
+    std::pair<size_t, size_t>
+    get_row_group_range(size_t cid) const {
+        size_t start = cid * kRowGroupsPerCell;
+        size_t end = std::min(start + kRowGroupsPerCell, total_row_groups_);
+        return {start, end};
     }
 };
 

internal/core/src/segcore/storagev2translator/GroupChunkTranslator.cpp

@@ -144,19 +144,39 @@ GroupChunkTranslator::GroupChunkTranslator(
                                              file_metas.size());
     }
 
-    meta_.num_rows_until_chunk_.reserve(total_row_groups + 1);
-    meta_.chunk_memory_size_.reserve(total_row_groups);
-
-    meta_.num_rows_until_chunk_.push_back(0);
+    // Collect row group sizes and row counts
+    std::vector<int64_t> row_group_row_counts;
+    row_group_sizes_.reserve(total_row_groups);
+    row_group_row_counts.reserve(total_row_groups);
     for (const auto& row_group_meta : row_group_meta_list_) {
         for (int i = 0; i < row_group_meta.size(); ++i) {
-            meta_.num_rows_until_chunk_.push_back(
-                meta_.num_rows_until_chunk_.back() +
-                row_group_meta.Get(i).row_num());
-            meta_.chunk_memory_size_.push_back(
-                row_group_meta.Get(i).memory_size());
+            row_group_sizes_.push_back(row_group_meta.Get(i).memory_size());
+            row_group_row_counts.push_back(row_group_meta.Get(i).row_num());
+        }
+    }
+
+    // Build cell mapping: each cell contains up to kRowGroupsPerCell row groups
+    meta_.total_row_groups_ = total_row_groups;
+    size_t num_cells =
+        (total_row_groups + kRowGroupsPerCell - 1) / kRowGroupsPerCell;
+
+    // Merge row groups into group chunks(cache cells)
+    meta_.num_rows_until_chunk_.reserve(num_cells + 1);
+    meta_.num_rows_until_chunk_.push_back(0);
+    meta_.chunk_memory_size_.reserve(num_cells);
+
+    int64_t cumulative_rows = 0;
+    for (size_t cell_id = 0; cell_id < num_cells; ++cell_id) {
+        auto [start, end] = meta_.get_row_group_range(cell_id);
+        int64_t cell_size = 0;
+        for (size_t i = start; i < end; ++i) {
+            cumulative_rows += row_group_row_counts[i];
+            cell_size += row_group_sizes_[i];
         }
+        meta_.num_rows_until_chunk_.push_back(cumulative_rows);
+        meta_.chunk_memory_size_.push_back(cell_size);
     }
+
     AssertInfo(
         meta_.num_rows_until_chunk_.back() == column_group_info_.row_count,
         fmt::format(
@@ -165,6 +185,14 @@ GroupChunkTranslator::GroupChunkTranslator(
             column_group_info_.field_id,
             meta_.num_rows_until_chunk_.back(),
             column_group_info_.row_count));
+
+    LOG_INFO(
+        "[StorageV2] translator {} merged {} row groups into {} cells ({} "
+        "row groups per cell)",
+        key_,
+        total_row_groups,
+        num_cells,
+        kRowGroupsPerCell);
 }
 
 GroupChunkTranslator::~GroupChunkTranslator() {
@@ -184,10 +212,14 @@ std::pair<milvus::cachinglayer::ResourceUsage,
           milvus::cachinglayer::ResourceUsage>
 GroupChunkTranslator::estimated_byte_size_of_cell(
     milvus::cachinglayer::cid_t cid) const {
-    auto [file_idx, row_group_idx] = get_file_and_row_group_index(cid);
-    auto& row_group_meta = row_group_meta_list_[file_idx].Get(row_group_idx);
+    AssertInfo(cid < meta_.chunk_memory_size_.size(),
+               fmt::format("[StorageV2] translator {} cid {} is out of range. "
+                           "Total cells: {}",
+                           key_,
+                           cid,
+                           meta_.chunk_memory_size_.size()));
 
-    auto cell_sz = static_cast<int64_t>(row_group_meta.memory_size());
+    auto cell_sz = meta_.chunk_memory_size_[cid];
 
     if (use_mmap_) {
         // why double the disk size for loading?
@@ -205,26 +237,29 @@ GroupChunkTranslator::key() const {
 }
 
 std::pair<size_t, size_t>
-GroupChunkTranslator::get_file_and_row_group_index(
-    milvus::cachinglayer::cid_t cid) const {
+GroupChunkTranslator::get_file_and_row_group_offset(
+    size_t global_row_group_idx) const {
     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]};
+        if (global_row_group_idx < file_row_group_prefix_sum_[file_idx + 1]) {
+            return {
+                file_idx,
+                global_row_group_idx - file_row_group_prefix_sum_[file_idx]};
         }
     }
 
-    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()));
+    AssertInfo(
+        false,
+        fmt::format("[StorageV2] translator {} global_row_group_idx {} is out "
+                    "of range. Total row groups across all files: {}",
+                    key_,
+                    global_row_group_idx,
+                    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 {
+GroupChunkTranslator::get_global_row_group_idx(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: {}",
@@ -254,12 +289,31 @@ GroupChunkTranslator::get_cells(const std::vector<cachinglayer::cid_t>& cids) {
         cells;
     cells.reserve(cids.size());
 
-    // Create row group lists for requested cids
-    std::vector<std::vector<int64_t>> row_group_lists(insert_files_.size());
+    std::vector<cachinglayer::cid_t> sorted_cids(cids.begin(), cids.end());
+    std::sort(sorted_cids.begin(), sorted_cids.end());
+
+    // Collect all row group indices needed for the requested cells
+    std::vector<size_t> needed_row_group_indices;
+    needed_row_group_indices.reserve(kRowGroupsPerCell * sorted_cids.size());
+    for (auto cid : sorted_cids) {
+        AssertInfo(cid < meta_.chunk_memory_size_.size(),
+                   fmt::format("[StorageV2] translator {} cid {} is out of "
+                               "range. Total cells: {}",
+                               key_,
+                               cid,
+                               meta_.chunk_memory_size_.size()));
+        auto [start, end] = meta_.get_row_group_range(cid);
+        for (size_t i = start; i < end; ++i) {
+            needed_row_group_indices.push_back(i);
+        }
+    }
 
-    for (auto cid : cids) {
-        auto [file_idx, row_group_idx] = get_file_and_row_group_index(cid);
-        row_group_lists[file_idx].push_back(row_group_idx);
+    // Create row group lists for file loading
+    std::vector<std::vector<int64_t>> row_group_lists(insert_files_.size());
+    // no need to sort the needed_row_group_indices because it is built from sorted_cids
+    for (auto rg_idx : needed_row_group_indices) {
+        auto [file_idx, row_group_off] = get_file_and_row_group_offset(rg_idx);
+        row_group_lists[file_idx].push_back(row_group_off);
     }
 
     auto parallel_degree =
@@ -288,57 +342,81 @@ GroupChunkTranslator::get_cells(const std::vector<cachinglayer::cid_t>& cids) {
         key_,
         column_group_info_.field_id);
 
+    // Collect loaded tables by row group index
+    std::unordered_map<size_t, std::shared_ptr<arrow::Table>> row_group_tables;
+    row_group_tables.reserve(needed_row_group_indices.size());
+
     std::shared_ptr<milvus::ArrowDataWrapper> r;
-    std::unordered_set<cachinglayer::cid_t> filled_cids;
-    filled_cids.reserve(cids.size());
     while (channel->pop(r)) {
         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);
+            // Convert file_index and row_group_index (file inner index, not global index) to global row group index
+            auto rg_idx = get_global_row_group_idx(table_info.file_index,
+                                                   table_info.row_group_index);
+            row_group_tables[rg_idx] = table_info.table;
         }
     }
 
     // access underlying feature to get exception if any
     load_future.get();
 
-    // Verify all requested cids have been filled
-    for (auto cid : cids) {
-        AssertInfo(filled_cids.find(cid) != filled_cids.end(),
-                   "[StorageV2] translator {} cid {} was not filled, missing "
-                   "row group id {}",
-                   key_,
-                   cid,
-                   cid);
+    // Build cells from collected tables
+    std::unordered_set<cachinglayer::cid_t> filled_cids;
+    filled_cids.reserve(sorted_cids.size());
+    for (auto cid : sorted_cids) {
+        if (filled_cids.count(cid) > 0) {
+            continue;  // Already processed this cell (handles duplicates)
+        }
+
+        auto [start, end] = meta_.get_row_group_range(cid);
+        std::vector<std::shared_ptr<arrow::Table>> tables;
+        tables.reserve(end - start);
+
+        for (size_t i = start; i < end; ++i) {
+            auto it = row_group_tables.find(i);
+            AssertInfo(it != row_group_tables.end(),
+                       fmt::format("[StorageV2] translator {} row group {} "
+                                   "for cell {} was not loaded",
+                                   key_,
+                                   i,
+                                   cid));
+            tables.push_back(it->second);
+        }
+
+        cells.emplace_back(cid, load_group_chunk(tables, cid));
+        filled_cids.insert(cid);
     }
+
     return cells;
 }
 
 std::unique_ptr<milvus::GroupChunk>
 GroupChunkTranslator::load_group_chunk(
-    const std::shared_ptr<arrow::Table>& table,
+    const std::vector<std::shared_ptr<arrow::Table>>& tables,
     const milvus::cachinglayer::cid_t cid) {
-    AssertInfo(table != nullptr, "arrow table is nullptr");
-    // Create chunks for each field in this batch
-    std::unordered_map<FieldId, std::shared_ptr<Chunk>> chunks;
-    // Iterate through field_id_list to get field_id and create chunk
+    AssertInfo(!tables.empty(), "tables vector is empty");
+    for (const auto& table : tables) {
+        AssertInfo(table != nullptr, "arrow table is nullptr");
+    }
+
+    // Use the first table's schema as reference for field iteration
+    const auto& schema = tables[0]->schema();
+
+    // Collect field info and merge array vectors from all tables
     std::vector<FieldId> field_ids;
     std::vector<FieldMeta> field_metas;
     std::vector<arrow::ArrayVector> array_vecs;
-    field_metas.reserve(table->schema()->num_fields());
-    array_vecs.reserve(table->schema()->num_fields());
+    field_ids.reserve(schema->num_fields());
+    field_metas.reserve(schema->num_fields());
+    array_vecs.reserve(schema->num_fields());
 
-    for (int i = 0; i < table->schema()->num_fields(); ++i) {
-        AssertInfo(table->schema()->field(i)->metadata()->Contains(
+    for (int i = 0; i < schema->num_fields(); ++i) {
+        AssertInfo(schema->field(i)->metadata()->Contains(
                        milvus_storage::ARROW_FIELD_ID_KEY),
                    "[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)
+                   schema->field(i)->name());
+        auto field_id = std::stoll(schema->field(i)
                                        ->metadata()
                                        ->Get(milvus_storage::ARROW_FIELD_ID_KEY)
                                        ->data());
@@ -355,12 +433,22 @@ GroupChunkTranslator::load_group_chunk(
             key_,
             fid.get());
         const auto& field_meta = it->second;
-        const arrow::ArrayVector& array_vec = table->column(i)->chunks();
+
+        // Merge array vectors from all tables for this field
+        // All tables in a cell come from the same column group with consistent schema
+        arrow::ArrayVector merged_array_vec;
+        for (const auto& table : tables) {
+            const arrow::ArrayVector& array_vec = table->column(i)->chunks();
+            merged_array_vec.insert(
+                merged_array_vec.end(), array_vec.begin(), array_vec.end());
+        }
+
         field_ids.push_back(fid);
         field_metas.push_back(field_meta);
-        array_vecs.push_back(array_vec);
+        array_vecs.push_back(std::move(merged_array_vec));
     }
 
+    std::unordered_map<FieldId, std::shared_ptr<Chunk>> chunks;
     if (!use_mmap_) {
         chunks = create_group_chunk(field_ids, field_metas, array_vecs);
     } else {

internal/core/src/segcore/storagev2translator/GroupChunkTranslator.h

@@ -66,11 +66,10 @@ class GroupChunkTranslator
     get_cells(const std::vector<milvus::cachinglayer::cid_t>& cids) override;
 
     std::pair<size_t, size_t>
-    get_file_and_row_group_index(milvus::cachinglayer::cid_t cid) const;
+    get_file_and_row_group_offset(size_t global_row_group_idx) const;
 
     milvus::cachinglayer::cid_t
-    get_cid_from_file_and_row_group_index(size_t file_idx,
-                                          size_t row_group_idx) const;
+    get_global_row_group_idx(size_t file_idx, size_t row_group_idx) const;
 
     milvus::cachinglayer::Meta*
     meta() override {
@@ -83,12 +82,8 @@ class GroupChunkTranslator
         constexpr int64_t MIN_STORAGE_BYTES = 1 * 1024 * 1024;
         int64_t total_size = 0;
         for (auto cid : cids) {
-            auto [file_idx, row_group_idx] = get_file_and_row_group_index(cid);
-            auto& row_group_meta =
-                row_group_meta_list_[file_idx].Get(row_group_idx);
             total_size +=
-                std::max(static_cast<int64_t>(row_group_meta.memory_size()),
-                         MIN_STORAGE_BYTES);
+                std::max(meta_.chunk_memory_size_[cid], MIN_STORAGE_BYTES);
         }
         return total_size;
     }
@@ -104,8 +99,9 @@ class GroupChunkTranslator
     }
 
  private:
+    // Load a single cell which may contain multiple row groups
     std::unique_ptr<milvus::GroupChunk>
-    load_group_chunk(const std::shared_ptr<arrow::Table>& table,
+    load_group_chunk(const std::vector<std::shared_ptr<arrow::Table>>& tables,
                      const milvus::cachinglayer::cid_t cid);
 
     int64_t segment_id_;
@@ -127,6 +123,9 @@ class GroupChunkTranslator
         milvus::proto::common::LoadPriority::HIGH};
     std::vector<std::shared_ptr<parquet::FileMetaData>> parquet_file_metadata_;
     std::map<int64_t, milvus_storage::ColumnOffset> field_id_mapping_;
+
+    // Row group sizes before cell merging
+    std::vector<int64_t> row_group_sizes_;
 };
 
 }  // namespace milvus::segcore::storagev2translator