Global WITH and WITH propagation.

src/Core/Settings.h

@@ -471,6 +471,7 @@ class IColumn;
     M(Bool, output_format_write_statistics, true, "Write statistics about read rows, bytes, time elapsed in suitable output formats.", 0) \
     M(Bool, allow_non_metadata_alters, true, "Allow to execute alters which affects not only tables metadata, but also data on disk", 0) \
     M(Bool, output_format_pretty_row_numbers, false, "Add row numbers before each row for pretty output format", 0) \
+    M(Bool, with_global, false, "Propagate WITH statements to UNION queries and all subqueries", 0) \
 
 #define LIST_OF_SETTINGS(M)    \
     COMMON_SETTINGS(M)         \

src/Interpreters/ApplyWithAliasVisitor.cpp

@@ -0,0 +1,53 @@
+#include <Interpreters/ApplyWithAliasVisitor.h>
+#include <Interpreters/misc.h>
+#include <Parsers/ASTSelectQuery.h>
+#include <Parsers/ASTSelectWithUnionQuery.h>
+
+namespace DB
+{
+void ApplyWithAliasVisitor::visit(ASTPtr & ast, const Data & data)
+{
+    if (auto * node_select = ast->as<ASTSelectQuery>())
+    {
+        std::optional<Data> new_data;
+        if (auto with = node_select->with())
+        {
+            for (auto & child : with->children)
+                visit(child, data);
+
+            std::set<String> current_names;
+            for (auto & child : with->children)
+            {
+                if (auto * ast_with_alias = dynamic_cast<ASTWithAlias *>(child.get()))
+                {
+                    if (!new_data)
+                        new_data = data;
+                    new_data->exprs[ast_with_alias->alias] = child;
+                    current_names.insert(ast_with_alias->alias);
+                }
+            }
+            for (const auto & with_alias : data.exprs)
+            {
+                if (!current_names.count(with_alias.first))
+                    with->children.push_back(with_alias.second->clone());
+            }
+        }
+        else if (!data.exprs.empty())
+        {
+            auto with_expression_list = std::make_shared<ASTExpressionList>();
+            for (const auto & with_alias : data.exprs)
+                with_expression_list->children.push_back(with_alias.second->clone());
+            node_select->setExpression(ASTSelectQuery::Expression::WITH, std::move(with_expression_list));
+        }
+        for (auto & child : node_select->children)
+        {
+            if (child != node_select->with())
+                visit(child, new_data ? *new_data : data);
+        }
+    }
+    else
+        for (auto & child : ast->children)
+            visit(child, data);
+}
+
+}

src/Interpreters/ApplyWithAliasVisitor.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <map>
+
+#include <Parsers/IAST.h>
+
+namespace DB
+{
+/// Propagate every WITH alias expression to its descendent subqueries, with correct scoping visibility.
+class ApplyWithAliasVisitor
+{
+public:
+    struct Data
+    {
+        std::map<String, ASTPtr> exprs;
+    };
+
+    static void visit(ASTPtr & ast) { visit(ast, {}); }
+
+private:
+    static void visit(ASTPtr & ast, const Data & data);
+};
+
+}

src/Interpreters/ApplyWithGlobalVisitor.cpp

@@ -0,0 +1,53 @@
+#include <Interpreters/ApplyWithGlobalVisitor.h>
+#include <Parsers/ASTSelectQuery.h>
+#include <Parsers/ASTSelectWithUnionQuery.h>
+#include <Parsers/ASTWithAlias.h>
+
+#include <map>
+
+namespace DB
+{
+void ApplyWithGlobalVisitor::visit(ASTPtr & ast)
+{
+    if (ASTSelectWithUnionQuery * node_union = ast->as<ASTSelectWithUnionQuery>())
+    {
+        auto & first_select = node_union->list_of_selects->children[0]->as<ASTSelectQuery &>();
+        ASTPtr with_expression_list = first_select.with();
+
+        if (with_expression_list)
+        {
+            std::map<String, ASTPtr> exprs;
+            for (auto & child : with_expression_list->children)
+            {
+                if (auto * ast_with_alias = dynamic_cast<ASTWithAlias *>(child.get()))
+                    exprs[ast_with_alias->alias] = child;
+            }
+            for (auto it = node_union->list_of_selects->children.begin() + 1; it != node_union->list_of_selects->children.end(); ++it)
+            {
+                auto & select = (*it)->as<ASTSelectQuery &>();
+                auto with = select.with();
+                if (with)
+                {
+                    std::set<String> current_names;
+                    for (auto & child : with->children)
+                    {
+                        if (auto * ast_with_alias = dynamic_cast<ASTWithAlias *>(child.get()))
+                            current_names.insert(ast_with_alias->alias);
+                    }
+                    for (auto & with_alias : exprs)
+                    {
+                        if (!current_names.count(with_alias.first))
+                            with->children.push_back(with_alias.second->clone());
+                    }
+                }
+                else
+                    select.setExpression(ASTSelectQuery::Expression::WITH, with_expression_list->clone());
+            }
+        }
+    }
+
+    for (auto & child : ast->children)
+        visit(child);
+}
+
+}

src/Interpreters/ApplyWithGlobalVisitor.h

@@ -0,0 +1,14 @@
+#pragma once
+
+#include <Parsers/IAST.h>
+
+namespace DB
+{
+/// Pull out the WITH statement from the first child of ASTSelectWithUnion query if any.
+class ApplyWithGlobalVisitor
+{
+public:
+    static void visit(ASTPtr & ast);
+};
+
+}

src/Interpreters/ApplyWithSubqueryVisitor.cpp

@@ -13,9 +13,8 @@ void ApplyWithSubqueryVisitor::visit(ASTPtr & ast, const Data & data)
 {
     if (auto * node_select = ast->as<ASTSelectQuery>())
     {
-        auto with = node_select->with();
         std::optional<Data> new_data;
-        if (with)
+        if (auto with = node_select->with())
         {
             for (auto & child : with->children)
                 visit(child, data);
@@ -32,12 +31,12 @@ void ApplyWithSubqueryVisitor::visit(ASTPtr & ast, const Data & data)
 
         for (auto & child : node_select->children)
         {
-            if (child != with)
+            if (child != node_select->with())
                 visit(child, new_data ? *new_data : data);
         }
-        return;
     }
-
+    else
+    {
         for (auto & child : ast->children)
             visit(child, data);
         if (auto * node_func = ast->as<ASTFunction>())
@@ -45,6 +44,7 @@ void ApplyWithSubqueryVisitor::visit(ASTPtr & ast, const Data & data)
         else if (auto * node_table = ast->as<ASTTableExpression>())
             visit(*node_table, data);
     }
+}
 
 void ApplyWithSubqueryVisitor::visit(ASTTableExpression & table, const Data & data)
 {

src/Interpreters/ApplyWithSubqueryVisitor.h

@@ -6,8 +6,6 @@
 
 namespace DB
 {
-// TODO After we support `union_with_global`, this visitor should also be extended to match ASTSelectQueryWithUnion.
-class ASTSelectQuery;
 class ASTFunction;
 struct ASTTableExpression;
 

src/Interpreters/InterpreterSelectQuery.cpp

@@ -14,6 +14,7 @@
 
 #include <Access/AccessFlags.h>
 
+#include <Interpreters/ApplyWithAliasVisitor.h>
 #include <Interpreters/ApplyWithSubqueryVisitor.h>
 #include <Interpreters/InterpreterSelectQuery.h>
 #include <Interpreters/InterpreterSelectWithUnionQuery.h>
@@ -245,6 +246,8 @@ InterpreterSelectQuery::InterpreterSelectQuery(
         source_header = input_pipe->getHeader();
     }
 
+    if (context->getSettingsRef().with_global)
+        ApplyWithAliasVisitor().visit(query_ptr);
     ApplyWithSubqueryVisitor().visit(query_ptr);
 
     JoinedTables joined_tables(getSubqueryContext(*context), getSelectQuery());

src/Interpreters/InterpreterSelectWithUnionQuery.cpp

@@ -1,6 +1,7 @@
 #include <Interpreters/InterpreterSelectWithUnionQuery.h>
 #include <Interpreters/InterpreterSelectQuery.h>
 #include <Interpreters/Context.h>
+#include <Interpreters/ApplyWithGlobalVisitor.h>
 #include <Parsers/ASTSelectWithUnionQuery.h>
 #include <Parsers/ASTSelectQuery.h>
 #include <Columns/getLeastSuperColumn.h>
@@ -38,6 +39,9 @@ InterpreterSelectWithUnionQuery::InterpreterSelectWithUnionQuery(
     if (!num_selects)
         throw Exception("Logical error: no children in ASTSelectWithUnionQuery", ErrorCodes::LOGICAL_ERROR);
 
+    if (context->getSettingsRef().with_global)
+        ApplyWithGlobalVisitor().visit(query_ptr);
+
     /// Initialize interpreters for each SELECT query.
     /// Note that we pass 'required_result_column_names' to first SELECT.
     /// And for the rest, we pass names at the corresponding positions of 'required_result_column_names' in the result of first SELECT,

src/Interpreters/ya.make

@@ -23,6 +23,8 @@ SRCS(
     addTypeConversionToAST.cpp
     AggregateDescription.cpp
     Aggregator.cpp
+    ApplyWithAliasVisitor.cpp
+    ApplyWithGlobalVisitor.cpp
     ApplyWithSubqueryVisitor.cpp
     ArithmeticOperationsInAgrFuncOptimize.cpp
     ArrayJoinAction.cpp

tests/queries/0_stateless/01515_with_global_and_with_propagation.reference

@@ -0,0 +1,7 @@
+1
+1
+2	1
+1
+1
+1
+2

tests/queries/0_stateless/01515_with_global_and_with_propagation.sql

@@ -0,0 +1,6 @@
+SET with_global = true;
+WITH 1 AS x SELECT x;
+WITH 1 AS x SELECT * FROM (SELECT x);
+WITH 1 AS x SELECT *, x FROM (WITH 2 AS x SELECT x AS y);
+WITH 1 AS x SELECT x UNION ALL SELECT x;
+WITH 1 AS x SELECT x UNION ALL WITH 2 AS x SELECT x;