Support (EXPLAIN SELECT ...) as a subquery (ported from clickhouse #40630) (#5857) (#826)

* Support (EXPLAIN SELECT ...) as a subquery (ported from clickhouse)

* minor adjustments and changes

* change stateless test file to pass the tests

* adjust format and add extra info

* small adjustments of format

Author: yuanhang zhao

src/Parsers/ASTExplainQuery.h

@@ -30,7 +30,8 @@ class ASTExplainQuery : public ASTQueryWithOutput
     {
         auto res = std::make_shared<ASTExplainQuery>(*this);
         res->children.clear();
-        res->children.push_back(children[0]->clone());
+        if (!children.empty())
+            res->children.push_back(children[0]->clone());
         cloneOutputOptions(*res);
         return res;
     }

src/Parsers/ExpressionElementParsers.cpp

@@ -11,6 +11,9 @@
 #include <Parsers/ASTAssignment.h>
 #include <Parsers/ASTAsterisk.h>
 #include <Parsers/ASTColumnsMatcher.h>
+#include <Parsers/ASTExplainQuery.h>
+#include <Parsers/ASTSelectQuery.h>
+#include <Parsers/ASTTablesInSelectQuery.h>
 #include <Parsers/ASTColumnsTransformers.h>
 #include <Parsers/ASTExpressionList.h>
 #include <Parsers/ASTFunction.h>
@@ -33,6 +36,7 @@
 
 #include <Parsers/ExpressionElementParsers.h>
 #include <Parsers/ParserCreateQuery.h>
+#include <Parsers/ParserExplainQuery.h>
 
 #include <Parsers/queryToString.h>
 
@@ -53,6 +57,52 @@ namespace ErrorCodes
     extern const int LOGICAL_ERROR;
 }
 
+/*
+ * Build an AST with the following structure:
+ *
+ * ```
+ * SelectWithUnionQuery (children 1)
+ *  ExpressionList (children 1)
+ *   SelectQuery (children 2)
+ *    ExpressionList (children 1)
+ *     Asterisk
+ *    TablesInSelectQuery (children 1)
+ *     TablesInSelectQueryElement (children 1)
+ *      TableExpression (children 1)
+ *       Function <...>
+ * ```
+ */
+static ASTPtr buildSelectFromTableFunction(const std::shared_ptr<ASTFunction> & ast_function)
+{
+    auto result_select_query = std::make_shared<ASTSelectWithUnionQuery>();
+
+    {
+        auto select_ast = std::make_shared<ASTSelectQuery>();
+        select_ast->setExpression(ASTSelectQuery::Expression::SELECT, std::make_shared<ASTExpressionList>());
+        select_ast->select()->children.push_back(std::make_shared<ASTAsterisk>());
+
+        auto list_of_selects = std::make_shared<ASTExpressionList>();
+        list_of_selects->children.push_back(select_ast);
+
+        result_select_query->children.push_back(std::move(list_of_selects));
+        result_select_query->list_of_selects = result_select_query->children.back();
+
+        {
+            auto tables = std::make_shared<ASTTablesInSelectQuery>();
+            select_ast->setExpression(ASTSelectQuery::Expression::TABLES, tables);
+            auto tables_elem = std::make_shared<ASTTablesInSelectQueryElement>();
+            auto table_expr = std::make_shared<ASTTableExpression>();
+            tables->children.push_back(tables_elem);
+            tables_elem->table_expression = table_expr;
+            tables_elem->children.push_back(table_expr);
+
+            table_expr->table_function = ast_function;
+            table_expr->children.push_back(table_expr->table_function);
+        }
+    }
+
+    return result_select_query;
+}
 
 bool ParserArray::parseImpl(Pos & pos, ASTPtr & node, Expected & expected, [[ maybe_unused ]] bool hint)
 {
@@ -140,22 +190,35 @@ bool ParserParenthesisExpression::parseImpl(Pos & pos, ASTPtr & node, Expected &
 
 bool ParserSubquery::parseImpl(Pos & pos, ASTPtr & node, Expected & expected, [[ maybe_unused ]] bool hint)
 {
-    ASTPtr select_node;
     ParserSelectWithUnionQuery select;
+    ParserExplainQuery explain;
 
     if (pos->type != TokenType::OpeningRoundBracket)
         return false;
     ++pos;
 
-    if (!select.parse(pos, select_node, expected))
+    ASTPtr result_node = nullptr;
+
+    if (ASTPtr select_node; select.parse(pos, select_node, expected))
+    {
+        result_node = std::move(select_node);
+    }
+    else if (ASTPtr explain_node; explain.parse(pos, explain_node, expected))
+    {
+        /// Replace SELECT * FROM (EXPLAIN SELECT ...) with SELECT * FROM viewExplain(EXPLAIN SELECT ...)
+        result_node = buildSelectFromTableFunction(makeASTFunction("viewExplain", explain_node));
+    }
+    else
+    {
         return false;
+    }
 
     if (pos->type != TokenType::ClosingRoundBracket)
         return false;
     ++pos;
 
     node = std::make_shared<ASTSubquery>();
-    node->children.push_back(select_node);
+    node->children.push_back(result_node);
     return true;
 }
 

src/Parsers/ParserExplainQuery.cpp

@@ -88,10 +88,19 @@ bool ParserExplainQuery::parseImpl(Pos & pos, ASTPtr & node, Expected & expected
     {
         /// Nothing to parse
     }
+    else if (select_only)
+    {
+        if (select_p.parse(pos, query, expected))
+            explain_query->setExplainedQuery(std::move(query));
+        else
+            return false;
+    }
     else if (select_p.parse(pos, query, expected) ||
         create_p.parse(pos, query, expected) ||
         insert_p.parse(pos, query, expected))
+    {
         explain_query->setExplainedQuery(std::move(query));
+    }
     else
         return false;
 

src/TableFunctions/TableFunctionExplain.cpp

@@ -0,0 +1,113 @@
+#include <Interpreters/InterpreterSelectWithUnionQuery.h>
+#include <Parsers/ASTFunction.h>
+#include <Parsers/ASTSelectWithUnionQuery.h>
+#include <Parsers/queryToString.h>
+#include <Storages/StorageValues.h>
+#include <TableFunctions/ITableFunction.h>
+#include <TableFunctions/TableFunctionFactory.h>
+#include <TableFunctions/TableFunctionExplain.h>
+#include <TableFunctions/registerTableFunctions.h>
+#include <Processors/Executors/PullingPipelineExecutor.h>
+
+namespace DB
+{
+namespace ErrorCodes
+{
+    extern const int LOGICAL_ERROR;
+    extern const int BAD_ARGUMENTS;
+}
+
+void TableFunctionExplain::parseArguments(const ASTPtr & ast_function, ContextPtr /*context*/)
+{
+    const auto * function = ast_function->as<ASTFunction>();
+    if (function && function->arguments && function->arguments->children.size() == 1)
+    {
+        const auto & query_arg = function->arguments->children[0];
+
+        if (!query_arg->as<ASTExplainQuery>())
+            throw Exception(ErrorCodes::BAD_ARGUMENTS,
+                "Table function '{}' requires a explain query argument, got '{}'",
+                getName(), queryToString(query_arg));
+
+        query = query_arg;
+    }
+    else
+    {
+        throw Exception(ErrorCodes::BAD_ARGUMENTS,
+            "Table function '{}' cannot be called directly, use `SELECT * FROM (EXPLAIN ...)` syntax", getName());
+    }
+}
+
+ColumnsDescription TableFunctionExplain::getActualTableStructure(ContextPtr context) const
+{
+    Block sample_block = getInterpreter(context).getSampleBlock(query->as<ASTExplainQuery>()->getKind());
+    ColumnsDescription columns_description;
+    for (const auto & column : sample_block.getColumnsWithTypeAndName())
+        columns_description.add(ColumnDescription(column.name, column.type));
+    return columns_description;
+}
+
+static Block executeMonoBlock(QueryPipeline & pipeline)
+{
+    if (!pipeline.pulling())
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Expected pulling pipeline");
+
+    PullingPipelineExecutor pulling_executor(pipeline);
+    std::vector<Block> blocks;
+    while (true)
+    {
+        Block block;
+        if (pulling_executor.pull(block))
+            blocks.push_back(std::move(block));
+        else
+            break;
+    }
+
+    if (blocks.size() == 1)
+        return blocks[0];
+
+    return concatenateBlocks(blocks);
+}
+
+StoragePtr TableFunctionExplain::executeImpl(
+    const ASTPtr & /*ast_function*/, ContextPtr context, const std::string & table_name, ColumnsDescription /*cached_columns*/) const
+{
+    BlockIO blockio = getInterpreter(context).execute();
+    Block block = executeMonoBlock(blockio.pipeline);
+
+    StorageID storage_id(getDatabaseName(), table_name);
+    /// proton: start cannot directly make shared DataTypePtr and StorageValues::create is the only way to do it
+    auto storage = StorageValues::create(storage_id,getActualTableStructure(context),std::move(block));
+    /// proton: end
+    storage->startup();
+    return storage;
+}
+
+InterpreterExplainQuery TableFunctionExplain::getInterpreter(ContextPtr context) const
+{
+    if (!query)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Table function '{}' requires a explain query argument", getName());
+
+    return InterpreterExplainQuery(query, context);
+}
+
+void registerTableFunctionExplain(TableFunctionFactory & factory)
+{
+    factory.registerFunction<TableFunctionExplain>({R"(
+Returns result of EXPLAIN query.
+
+The function should not be called directly but can be invoked via `SELECT * FROM (EXPLAIN <query>)`.
+
+You can use this query to process the result of EXPLAIN further using SQL (e.g., in tests).
+
+Example:
+[example:1]
+
+)",
+{{"1", "SELECT explain FROM (EXPLAIN AST SELECT * FROM system.numbers) WHERE explain LIKE '%Asterisk%'"}}
+});
+
+}
+
+}
+

src/TableFunctions/TableFunctionExplain.h

@@ -0,0 +1,32 @@
+#pragma once
+
+#include <TableFunctions/ITableFunction.h>
+#include <Parsers/ASTExplainQuery.h>
+#include <Interpreters/InterpreterExplainQuery.h>
+#include <base/types.h>
+
+
+namespace DB
+{
+
+class TableFunctionExplain : public ITableFunction
+{
+public:
+    static constexpr auto name = "viewExplain";
+    std::string getName() const override { return name; }
+
+private:
+    StoragePtr executeImpl(const ASTPtr & ast_function, ContextPtr context, const String & table_name, ColumnsDescription cached_columns) const override;
+    const char * getStorageTypeName() const override { return "Explain"; }
+
+    void parseArguments(const ASTPtr & ast_function, ContextPtr context) override;
+    ColumnsDescription getActualTableStructure(ContextPtr context) const override;
+
+    InterpreterExplainQuery getInterpreter(ContextPtr context) const;
+
+    ASTPtr query = nullptr;
+};
+
+
+}
+

src/TableFunctions/registerTableFunctions.cpp

@@ -28,6 +28,8 @@ void registerTableFunctions()
 
     registerTableFunctionDictionary(factory);
 
+    registerTableFunctionExplain(factory);
+
     /// proton: starts
     Streaming::registerTableFunctionHop(factory);
     Streaming::registerTableFunctionTumble(factory);

src/TableFunctions/registerTableFunctions.h

@@ -26,6 +26,8 @@ void registerTableFunctionView(TableFunctionFactory & factory);
 
 void registerTableFunctionDictionary(TableFunctionFactory & factory);
 
+void registerTableFunctionExplain(TableFunctionFactory & factory);
+
 /// proton: starts
 namespace Streaming
 {

tests/queries_ported/0_stateless/02421_explain_subquery.reference

@@ -0,0 +1,8 @@
+true
+true
+true
+true
+true
+true
+true
+true

tests/queries_ported/0_stateless/02421_explain_subquery.sql

@@ -0,0 +1,24 @@
+SELECT count() > 3 FROM (EXPLAIN PIPELINE header = 1 SELECT * FROM system.numbers ORDER BY number DESC) WHERE explain LIKE '%Header: number uint64%';
+SELECT count() > 0 FROM (EXPLAIN PLAN SELECT * FROM system.numbers ORDER BY number DESC) WHERE explain ILIKE '%Sort%';
+SELECT count() > 0 FROM (EXPLAIN SELECT * FROM system.numbers ORDER BY number DESC) WHERE explain ILIKE '%Sort%';
+SELECT count() > 0 FROM (EXPLAIN CURRENT TRANSACTION);
+SELECT count() == 1 FROM (EXPLAIN SYNTAX SELECT number FROM system.numbers ORDER BY number DESC) WHERE explain ILIKE 'SELECT%';
+SELECT trim(explain) == 'Asterisk' FROM (EXPLAIN AST SELECT * FROM system.numbers LIMIT 10) WHERE explain LIKE '%Asterisk%';
+
+SELECT * FROM (
+    EXPLAIN AST SELECT * FROM (
+        EXPLAIN PLAN SELECT * FROM (
+            EXPLAIN SYNTAX SELECT trim(explain) == 'Asterisk' FROM (
+                EXPLAIN AST SELECT * FROM system.numbers LIMIT 10
+            ) WHERE explain LIKE '%Asterisk%'
+        )
+    )
+) FORMAT Null;
+
+CREATE STREAM t1 ( a uint64 ) AS SELECT number AS a,now64() as _tp_time FROM system.numbers LIMIT 10000;
+SELECT * FROM (SELECT sleep(3) AS n) AS a JOIN (SELECT 1 + sleep(3) AS f) AS b ON a.n == b.f; --- sleep to make sure that the stream t1 finishes being created and inserted into
+SELECT rows > 1000 FROM (EXPLAIN ESTIMATE SELECT sum(a) FROM table(t1));
+SELECT count() == 1 FROM (EXPLAIN ESTIMATE SELECT sum(a) FROM table(t1));
+
+DROP STREAM IF EXISTS t1;
+