From 29346f607398dba5cb7796d53dc525fd937af104 Mon Sep 17 00:00:00 2001
From: Igor Nikonov <igor@clickhouse.com>
Date: Tue, 28 May 2024 12:16:18 +0000
Subject: [PATCH] Init coordinator separately

---
 src/Planner/PlannerJoinTree.cpp               |  9 +++++---
 .../QueryPlan/ParallelReplicasLocalPlan.cpp   | 23 ++++++++++++++++++-
 src/Processors/QueryPlan/ReadFromMergeTree.h  |  6 +++++
 .../ParallelReplicasReadingCoordinator.h      |  2 +-
 4 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/Planner/PlannerJoinTree.cpp b/src/Planner/PlannerJoinTree.cpp
index 275461fa109..01455d4b955 100644
--- a/src/Planner/PlannerJoinTree.cpp
+++ b/src/Planner/PlannerJoinTree.cpp
@@ -892,16 +892,19 @@ JoinTreeQueryPlan buildQueryPlanForTableExpression(QueryTreeNodePtr table_expres
 
                         if (!node->children.empty())
                             node = node->children.at(0);
+                        else
+                            node = nullptr;
                     }
 
                     chassert(reading);
 
+                    auto result_ptr = reading->selectRangesToRead();
+                    UInt64 rows_to_read = result_ptr->selected_rows;
+                    reading->setAnalyzedResult(std::move(result_ptr));
+
                     // (2) if it's ReadFromMergeTree - run index analysis and check number of rows to read
                     if (settings.parallel_replicas_min_number_of_rows_per_replica > 0)
                     {
-                        auto result_ptr = reading->selectRangesToRead();
-
-                        UInt64 rows_to_read = result_ptr->selected_rows;
                         if (table_expression_query_info.limit > 0 && table_expression_query_info.limit < rows_to_read)
                             rows_to_read = table_expression_query_info.limit;
 
diff --git a/src/Processors/QueryPlan/ParallelReplicasLocalPlan.cpp b/src/Processors/QueryPlan/ParallelReplicasLocalPlan.cpp
index 2bfd8965269..1c5e0861530 100644
--- a/src/Processors/QueryPlan/ParallelReplicasLocalPlan.cpp
+++ b/src/Processors/QueryPlan/ParallelReplicasLocalPlan.cpp
@@ -268,10 +268,32 @@ std::unique_ptr<QueryPlan> createLocalPlanForParallelReplicas(
 
         if (!node->children.empty())
             node = node->children.at(0);
+        else
+            node = nullptr;
     }
 
     chassert(reading);
 
+    const auto * analyzed_merge_tree = typeid_cast<const ReadFromMergeTree *>(read_from_merge_tree.get());
+    if (!analyzed_merge_tree->hasAnalyzedResult())
+        analyzed_merge_tree->selectRangesToRead();
+
+    switch (analyzed_merge_tree->getReadType())
+    {
+    case ReadFromMergeTree::ReadType::Default:
+        coordinator->initialize(CoordinationMode::Default);
+        break;
+    case ReadFromMergeTree::ReadType::InOrder:
+        coordinator->initialize(CoordinationMode::WithOrder);
+        break;
+    case ReadFromMergeTree::ReadType::InReverseOrder:
+        coordinator->initialize(CoordinationMode::ReverseOrder);
+        break;
+    case ReadFromMergeTree::ReadType::ParallelReplicas:
+        chassert(false);
+        UNREACHABLE();
+    }
+
     MergeTreeAllRangesCallback all_ranges_cb = [coordinator](InitialAllRangesAnnouncement announcement)
     {
         chassert(coordinator);
@@ -281,7 +303,6 @@ std::unique_ptr<QueryPlan> createLocalPlanForParallelReplicas(
     MergeTreeReadTaskCallback read_task_cb = [coordinator](ParallelReadRequest req) -> std::optional<ParallelReadResponse>
     { return coordinator->handleRequest(std::move(req)); };
 
-    const auto * analyzed_merge_tree = typeid_cast<const ReadFromMergeTree *>(read_from_merge_tree.get());
     auto read_from_merge_tree_parallel_replicas = reading->createLocalParallelReplicasReadingStep(analyzed_merge_tree, true, all_ranges_cb, read_task_cb);
     node->step = std::move(read_from_merge_tree_parallel_replicas);
 
diff --git a/src/Processors/QueryPlan/ReadFromMergeTree.h b/src/Processors/QueryPlan/ReadFromMergeTree.h
index bba5293e863..adc3818c3ab 100644
--- a/src/Processors/QueryPlan/ReadFromMergeTree.h
+++ b/src/Processors/QueryPlan/ReadFromMergeTree.h
@@ -203,6 +203,12 @@ public:
 
     void applyFilters(ActionDAGNodes added_filter_nodes) override;
 
+    ReadType getReadType() const
+    {
+        chassert(analyzed_result_ptr);
+        return analyzed_result_ptr->read_type;
+    }
+
 private:
     static AnalysisResultPtr selectRangesToReadImpl(
         MergeTreeData::DataPartsVector parts,
diff --git a/src/Storages/MergeTree/ParallelReplicasReadingCoordinator.h b/src/Storages/MergeTree/ParallelReplicasReadingCoordinator.h
index 60343988f03..ff72decbf8d 100644
--- a/src/Storages/MergeTree/ParallelReplicasReadingCoordinator.h
+++ b/src/Storages/MergeTree/ParallelReplicasReadingCoordinator.h
@@ -30,8 +30,8 @@ public:
     /// needed to report total rows to read
     void setProgressCallback(ProgressCallback callback);
 
-private:
     void initialize(CoordinationMode mode);
+private:
 
     std::mutex mutex;
     size_t replicas_count{0};