Share the block_queue for partition table scan

Signed-off-by: JaySon-Huang <tshent@qq.com>

Author: JaySon-Huang

dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp

@@ -32,6 +32,7 @@
 #include <Flash/Coprocessor/DAGStorageInterpreter.h>
 #include <Flash/Coprocessor/InterpreterUtils.h>
 #include <Flash/Coprocessor/RemoteRequest.h>
+#include <Flash/Coprocessor/TableScanPipelineExecBuilder.h>
 #include <Flash/Coprocessor/collectOutputFieldTypes.h>
 #include <Interpreters/Context.h>
 #include <Interpreters/SharedContexts/Disagg.h>
@@ -43,6 +44,7 @@
 #include <Operators/UnorderedSourceOp.h>
 #include <Parsers/makeDummyQuery.h>
 #include <Storages/DeltaMerge/Index/VectorIndex/Stream/Ctx.h>
+#include <Storages/DeltaMerge/ReadThread/SharedBlockQueue.h>
 #include <Storages/DeltaMerge/Remote/DisaggSnapshot.h>
 #include <Storages/DeltaMerge/Remote/WNDisaggSnapshotManager.h>
 #include <Storages/DeltaMerge/ScanContext.h>
@@ -917,6 +919,7 @@ LearnerReadSnapshot DAGStorageInterpreter::doBatchCopLearnerRead()
 std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSelectQueryInfos()
 {
     std::unordered_map<TableID, SelectQueryInfo> ret;
+    auto read_queue = std::make_shared<DM::SharedBlockQueue>(log);
     auto create_query_info = [&](Int64 table_id) -> SelectQueryInfo {
         SelectQueryInfo query_info;
         /// to avoid null point exception
@@ -934,6 +937,7 @@ std::unordered_map<TableID, SelectQueryInfo> DAGStorageInterpreter::generateSele
         query_info.req_id = fmt::format("{} table_id={}", log->identifier(), table_id);
         query_info.keep_order = table_scan.keepOrder();
         query_info.is_fast_scan = table_scan.isFastScan();
+        query_info.read_queue = read_queue;
         return query_info;
     };
     RUNTIME_CHECK_MSG(mvcc_query_info->scan_context != nullptr, "Unexpected null scan_context");
@@ -1060,16 +1064,16 @@ Int32 getMaxAllowRetryForLocalRead(const SelectQueryInfo & query_info)
 }
 } // namespace
 
-DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocalStreamsForPhysicalTable(
+void DAGStorageInterpreter::buildLocalStreamsForPhysicalTable(
     const TableID & table_id,
     const SelectQueryInfo & query_info,
     DAGPipeline & pipeline,
+    DM::Remote::DisaggReadSnapshotPtr & disagg_snap,
     size_t max_block_size)
 {
-    DM::Remote::DisaggPhysicalTableReadSnapshotPtr table_snap;
     size_t region_num = query_info.mvcc_query_info->regions_query_info.size();
     if (region_num == 0)
-        return table_snap;
+        return;
 
     assert(storages_with_structure_lock.find(table_id) != storages_with_structure_lock.end());
     auto & storage = storages_with_structure_lock[table_id].storage;
@@ -1087,6 +1091,7 @@ DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocal
     {
         try
         {
+            DM::Remote::DisaggPhysicalTableReadSnapshotPtr table_snap;
             if (!dag_context.is_disaggregated_task)
             {
                 // build local inputstreams
@@ -1114,6 +1119,12 @@ DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocal
             // (by calling `validateQueryInfo`). In case the key ranges of Regions have changed (Region merge/split), those `streams`
             // may contain different data other than expected.
             validateQueryInfo(*query_info.mvcc_query_info, learner_read_snapshot, tmt, log);
+
+            // Only after all streams are built successfully, we add the task to mvcc_query_info
+            if (table_snap)
+            {
+                disagg_snap->addTask(table_id, std::move(table_snap));
+            }
             break;
         }
         catch (RegionException & e)
@@ -1124,6 +1135,8 @@ DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocal
             {
                 // clean all streams from local because we are not sure the correctness of those streams
                 pipeline.streams.clear();
+                // clean table task from read_queue
+                query_info.read_queue->resetTableTask(table_id);
                 if (likely(checkRetriableForBatchCopOrMPP(table_id, query_info, e, num_allow_retry)))
                     continue; // next retry to read from local storage
                 else
@@ -1143,20 +1156,18 @@ DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocal
             throw;
         }
     }
-    return table_snap;
 }
 
-DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocalExecForPhysicalTable(
+void DAGStorageInterpreter::buildLocalExecForPhysicalTable(
     PipelineExecutorContext & exec_context,
-    PipelineExecGroupBuilder & group_builder,
+    TableScanPipelineExecGroupBuilder & group_builder,
     const TableID & table_id,
     const SelectQueryInfo & query_info,
     size_t max_block_size)
 {
-    DM::Remote::DisaggPhysicalTableReadSnapshotPtr table_snap;
     size_t region_num = query_info.mvcc_query_info->regions_query_info.size();
     if (region_num == 0)
-        return table_snap;
+        return;
 
     RUNTIME_CHECK(storages_with_structure_lock.find(table_id) != storages_with_structure_lock.end());
     auto & storage = storages_with_structure_lock[table_id].storage;
@@ -1167,6 +1178,7 @@ DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocal
     {
         try
         {
+            DM::Remote::DisaggPhysicalTableReadSnapshotPtr table_snap;
             if (!dag_context.is_disaggregated_task)
             {
                 storage->read(
@@ -1197,6 +1209,9 @@ DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocal
             // (by calling `validateQueryInfo`). In case the key ranges of Regions have changed (Region merge/split), those `sourceOps`
             // may contain different data other than expected.
             validateQueryInfo(*query_info.mvcc_query_info, learner_read_snapshot, tmt, log);
+
+            // Only after all sourceOps are built and verified, we add the snapshot to group_builder
+            group_builder.addPhysicalTableTask(table_id, std::move(table_snap));
             break;
         }
         catch (RegionException & e)
@@ -1207,6 +1222,8 @@ DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocal
             {
                 // clean all operator from local because we are not sure the correctness of those operators
                 group_builder.reset();
+                // clean table task from read_queue
+                query_info.read_queue->resetTableTask(table_id);
                 if (likely(checkRetriableForBatchCopOrMPP(table_id, query_info, e, num_allow_retry)))
                     continue;
                 else
@@ -1226,7 +1243,6 @@ DM::Remote::DisaggPhysicalTableReadSnapshotPtr DAGStorageInterpreter::buildLocal
             throw;
         }
     }
-    return table_snap;
 }
 
 void DAGStorageInterpreter::buildLocalStreams(DAGPipeline & pipeline, size_t max_block_size)
@@ -1237,6 +1253,7 @@ void DAGStorageInterpreter::buildLocalStreams(DAGPipeline & pipeline, size_t max
         return;
     mvcc_query_info->scan_context->setRegionNumOfCurrentInstance(total_local_region_num);
     const auto table_query_infos = generateSelectQueryInfos();
+    RUNTIME_CHECK_MSG(!table_query_infos.empty(), "No table query info generated for local read");
     bool has_multiple_partitions = table_query_infos.size() > 1;
     // MultiPartitionStreamPool will be disabled in no partition mode or single-partition case
     std::shared_ptr<MultiPartitionStreamPool> stream_pool
@@ -1248,12 +1265,17 @@ void DAGStorageInterpreter::buildLocalStreams(DAGPipeline & pipeline, size_t max
         DAGPipeline current_pipeline;
         const TableID physical_table_id = table_query_info.first;
         const SelectQueryInfo & query_info = table_query_info.second;
-        auto table_snap
-            = buildLocalStreamsForPhysicalTable(physical_table_id, query_info, current_pipeline, max_block_size);
-        if (table_snap)
-        {
-            disaggregated_snap->addTask(physical_table_id, std::move(table_snap));
-        }
+        RUNTIME_CHECK_MSG(
+            query_info.read_queue != nullptr,
+            "read_queue should not be null, table_id={}",
+            physical_table_id);
+        buildLocalStreamsForPhysicalTable(
+            physical_table_id,
+            query_info,
+            current_pipeline,
+            disaggregated_snap,
+            max_block_size);
+
         if (has_multiple_partitions)
             stream_pool->addPartitionStreams(current_pipeline.streams);
         else
@@ -1263,6 +1285,10 @@ void DAGStorageInterpreter::buildLocalStreams(DAGPipeline & pipeline, size_t max
                 current_pipeline.streams.end());
     }
 
+    assert(!table_query_infos.empty()); // check at start of this function
+    auto read_queue = table_query_infos.begin()->second.read_queue;
+    read_queue->finishQueueIfEmpty();
+
     LOG_DEBUG(
         log,
         "local streams built, is_disaggregated_task={} snap_id={}",
@@ -1304,30 +1330,33 @@ void DAGStorageInterpreter::buildLocalExec(
         return;
     mvcc_query_info->scan_context->setRegionNumOfCurrentInstance(total_local_region_num);
     const auto table_query_infos = generateSelectQueryInfos();
+    RUNTIME_CHECK_MSG(!table_query_infos.empty(), "No table query info generated for local read");
     bool has_multiple_partitions = table_query_infos.size() > 1;
     ConcatBuilderPool builder_pool{max_streams, context.getSettingsRef().dt_enable_unordered_concat};
 
     auto disaggregated_snap = std::make_shared<DM::Remote::DisaggReadSnapshot>();
     for (const auto & table_query_info : table_query_infos)
     {
-        PipelineExecGroupBuilder builder;
         const TableID physical_table_id = table_query_info.first;
         const SelectQueryInfo & query_info = table_query_info.second;
-        auto table_snap
-            = buildLocalExecForPhysicalTable(exec_context, builder, physical_table_id, query_info, max_block_size);
-        if (table_snap)
-        {
-            disaggregated_snap->addTask(physical_table_id, std::move(table_snap));
-        }
+        RUNTIME_CHECK_MSG(
+            query_info.read_queue != nullptr,
+            "read_queue should not be null, table_id={}",
+            physical_table_id);
+        TableScanPipelineExecGroupBuilder builder(query_info.read_queue, disaggregated_snap);
+        buildLocalExecForPhysicalTable(exec_context, builder, physical_table_id, query_info, max_block_size);
 
         if (has_multiple_partitions)
             builder_pool.add(builder);
         else
             group_builder.merge(std::move(builder));
     }
 
-    LOG_DEBUG(log, "local sourceOps built, is_disaggregated_task={}", dag_context.is_disaggregated_task);
+    assert(!table_query_infos.empty()); // check at start of this function
+    auto read_queue = table_query_infos.begin()->second.read_queue;
+    read_queue->finishQueueIfEmpty();
 
+    LOG_DEBUG(log, "local sourceOps built, is_disaggregated_task={}", dag_context.is_disaggregated_task);
     if (dag_context.is_disaggregated_task)
     {
         // register the snapshot to manager

dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h

@@ -20,6 +20,7 @@
 #include <Flash/Coprocessor/DAGPipeline.h>
 #include <Flash/Coprocessor/FilterConditions.h>
 #include <Flash/Coprocessor/RemoteRequest.h>
+#include <Flash/Coprocessor/TableScanPipelineExecBuilder.h>
 #include <Flash/Coprocessor/TiDBTableScan.h>
 #include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
 #include <Storages/DeltaMerge/Remote/DisaggSnapshot_fwd.h>
@@ -73,15 +74,16 @@ class DAGStorageInterpreter
         const RegionException & e,
         int num_allow_retry);
 
-    DM::Remote::DisaggPhysicalTableReadSnapshotPtr buildLocalStreamsForPhysicalTable(
+    void buildLocalStreamsForPhysicalTable(
         const TableID & table_id,
         const SelectQueryInfo & query_info,
         DAGPipeline & pipeline,
+        DM::Remote::DisaggReadSnapshotPtr & disagg_snap,
         size_t max_block_size);
 
-    DM::Remote::DisaggPhysicalTableReadSnapshotPtr buildLocalExecForPhysicalTable(
+    void buildLocalExecForPhysicalTable(
         PipelineExecutorContext & exec_context,
-        PipelineExecGroupBuilder & group_builder,
+        TableScanPipelineExecGroupBuilder & group_builder,
         const TableID & table_id,
         const SelectQueryInfo & query_info,
         size_t max_block_size);

dbms/src/Flash/Coprocessor/TableScanPipelineExecBuilder.h

@@ -0,0 +1,48 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// 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.
+
+#pragma once
+
+#include <Flash/Pipeline/Exec/PipelineExec.h>
+#include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
+#include <Storages/DeltaMerge/ReadThread/SharedBlockQueue.h>
+#include <Storages/DeltaMerge/Remote/DisaggSnapshot.h>
+
+namespace DB
+{
+
+class TableScanPipelineExecGroupBuilder : public PipelineExecGroupBuilder
+{
+public:
+    TableScanPipelineExecGroupBuilder(DM::SharedBlockQueuePtr read_queue_, DM::Remote::DisaggReadSnapshotPtr disagg_snap_)
+        : read_queue(std::move(read_queue_))
+        , disagg_snapshot(std::move(disagg_snap_))
+    {}
+
+    void addPhysicalTableTask(TableID table_id, DM::Remote::DisaggPhysicalTableReadSnapshotPtr table_snap)
+    {
+        if (table_snap != nullptr)
+        {
+            RUNTIME_CHECK(disagg_snapshot != nullptr);
+            disagg_snapshot->addTask(table_id, std::move(table_snap));
+        }
+        read_queue->addTableTask(table_id);
+    }
+
+private:
+    DM::SharedBlockQueuePtr read_queue;
+    DM::Remote::DisaggReadSnapshotPtr disagg_snapshot;
+};
+
+} // namespace DB

dbms/src/Operators/ConcatSourceOp.h

@@ -390,10 +390,7 @@ class ConcatBuilderPool
                 }
                 else
                 {
-                    LOG_INFO(
-                        Logger::get(req_id),
-                        "Using ConcatSourceOp to concat {} source ops",
-                        builders.size());
+                    LOG_INFO(Logger::get(req_id), "Using ConcatSourceOp to concat {} source ops", builders.size());
                     result_builder.addConcurrency(std::make_unique<ConcatSourceOp>(exec_context, req_id, builders));
                 }
             }