Fix disagg

dbms/src/Storages/StorageDisaggregated.cpp

@@ -64,16 +64,18 @@ void StorageDisaggregated::read(
     return readThroughS3(exec_context, group_builder, db_context, num_streams);
 }
 
-std::vector<StorageDisaggregated::RemoteTableRange> StorageDisaggregated::buildRemoteTableRanges()
+std::tuple<std::vector<StorageDisaggregated::RemoteTableRange>, UInt64> StorageDisaggregated::buildRemoteTableRanges()
 {
     std::unordered_map<TableID, RegionRetryList> all_remote_regions;
+    UInt64 region_num = 0;
     for (auto physical_table_id : table_scan.getPhysicalTableIDs())
     {
         const auto & table_regions_info = context.getDAGContext()->getTableRegionsInfoByTableID(physical_table_id);
 
         RUNTIME_CHECK_MSG(
             table_regions_info.local_regions.empty(),
             "in disaggregated_compute_mode, local_regions should be empty");
+        region_num += table_regions_info.remote_regions.size();
         for (const auto & reg : table_regions_info.remote_regions)
             all_remote_regions[physical_table_id].emplace_back(std::cref(reg));
     }
@@ -87,7 +89,7 @@ std::vector<StorageDisaggregated::RemoteTableRange> StorageDisaggregated::buildR
         auto key_ranges = RemoteRequest::buildKeyRanges(remote_regions);
         remote_table_ranges.emplace_back(RemoteTableRange{physical_table_id, key_ranges});
     }
-    return remote_table_ranges;
+    return std::make_tuple(std::move(remote_table_ranges), region_num);
 }
 
 std::vector<pingcap::coprocessor::BatchCopTask> StorageDisaggregated::buildBatchCopTasks(

dbms/src/Storages/StorageDisaggregated.h

@@ -125,7 +125,7 @@ class StorageDisaggregated : public IStorage
         size_t num_streams);
 
     using RemoteTableRange = std::pair<TableID, pingcap::coprocessor::KeyRanges>;
-    std::vector<RemoteTableRange> buildRemoteTableRanges();
+    std::tuple<std::vector<RemoteTableRange>, UInt64> buildRemoteTableRanges();
     std::vector<pingcap::coprocessor::BatchCopTask> buildBatchCopTasks(
         const std::vector<RemoteTableRange> & remote_table_ranges,
         const pingcap::kv::LabelFilter & label_filter);

dbms/src/Storages/StorageDisaggregatedRemote.cpp

@@ -199,7 +199,8 @@ DM::SegmentReadTasks StorageDisaggregated::buildReadTask(
     // First split the read task for different write nodes.
     // For each write node, a BatchCopTask is built.
     {
-        auto remote_table_ranges = buildRemoteTableRanges();
+        auto [remote_table_ranges, region_num] = buildRemoteTableRanges();
+        scan_context->setRegionNumOfCurrentInstance(region_num);
         // only send to tiflash node with label [{"engine":"tiflash"}, {"engine-role":"write"}]
         const auto label_filter = pingcap::kv::labelFilterOnlyTiFlashWriteNode;
         batch_cop_tasks = buildBatchCopTasks(remote_table_ranges, label_filter);
@@ -227,7 +228,7 @@ DM::SegmentReadTasks StorageDisaggregated::buildReadTask(
     {
         // TODO
     }
-
+    scan_context->num_segments = output_seg_tasks.size();
     return output_seg_tasks;
 }