diff --git a/src/Core/Settings.h b/src/Core/Settings.h index 6d3437defb2..23dd951aff8 100644 --- a/src/Core/Settings.h +++ b/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) \ diff --git a/src/Interpreters/ApplyWithAliasVisitor.cpp b/src/Interpreters/ApplyWithAliasVisitor.cpp new file mode 100644 index 00000000000..ab1656a5dde --- /dev/null +++ b/src/Interpreters/ApplyWithAliasVisitor.cpp @@ -0,0 +1,53 @@ +#include +#include +#include +#include + +namespace DB +{ +void ApplyWithAliasVisitor::visit(ASTPtr & ast, const Data & data) +{ + if (auto * node_select = ast->as()) + { + std::optional new_data; + if (auto with = node_select->with()) + { + for (auto & child : with->children) + visit(child, data); + + std::set current_names; + for (auto & child : with->children) + { + if (auto * ast_with_alias = dynamic_cast(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(); + 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); +} + +} diff --git a/src/Interpreters/ApplyWithAliasVisitor.h b/src/Interpreters/ApplyWithAliasVisitor.h new file mode 100644 index 00000000000..8ebe8a66182 --- /dev/null +++ b/src/Interpreters/ApplyWithAliasVisitor.h @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include + +namespace DB +{ +/// Propagate every WITH alias expression to its descendent subqueries, with correct scoping visibility. +class ApplyWithAliasVisitor +{ +public: + struct Data + { + std::map exprs; + }; + + static void visit(ASTPtr & ast) { visit(ast, {}); } + +private: + static void visit(ASTPtr & ast, const Data & data); +}; + +} diff --git a/src/Interpreters/ApplyWithGlobalVisitor.cpp b/src/Interpreters/ApplyWithGlobalVisitor.cpp new file mode 100644 index 00000000000..70256adaaa2 --- /dev/null +++ b/src/Interpreters/ApplyWithGlobalVisitor.cpp @@ -0,0 +1,53 @@ +#include +#include +#include +#include + +#include + +namespace DB +{ +void ApplyWithGlobalVisitor::visit(ASTPtr & ast) +{ + if (ASTSelectWithUnionQuery * node_union = ast->as()) + { + auto & first_select = node_union->list_of_selects->children[0]->as(); + ASTPtr with_expression_list = first_select.with(); + + if (with_expression_list) + { + std::map exprs; + for (auto & child : with_expression_list->children) + { + if (auto * ast_with_alias = dynamic_cast(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(); + auto with = select.with(); + if (with) + { + std::set current_names; + for (auto & child : with->children) + { + if (auto * ast_with_alias = dynamic_cast(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); +} + +} diff --git a/src/Interpreters/ApplyWithGlobalVisitor.h b/src/Interpreters/ApplyWithGlobalVisitor.h new file mode 100644 index 00000000000..c4dd2affd20 --- /dev/null +++ b/src/Interpreters/ApplyWithGlobalVisitor.h @@ -0,0 +1,14 @@ +#pragma once + +#include + +namespace DB +{ +/// Pull out the WITH statement from the first child of ASTSelectWithUnion query if any. +class ApplyWithGlobalVisitor +{ +public: + static void visit(ASTPtr & ast); +}; + +} diff --git a/src/Interpreters/ApplyWithSubqueryVisitor.cpp b/src/Interpreters/ApplyWithSubqueryVisitor.cpp index e03682dafb3..805f425beac 100644 --- a/src/Interpreters/ApplyWithSubqueryVisitor.cpp +++ b/src/Interpreters/ApplyWithSubqueryVisitor.cpp @@ -13,9 +13,8 @@ void ApplyWithSubqueryVisitor::visit(ASTPtr & ast, const Data & data) { if (auto * node_select = ast->as()) { - auto with = node_select->with(); std::optional new_data; - if (with) + if (auto with = node_select->with()) { for (auto & child : with->children) visit(child, data); @@ -32,18 +31,19 @@ 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; } - - for (auto & child : ast->children) - visit(child, data); - if (auto * node_func = ast->as()) - visit(*node_func, data); - else if (auto * node_table = ast->as()) - visit(*node_table, data); + else + { + for (auto & child : ast->children) + visit(child, data); + if (auto * node_func = ast->as()) + visit(*node_func, data); + else if (auto * node_table = ast->as()) + visit(*node_table, data); + } } void ApplyWithSubqueryVisitor::visit(ASTTableExpression & table, const Data & data) diff --git a/src/Interpreters/ApplyWithSubqueryVisitor.h b/src/Interpreters/ApplyWithSubqueryVisitor.h index 2aecd6aee01..ed3f9fbcd3e 100644 --- a/src/Interpreters/ApplyWithSubqueryVisitor.h +++ b/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; diff --git a/src/Interpreters/InterpreterSelectQuery.cpp b/src/Interpreters/InterpreterSelectQuery.cpp index 556070d0360..53bdafabcb9 100644 --- a/src/Interpreters/InterpreterSelectQuery.cpp +++ b/src/Interpreters/InterpreterSelectQuery.cpp @@ -14,6 +14,7 @@ #include +#include #include #include #include @@ -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()); diff --git a/src/Interpreters/InterpreterSelectWithUnionQuery.cpp b/src/Interpreters/InterpreterSelectWithUnionQuery.cpp index ba0ebfaaf27..8a94e83b997 100644 --- a/src/Interpreters/InterpreterSelectWithUnionQuery.cpp +++ b/src/Interpreters/InterpreterSelectWithUnionQuery.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -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, diff --git a/src/Interpreters/ya.make b/src/Interpreters/ya.make index bd16843517f..411100b0b9d 100644 --- a/src/Interpreters/ya.make +++ b/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 diff --git a/tests/queries/0_stateless/01515_with_global_and_with_propagation.reference b/tests/queries/0_stateless/01515_with_global_and_with_propagation.reference new file mode 100644 index 00000000000..311e8f318d9 --- /dev/null +++ b/tests/queries/0_stateless/01515_with_global_and_with_propagation.reference @@ -0,0 +1,7 @@ +1 +1 +2 1 +1 +1 +1 +2 diff --git a/tests/queries/0_stateless/01515_with_global_and_with_propagation.sql b/tests/queries/0_stateless/01515_with_global_and_with_propagation.sql new file mode 100644 index 00000000000..fb6b91dd7b9 --- /dev/null +++ b/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;