From 83787e26f2f1a287c12886ad316733dcbc3676cf Mon Sep 17 00:00:00 2001 From: Maksim Kita Date: Wed, 20 Oct 2021 12:43:03 +0300 Subject: [PATCH] SQLUserDefinedFunctions support lambdas --- .../InterpreterCreateFunctionQuery.cpp | 40 +++---------------- .../InterpreterCreateFunctionQuery.h | 1 - .../UserDefinedSQLFunctionVisitor.cpp | 10 ++++- ...l_user_defined_functions_aliases.reference | 1 + ...098_sql_user_defined_functions_aliases.sql | 4 ++ ...ql_user_defined_functions_lambda.reference | 1 + ...2099_sql_user_defined_functions_lambda.sql | 4 ++ 7 files changed, 24 insertions(+), 37 deletions(-) create mode 100644 tests/queries/0_stateless/02098_sql_user_defined_functions_aliases.reference create mode 100644 tests/queries/0_stateless/02098_sql_user_defined_functions_aliases.sql create mode 100644 tests/queries/0_stateless/02099_sql_user_defined_functions_lambda.reference create mode 100644 tests/queries/0_stateless/02099_sql_user_defined_functions_lambda.sql diff --git a/src/Interpreters/InterpreterCreateFunctionQuery.cpp b/src/Interpreters/InterpreterCreateFunctionQuery.cpp index ccb5f4040ec..9d92466c440 100644 --- a/src/Interpreters/InterpreterCreateFunctionQuery.cpp +++ b/src/Interpreters/InterpreterCreateFunctionQuery.cpp @@ -1,14 +1,17 @@ +#include + +#include + #include #include #include #include #include #include -#include #include #include #include -#include + namespace DB { @@ -66,42 +69,9 @@ void InterpreterCreateFunctionQuery::validateFunction(ASTPtr function, const Str } ASTPtr function_body = function->as()->children.at(0)->children.at(1); - std::unordered_set identifiers_in_body = getIdentifiers(function_body); - - for (const auto & identifier : identifiers_in_body) - { - if (!arguments.contains(identifier)) - throw Exception(ErrorCodes::UNKNOWN_IDENTIFIER, "Identifier {} does not exist in arguments", backQuote(identifier)); - } - validateFunctionRecursiveness(function_body, name); } -std::unordered_set InterpreterCreateFunctionQuery::getIdentifiers(ASTPtr node) -{ - std::unordered_set identifiers; - - std::stack ast_nodes_to_process; - ast_nodes_to_process.push(node); - - while (!ast_nodes_to_process.empty()) - { - auto ast_node_to_process = ast_nodes_to_process.top(); - ast_nodes_to_process.pop(); - - for (const auto & child : ast_node_to_process->children) - { - auto identifier_name_opt = tryGetIdentifierName(child); - if (identifier_name_opt) - identifiers.insert(identifier_name_opt.value()); - - ast_nodes_to_process.push(child); - } - } - - return identifiers; -} - void InterpreterCreateFunctionQuery::validateFunctionRecursiveness(ASTPtr node, const String & function_to_create) { for (const auto & child : node->children) diff --git a/src/Interpreters/InterpreterCreateFunctionQuery.h b/src/Interpreters/InterpreterCreateFunctionQuery.h index fdc03b379db..a67fdb9605d 100644 --- a/src/Interpreters/InterpreterCreateFunctionQuery.h +++ b/src/Interpreters/InterpreterCreateFunctionQuery.h @@ -22,7 +22,6 @@ public: private: static void validateFunction(ASTPtr function, const String & name); - static std::unordered_set getIdentifiers(ASTPtr node); static void validateFunctionRecursiveness(ASTPtr node, const String & function_to_create); ASTPtr query_ptr; diff --git a/src/Interpreters/UserDefinedSQLFunctionVisitor.cpp b/src/Interpreters/UserDefinedSQLFunctionVisitor.cpp index cc5db020387..8df6932b153 100644 --- a/src/Interpreters/UserDefinedSQLFunctionVisitor.cpp +++ b/src/Interpreters/UserDefinedSQLFunctionVisitor.cpp @@ -25,6 +25,7 @@ void UserDefinedSQLFunctionMatcher::visit(ASTPtr & ast, Data &) return; auto result = tryToReplaceFunction(*function); + if (result) ast = result; } @@ -83,9 +84,16 @@ ASTPtr UserDefinedSQLFunctionMatcher::tryToReplaceFunction(const ASTFunction & f if (identifier_name_opt) { auto function_argument_it = identifier_name_to_function_argument.find(*identifier_name_opt); - assert(function_argument_it != identifier_name_to_function_argument.end()); + if (function_argument_it == identifier_name_to_function_argument.end()) + continue; + + auto child_alias = child->tryGetAlias(); child = function_argument_it->second->clone(); + + if (!child_alias.empty()) + child->setAlias(child_alias); + continue; } diff --git a/tests/queries/0_stateless/02098_sql_user_defined_functions_aliases.reference b/tests/queries/0_stateless/02098_sql_user_defined_functions_aliases.reference new file mode 100644 index 00000000000..8ab2f6d0ac6 --- /dev/null +++ b/tests/queries/0_stateless/02098_sql_user_defined_functions_aliases.reference @@ -0,0 +1 @@ +8 4 diff --git a/tests/queries/0_stateless/02098_sql_user_defined_functions_aliases.sql b/tests/queries/0_stateless/02098_sql_user_defined_functions_aliases.sql new file mode 100644 index 00000000000..fef2daf8fd1 --- /dev/null +++ b/tests/queries/0_stateless/02098_sql_user_defined_functions_aliases.sql @@ -0,0 +1,4 @@ +-- Tags: no-parallel +CREATE FUNCTION alias_function AS x -> (((x * 2) AS x_doubled) + x_doubled); +SELECT alias_function(2); +DROP FUNCTION alias_function; diff --git a/tests/queries/0_stateless/02099_sql_user_defined_functions_lambda.reference b/tests/queries/0_stateless/02099_sql_user_defined_functions_lambda.reference new file mode 100644 index 00000000000..8f6cd5ccd03 --- /dev/null +++ b/tests/queries/0_stateless/02099_sql_user_defined_functions_lambda.reference @@ -0,0 +1 @@ +[2,4,6] diff --git a/tests/queries/0_stateless/02099_sql_user_defined_functions_lambda.sql b/tests/queries/0_stateless/02099_sql_user_defined_functions_lambda.sql new file mode 100644 index 00000000000..e66651c1e5a --- /dev/null +++ b/tests/queries/0_stateless/02099_sql_user_defined_functions_lambda.sql @@ -0,0 +1,4 @@ +-- Tags: no-parallel +CREATE FUNCTION lambda_function AS x -> arrayMap(array_element -> array_element * 2, x); +SELECT lambda_function([1,2,3]); +DROP FUNCTION lambda_function;