From 25da9dcef713f9a7860eb00c4aa23ea9a3263b9c Mon Sep 17 00:00:00 2001
From: Maksim Kita <kitaetoya@gmail.com>
Date: Tue, 24 Jan 2023 19:15:41 +0100
Subject: [PATCH] StorageDistributed Planner initialization fix

---
 src/Planner/Planner.cpp             | 11 +++++++++++
 src/Planner/Planner.h               |  6 ++++++
 src/Storages/StorageDistributed.cpp | 12 ++++++++++--
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/Planner/Planner.cpp b/src/Planner/Planner.cpp
index 2a1bd93391f..6bf51ee9cf3 100644
--- a/src/Planner/Planner.cpp
+++ b/src/Planner/Planner.cpp
@@ -959,6 +959,17 @@ Planner::Planner(const QueryTreeNodePtr & query_tree_,
 {
 }
 
+Planner::Planner(const QueryTreeNodePtr & query_tree_,
+    const SelectQueryOptions & select_query_options_,
+    PlannerContextPtr planner_context_,
+    PlannerConfiguration planner_configuration_)
+    : query_tree(query_tree_)
+    , select_query_options(select_query_options_)
+    , planner_context(std::move(planner_context_))
+    , planner_configuration(std::move(planner_configuration_))
+{
+}
+
 void Planner::buildQueryPlanIfNeeded()
 {
     if (query_plan.isInitialized())
diff --git a/src/Planner/Planner.h b/src/Planner/Planner.h
index 6e225bbf905..a87880eab6b 100644
--- a/src/Planner/Planner.h
+++ b/src/Planner/Planner.h
@@ -35,6 +35,12 @@ public:
         GlobalPlannerContextPtr global_planner_context_,
         PlannerConfiguration planner_configuration_ = {});
 
+    /// Initialize planner with query tree after query analysis phase and planner context
+    Planner(const QueryTreeNodePtr & query_tree_,
+        const SelectQueryOptions & select_query_options_,
+        PlannerContextPtr planner_context_,
+        PlannerConfiguration planner_configuration_ = {});
+
     const QueryPlan & getQueryPlan() const
     {
         return query_plan;
diff --git a/src/Storages/StorageDistributed.cpp b/src/Storages/StorageDistributed.cpp
index 954dddd07f8..e7468ae9d1e 100644
--- a/src/Storages/StorageDistributed.cpp
+++ b/src/Storages/StorageDistributed.cpp
@@ -663,8 +663,12 @@ void StorageDistributed::read(
         StorageID remote_storage_id{remote_database, remote_table};
         auto query_tree_with_replaced_distributed_table = buildQueryTreeDistributedTableReplacedWithLocalTable(query_info, remote_storage_id);
         query_ast = queryNodeToSelectQuery(query_tree_with_replaced_distributed_table);
-        Planner planner(query_tree_with_replaced_distributed_table, SelectQueryOptions(processed_stage), PlannerConfiguration{.only_analyze = true});
+
+        Planner planner(query_tree_with_replaced_distributed_table,
+            SelectQueryOptions(processed_stage),
+            PlannerConfiguration{.only_analyze = true});
         planner.buildQueryPlanIfNeeded();
+
         header = planner.getQueryPlan().getCurrentDataStream().header;
     }
     else
@@ -716,8 +720,12 @@ void StorageDistributed::read(
 
     if (local_context->getSettingsRef().allow_experimental_analyzer)
     {
-        Planner planner(query_info.query_tree, SelectQueryOptions(processed_stage), PlannerConfiguration{.only_analyze = true});
+        Planner planner(query_info.query_tree,
+            SelectQueryOptions(processed_stage),
+            query_info.planner_context,
+            PlannerConfiguration{.only_analyze = true});
         planner.buildQueryPlanIfNeeded();
+
         auto expected_header = planner.getQueryPlan().getCurrentDataStream().header;
 
         auto rename_actions_dag = ActionsDAG::makeConvertingActions(