Analyzer: Resolve GROUPING function on shards

src/Analyzer/Passes/GroupingFunctionsResolvePass.cpp

@@ -227,19 +227,20 @@ void resolveGroupingFunctions(QueryTreeNodePtr & query_node, ContextPtr context)
     visitor.visit(query_node);
 }
 
-class GroupingFunctionsResolveVisitor : public InDepthQueryTreeVisitor<GroupingFunctionsResolveVisitor>
+class GroupingFunctionsResolveVisitor : public InDepthQueryTreeVisitorWithContext<GroupingFunctionsResolveVisitor>
 {
+    using Base = InDepthQueryTreeVisitorWithContext<GroupingFunctionsResolveVisitor>;
 public:
     explicit GroupingFunctionsResolveVisitor(ContextPtr context_)
-        : context(std::move(context_))
+        : Base(std::move(context_))
     {}
 
-    void visitImpl(QueryTreeNodePtr & node)
+    void enterImpl(QueryTreeNodePtr & node)
     {
         if (node->getNodeType() != QueryTreeNodeType::QUERY)
             return;
 
-        resolveGroupingFunctions(node, context);
+        resolveGroupingFunctions(node, getContext());
     }
 
 private:

src/Analyzer/QueryTreePassManager.cpp

@@ -190,6 +190,12 @@ void QueryTreePassManager::run(QueryTreeNodePtr query_tree_node)
     }
 }
 
+void QueryTreePassManager::runOnlyResolve(QueryTreeNodePtr query_tree_node)
+{
+    // Run only QueryAnalysisPass and GroupingFunctionsResolvePass passes.
+    run(query_tree_node, 2);
+}
+
 void QueryTreePassManager::run(QueryTreeNodePtr query_tree_node, size_t up_to_pass_index)
 {
     size_t passes_size = passes.size();
@@ -243,6 +249,8 @@ void QueryTreePassManager::dump(WriteBuffer & buffer, size_t up_to_pass_index)
 void addQueryTreePasses(QueryTreePassManager & manager)
 {
     manager.addPass(std::make_unique<QueryAnalysisPass>());
+    manager.addPass(std::make_unique<GroupingFunctionsResolvePass>());
+
     manager.addPass(std::make_unique<RemoveUnusedProjectionColumnsPass>());
     manager.addPass(std::make_unique<FunctionToSubcolumnsPass>());
 
@@ -278,7 +286,6 @@ void addQueryTreePasses(QueryTreePassManager & manager)
 
     manager.addPass(std::make_unique<LogicalExpressionOptimizerPass>());
 
-    manager.addPass(std::make_unique<GroupingFunctionsResolvePass>());
     manager.addPass(std::make_unique<AutoFinalOnQueryPass>());
     manager.addPass(std::make_unique<CrossToInnerJoinPass>());
     manager.addPass(std::make_unique<ShardNumColumnToFunctionPass>());

src/Analyzer/QueryTreePassManager.h

@@ -27,6 +27,9 @@ public:
     /// Run query tree passes on query tree
     void run(QueryTreeNodePtr query_tree_node);
 
+    /// Run only query tree passes responsible to name resolution.
+    void runOnlyResolve(QueryTreeNodePtr query_tree_node);
+
     /** Run query tree passes on query tree up to up_to_pass_index.
       * Throws exception if up_to_pass_index is greater than passes size.
       */

src/Interpreters/InterpreterSelectQueryAnalyzer.cpp

@@ -143,7 +143,7 @@ QueryTreeNodePtr buildQueryTreeAndRunPasses(const ASTPtr & query,
     /// because it can lead to a changed header.
     if (select_query_options.ignore_ast_optimizations
         || context->getClientInfo().query_kind == ClientInfo::QueryKind::SECONDARY_QUERY)
-        query_tree_pass_manager.run(query_tree, 1 /*up_to_pass_index*/);
+        query_tree_pass_manager.runOnlyResolve(query_tree);
     else
         query_tree_pass_manager.run(query_tree);