From 882b2a33488c4dbabbe96fff40c01a065fe0a860 Mon Sep 17 00:00:00 2001 From: Amos Bird Date: Sun, 13 Sep 2020 01:00:04 +0800 Subject: [PATCH] CTE --- src/Interpreters/ApplyWithSubqueryVisitor.cpp | 90 +++++++++++++++++++ src/Interpreters/ApplyWithSubqueryVisitor.h | 30 +++++++ .../ExecuteScalarSubqueriesVisitor.cpp | 5 ++ src/Interpreters/InterpreterSelectQuery.cpp | 3 + src/Interpreters/ya.make | 1 + src/Parsers/ASTWithElement.cpp | 21 +++++ src/Parsers/ASTWithElement.h | 25 ++++++ src/Parsers/ParserSelectQuery.cpp | 6 +- src/Parsers/ParserWithElement.cpp | 39 ++++++++ src/Parsers/ParserWithElement.h | 18 ++++ src/Parsers/ya.make | 2 + ...495_subqueries_in_with_statement.reference | 14 +++ .../01495_subqueries_in_with_statement.sql | 13 +++ 13 files changed, 266 insertions(+), 1 deletion(-) create mode 100644 src/Interpreters/ApplyWithSubqueryVisitor.cpp create mode 100644 src/Interpreters/ApplyWithSubqueryVisitor.h create mode 100644 src/Parsers/ASTWithElement.cpp create mode 100644 src/Parsers/ASTWithElement.h create mode 100644 src/Parsers/ParserWithElement.cpp create mode 100644 src/Parsers/ParserWithElement.h create mode 100644 tests/queries/0_stateless/01495_subqueries_in_with_statement.reference create mode 100644 tests/queries/0_stateless/01495_subqueries_in_with_statement.sql diff --git a/src/Interpreters/ApplyWithSubqueryVisitor.cpp b/src/Interpreters/ApplyWithSubqueryVisitor.cpp new file mode 100644 index 00000000000..e03682dafb3 --- /dev/null +++ b/src/Interpreters/ApplyWithSubqueryVisitor.cpp @@ -0,0 +1,90 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +namespace DB +{ +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) + { + for (auto & child : with->children) + visit(child, data); + for (auto & child : with->children) + { + if (auto * ast_with_elem = child->as()) + { + if (!new_data) + new_data = data; + new_data->subqueries[ast_with_elem->name] = ast_with_elem->subquery; + } + } + } + + for (auto & child : node_select->children) + { + if (child != 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); +} + +void ApplyWithSubqueryVisitor::visit(ASTTableExpression & table, const Data & data) +{ + if (table.database_and_table_name) + { + auto table_id = IdentifierSemantic::extractDatabaseAndTable(table.database_and_table_name->as()); + if (table_id.database_name.empty()) + { + auto subquery_it = data.subqueries.find(table_id.table_name); + if (subquery_it != data.subqueries.end()) + { + table.children.clear(); + table.database_and_table_name.reset(); + table.subquery = subquery_it->second->clone(); + dynamic_cast(*table.subquery).alias = table_id.table_name; + table.children.emplace_back(table.subquery); + } + } + } +} + +void ApplyWithSubqueryVisitor::visit(ASTFunction & func, const Data & data) +{ + if (checkFunctionIsInOrGlobalInOperator(func)) + { + auto & ast = func.arguments->children.at(1); + if (const auto * ident = ast->as()) + { + auto table_id = IdentifierSemantic::extractDatabaseAndTable(*ident); + if (table_id.database_name.empty()) + { + auto subquery_it = data.subqueries.find(table_id.table_name); + if (subquery_it != data.subqueries.end()) + { + func.arguments->children[1] = subquery_it->second->clone(); + dynamic_cast(*func.arguments->children[1]).alias = table_id.table_name; + } + } + } + } +} + +} diff --git a/src/Interpreters/ApplyWithSubqueryVisitor.h b/src/Interpreters/ApplyWithSubqueryVisitor.h new file mode 100644 index 00000000000..2aecd6aee01 --- /dev/null +++ b/src/Interpreters/ApplyWithSubqueryVisitor.h @@ -0,0 +1,30 @@ +#pragma once + +#include + +#include + +namespace DB +{ +// TODO After we support `union_with_global`, this visitor should also be extended to match ASTSelectQueryWithUnion. +class ASTSelectQuery; +class ASTFunction; +struct ASTTableExpression; + +class ApplyWithSubqueryVisitor +{ +public: + struct Data + { + std::map subqueries; + }; + + static void visit(ASTPtr & ast) { visit(ast, {}); } + +private: + static void visit(ASTPtr & ast, const Data & data); + static void visit(ASTTableExpression & table, const Data & data); + static void visit(ASTFunction & func, const Data & data); +}; + +} diff --git a/src/Interpreters/ExecuteScalarSubqueriesVisitor.cpp b/src/Interpreters/ExecuteScalarSubqueriesVisitor.cpp index ee29d301c6b..f7a1fc83182 100644 --- a/src/Interpreters/ExecuteScalarSubqueriesVisitor.cpp +++ b/src/Interpreters/ExecuteScalarSubqueriesVisitor.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include @@ -40,6 +41,10 @@ bool ExecuteScalarSubqueriesMatcher::needChildVisit(ASTPtr & node, const ASTPtr if (node->as()) return false; + /// Do not go to subqueries defined in with statement + if (node->as()) + return false; + if (node->as()) { /// Do not go to FROM, JOIN, UNION. diff --git a/src/Interpreters/InterpreterSelectQuery.cpp b/src/Interpreters/InterpreterSelectQuery.cpp index dbf6d5ae8d3..603476ac1ba 100644 --- a/src/Interpreters/InterpreterSelectQuery.cpp +++ b/src/Interpreters/InterpreterSelectQuery.cpp @@ -14,6 +14,7 @@ #include +#include #include #include #include @@ -249,6 +250,8 @@ InterpreterSelectQuery::InterpreterSelectQuery( source_header = input_pipe->getHeader(); } + ApplyWithSubqueryVisitor().visit(query_ptr); + JoinedTables joined_tables(getSubqueryContext(*context), getSelectQuery()); if (!has_input && !storage) diff --git a/src/Interpreters/ya.make b/src/Interpreters/ya.make index 4eacdab1dcf..5f520505a8a 100644 --- a/src/Interpreters/ya.make +++ b/src/Interpreters/ya.make @@ -23,6 +23,7 @@ SRCS( addTypeConversionToAST.cpp AggregateDescription.cpp Aggregator.cpp + ApplyWithSubqueryVisitor.cpp ArithmeticOperationsInAgrFuncOptimize.cpp ArrayJoinAction.cpp AsynchronousMetricLog.cpp diff --git a/src/Parsers/ASTWithElement.cpp b/src/Parsers/ASTWithElement.cpp new file mode 100644 index 00000000000..e8dd4ff0498 --- /dev/null +++ b/src/Parsers/ASTWithElement.cpp @@ -0,0 +1,21 @@ +#include + +namespace DB +{ + +ASTPtr ASTWithElement::clone() const +{ + const auto res = std::make_shared(*this); + res->name = name; + res->subquery = subquery->clone(); + res->children.emplace_back(res->subquery); + return res; +} + +void ASTWithElement::formatImpl(const FormatSettings & settings, FormatState & state, FormatStateStacked frame) const +{ + settings.writeIdentifier(name); + settings.ostr << (settings.hilite ? hilite_keyword : "") << " AS " << (settings.hilite ? hilite_none : ""); + subquery->formatImpl(settings, state, frame); +} +} diff --git a/src/Parsers/ASTWithElement.h b/src/Parsers/ASTWithElement.h new file mode 100644 index 00000000000..97c68579fa1 --- /dev/null +++ b/src/Parsers/ASTWithElement.h @@ -0,0 +1,25 @@ +#pragma once + +#include + + +namespace DB +{ +/** subquery in with statement + */ +class ASTWithElement : public IAST +{ +public: + String name; + ASTPtr subquery; + + /** Get the text that identifies this element. */ + String getID(char) const override { return "WithElement"; } + + ASTPtr clone() const override; + +protected: + void formatImpl(const FormatSettings & settings, FormatState & state, FormatStateStacked frame) const override; +}; + +} diff --git a/src/Parsers/ParserSelectQuery.cpp b/src/Parsers/ParserSelectQuery.cpp index d2d7bbf9f21..9f2df82b4b4 100644 --- a/src/Parsers/ParserSelectQuery.cpp +++ b/src/Parsers/ParserSelectQuery.cpp @@ -8,6 +8,7 @@ #include #include #include +#include namespace DB @@ -74,7 +75,10 @@ bool ParserSelectQuery::parseImpl(Pos & pos, ASTPtr & node, Expected & expected) { if (s_with.ignore(pos, expected)) { - if (!exp_list_for_with_clause.parse(pos, with_expression_list, expected)) + if (!ParserList(std::make_unique(), std::make_unique(TokenType::Comma)) + .parse(pos, with_expression_list, expected)) + return false; + if (with_expression_list->children.empty()) return false; } } diff --git a/src/Parsers/ParserWithElement.cpp b/src/Parsers/ParserWithElement.cpp new file mode 100644 index 00000000000..048e891f0df --- /dev/null +++ b/src/Parsers/ParserWithElement.cpp @@ -0,0 +1,39 @@ +#include +#include +#include +#include +#include +#include +#include +#include + + +namespace DB +{ +bool ParserWithElement::parseImpl(Pos & pos, ASTPtr & node, Expected & expected) +{ + ParserIdentifier s_ident; + ParserKeyword s_as("AS"); + ParserSubquery s_subquery; + + auto old_pos = pos; + if (ASTPtr name, subquery; + s_ident.parse(pos, name, expected) && s_as.ignore(pos, expected) && s_subquery.parse(pos, subquery, expected)) + { + auto with_element = std::make_shared(); + tryGetIdentifierNameInto(name, with_element->name); + with_element->subquery = subquery; + node = with_element; + } + else + { + pos = old_pos; + ParserExpressionWithOptionalAlias s_expr(false); + if (!s_expr.parse(pos, node, expected)) + return false; + } + return true; +} + + +} diff --git a/src/Parsers/ParserWithElement.h b/src/Parsers/ParserWithElement.h new file mode 100644 index 00000000000..75ad11f5deb --- /dev/null +++ b/src/Parsers/ParserWithElement.h @@ -0,0 +1,18 @@ +#pragma once + +#include + + +namespace DB +{ +/** WITH (scalar query) AS identifier + * or WITH identifier AS (subquery) + */ +class ParserWithElement : public IParserBase +{ +protected: + const char * getName() const override { return "WITH element"; } + bool parseImpl(Pos & pos, ASTPtr & node, Expected & expected) override; +}; + +} diff --git a/src/Parsers/ya.make b/src/Parsers/ya.make index fabf2bbb8fd..0a0c301b722 100644 --- a/src/Parsers/ya.make +++ b/src/Parsers/ya.make @@ -61,6 +61,7 @@ SRCS( ASTTTLElement.cpp ASTUserNameWithHost.cpp ASTWithAlias.cpp + ASTWithElement.cpp CommonParsers.cpp ExpressionElementParsers.cpp ExpressionListParsers.cpp @@ -133,6 +134,7 @@ SRCS( ParserUseQuery.cpp ParserUserNameWithHost.cpp ParserWatchQuery.cpp + ParserWithElement.cpp parseUserName.cpp queryToString.cpp QueryWithOutputSettingsPushDownVisitor.cpp diff --git a/tests/queries/0_stateless/01495_subqueries_in_with_statement.reference b/tests/queries/0_stateless/01495_subqueries_in_with_statement.reference new file mode 100644 index 00000000000..8e851cd3ba1 --- /dev/null +++ b/tests/queries/0_stateless/01495_subqueries_in_with_statement.reference @@ -0,0 +1,14 @@ +0 +1 +2 +3 +4 +2 3 +4 5 +2 3 +4 5 +1 1 2 +3 3 4 +4 5 +4 5 +4 5 diff --git a/tests/queries/0_stateless/01495_subqueries_in_with_statement.sql b/tests/queries/0_stateless/01495_subqueries_in_with_statement.sql new file mode 100644 index 00000000000..9ec921a9d4c --- /dev/null +++ b/tests/queries/0_stateless/01495_subqueries_in_with_statement.sql @@ -0,0 +1,13 @@ +DROP TABLE IF EXISTS test1; + +CREATE TABLE test1(i int, j int) ENGINE Log; + +INSERT INTO test1 VALUES (1, 2), (3, 4); + +WITH test1 AS (SELECT * FROM numbers(5)) SELECT * FROM test1; +WITH test1 AS (SELECT i + 1, j + 1 FROM test1) SELECT * FROM test1; +WITH test1 AS (SELECT i + 1, j + 1 FROM test1) SELECT * FROM (SELECT * FROM test1); +SELECT * FROM (WITH test1 AS (SELECT toInt32(*) i FROM numbers(5)) SELECT * FROM test1) l ANY INNER JOIN test1 r on (l.i == r.i); +WITH test1 AS (SELECT i + 1, j + 1 FROM test1) SELECT toInt64(4) i, toInt64(5) j FROM numbers(3) WHERE (i, j) IN test1; + +DROP TABLE IF EXISTS test1;