Merge pull request #36766 from kitaisreal/evaluate-constant-expression-fix

Fix evaluateConstantExpression for subqueries

src/Interpreters/UserDefinedSQLFunctionFactory.cpp

@@ -48,7 +48,7 @@ void UserDefinedSQLFunctionFactory::registerFunction(ContextPtr context, const S
         if (if_not_exists)
             return;
 
-        throw Exception(ErrorCodes::CANNOT_DROP_FUNCTION, "User defined executable function '{}'", function_name);
+        throw Exception(ErrorCodes::FUNCTION_ALREADY_EXISTS, "User defined executable function '{}' already exists", function_name);
     }
 
     std::lock_guard lock(mutex);

src/Interpreters/evaluateConstantExpression.cpp

@@ -13,6 +13,7 @@
 #include <Parsers/ASTFunction.h>
 #include <Parsers/ASTIdentifier.h>
 #include <Parsers/ASTLiteral.h>
+#include <Parsers/ASTSubquery.h>
 #include <Parsers/ExpressionElementParsers.h>
 #include <TableFunctions/TableFunctionFactory.h>
 #include <Common/typeid_cast.h>
@@ -37,7 +38,19 @@ std::pair<Field, std::shared_ptr<const IDataType>> evaluateConstantExpression(co
         return std::make_pair(literal->value, applyVisitor(FieldToDataType(), literal->value));
 
     NamesAndTypesList source_columns = {{ "_dummy", std::make_shared<DataTypeUInt8>() }};
+
     auto ast = node->clone();
+
+    if (ast->as<ASTSubquery>() != nullptr)
+    {
+        /** For subqueries getColumnName if there are no alias will return __subquery_ + 'hash'.
+          * If there is alias getColumnName for subquery will return alias.
+          * In result block name of subquery after QueryAliasesVisitor pass will be _subquery1.
+          * We specify alias for subquery, because we need to get column from result block.
+          */
+        ast->setAlias("constant_expression");
+    }
+
     ReplaceQueryParameterVisitor param_visitor(context->getQueryParameters());
     param_visitor.visit(ast);
 
@@ -46,6 +59,12 @@ std::pair<Field, std::shared_ptr<const IDataType>> evaluateConstantExpression(co
 
     String name = ast->getColumnName();
     auto syntax_result = TreeRewriter(context).analyze(ast, source_columns);
+
+    /// AST potentially could be transformed to literal during TreeRewriter analyze.
+    /// For example if we have SQL user defined function that return literal AS subquery.
+    if (ASTLiteral * literal = ast->as<ASTLiteral>())
+        return std::make_pair(literal->value, applyVisitor(FieldToDataType(), literal->value));
+
     ExpressionActionsPtr expr_for_constant_folding = ExpressionAnalyzer(ast, syntax_result, context).getConstActions();
 
     /// There must be at least one column in the block so that it knows the number of rows.

tests/integration/test_executable_table_function/test.py

@@ -92,6 +92,22 @@ def test_executable_function_input_python(started_cluster):
         == "Key 0\nKey 1\nKey 2\n"
     )
 
+    query = "SELECT * FROM executable('input.py', 'TabSeparated', (SELECT 'value String'), {source})"
+    assert node.query(query.format(source="(SELECT 1)")) == "Key 1\n"
+    assert (
+        node.query(query.format(source="(SELECT id FROM test_data_table)"))
+        == "Key 0\nKey 1\nKey 2\n"
+    )
+
+    node.query("CREATE FUNCTION test_function AS () -> 'value String';")
+    query = "SELECT * FROM executable('input.py', 'TabSeparated', (SELECT test_function()), {source})"
+    assert node.query(query.format(source="(SELECT 1)")) == "Key 1\n"
+    assert (
+        node.query(query.format(source="(SELECT id FROM test_data_table)"))
+        == "Key 0\nKey 1\nKey 2\n"
+    )
+    node.query("DROP FUNCTION test_function;")
+
 
 def test_executable_function_input_sum_python(started_cluster):
     skip_test_msan(node)