Merge pull request #58833 from ClickHouse/fix-analyzer-remote-aggregation

Analyzer: Resolve GROUPING function on shards

src/Analyzer/Passes/FunctionToSubcolumnsPass.cpp

@@ -49,6 +49,9 @@ public:
         if (!first_argument_column_node)
             return;
 
+        if (first_argument_column_node->getColumnName() == "__grouping_set")
+            return;
+
         auto column_source = first_argument_column_node->getColumnSource();
         auto * table_node = column_source->as<TableNode>();
 

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/Passes/OptimizeDateOrDateTimeConverterWithPreimagePass.cpp

@@ -91,6 +91,9 @@ public:
         const auto * column_id = func_node->getArguments().getNodes()[0]->as<ColumnNode>();
         if (!column_id) return;
 
+        if (column_id->getColumnName() == "__grouping_set")
+            return;
+
         const auto * column_type = column_id->getColumnType().get();
         if (!isDateOrDate32(column_type) && !isDateTime(column_type) && !isDateTime64(column_type)) return;
 

src/Analyzer/Passes/RemoveUnusedProjectionColumnsPass.cpp

@@ -52,6 +52,9 @@ public:
             return;
 
         auto & column_node = node->as<ColumnNode &>();
+        if (column_node.getColumnName() == "__grouping_set")
+            return;
+
         auto column_source_node = column_node.getColumnSource();
         auto column_source_node_type = column_source_node->getNodeType();
 

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

@@ -144,7 +144,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);