Repository: milvus-io/milvus
File: internal/core/src/segcore/memory_planner.cpp (C++)

Code before:
    auto& pool = ThreadPools::GetThreadPool(milvus::PriorityForLoad(priority));
    auto remaining = std::make_shared<std::atomic<size_t>>(batches.size());
    auto reader_memory_limit =
        std::max<int64_t>(memory_limit / static_cast<int64_t>(batches.size()),
                          FILE_SLICE_SIZE.load());
    auto shared_factory =
        std::make_shared<BatchReaderFactory>(std::move(reader_factory));

    std::vector<std::future<void>> futures;
    futures.reserve(batches.size());

    for (auto& batch : batches) {
        futures.emplace_back(pool.Submit([batch = std::move(batch),
                                          shared_factory,
                                          reader_memory_limit,
                                          channel,
                                          remaining,
                                          op_ctx]() {
            auto task_guard = folly::makeGuard([&channel, &remaining]() {
                if (remaining->fetch_sub(1) == 1) {
                    channel->close();
                }
            });
            CheckCancellation(op_ctx, -1, "LoadCellBatchAsync");

            auto tables_result = (*shared_factory)(batch.file_idx,
                                                   batch.rg_offset,
                                                   batch.rg_count,
                                                   reader_memory_limit);
            AssertInfo(tables_result.ok(),
                       "[StorageV2] Failed to read batch: " +
                           tables_result.status().ToString());
            auto all_tables = std::move(tables_result).ValueOrDie();

            int64_t table_offset = 0;
            for (const auto& cell : batch.cells) {
                auto cell_result = std::make_shared<CellLoadResult>();
                cell_result->cid = cell.cid;
                cell_result->tables.reserve(cell.rg_count);
                for (int64_t i = 0; i < cell.rg_count; ++i) {
                    cell_result->tables.push_back(
                        std::move(all_tables[table_offset + i]));
                }
                table_offset += cell.rg_count;
                channel->push(std::move(cell_result));
            }
        }));
    }

    return futures;
}

Review comment (performance):
No size check on `all_tables` before indexing — if the factory ever returns fewer tables than `batch.rg_count`, this is an out-of-bounds access. Consider adding `AssertInfo(all_tables.size() == static_cast<size_t>(batch.rg_count), ...)` after the factory call.

Code after:

    std::vector<std::future<void>> futures;
    futures.reserve(batches.size());

    for (auto& batch : batches) {
        futures.emplace_back(pool.Submit([batch = std::move(batch),
                                          shared_factory,
                                          reader_memory_limit,
                                          channel,
                                          remaining,
                                          op_ctx]() {
            auto task_guard = folly::makeGuard([&channel, &remaining]() {
                if (remaining->fetch_sub(1) == 1) {
                    channel->close();
                }
            });
            CheckCancellation(op_ctx, -1, "LoadCellBatchAsync");

            auto tables_result = (*shared_factory)(batch.file_idx,
                                                   batch.rg_offset,
                                                   batch.rg_count,
                                                   reader_memory_limit);
            AssertInfo(tables_result.ok(),
                       "[StorageV2] Failed to read batch: " +
                           tables_result.status().ToString());
            auto all_tables = std::move(tables_result).ValueOrDie();
            AssertInfo(all_tables.size() == static_cast<size_t>(batch.rg_count),
                       "reader returns less tables than expected, batch rg "
                       "count: {}, result size: {}",
                       batch.rg_count,
                       all_tables.size());

            int64_t table_offset = 0;
            for (const auto& cell : batch.cells) {
                auto cell_result = std::make_shared<CellLoadResult>();
                cell_result->cid = cell.cid;
                cell_result->tables.reserve(cell.rg_count);
                for (int64_t i = 0; i < cell.rg_count; ++i) {
                    cell_result->tables.push_back(
                        std::move(all_tables[table_offset + i]));
                }
                table_offset += cell.rg_count;
                channel->push(std::move(cell_result));
            }
        }));
    }

    return futures;
}

BatchReaderFactory