Normalize the query before parsing aliases.

This commit is contained in:
Yarik Briukhovetskyi 2024-11-20 19:48:10 +01:00 committed by GitHub
parent e0fb20f6d4
commit 7a30ea2e48
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 29 additions and 17 deletions

View File

@ -11,7 +11,9 @@
#include <Parsers/ASTExpressionList.h>
#include <Parsers/ASTIdentifier.h>
#include <Functions/UserDefined/UserDefinedSQLFunctionFactory.h>
#include "Parsers/ASTColumnDeclaration.h"
#include <Interpreters/QueryAliasesVisitor.h>
#include <Interpreters/MarkTableIdentifiersVisitor.h>
#include <Interpreters/QueryNormalizer.h>
namespace DB
@ -22,14 +24,14 @@ namespace ErrorCodes
extern const int UNSUPPORTED_METHOD;
}
void UserDefinedSQLFunctionVisitor::visit(ASTPtr & ast)
void UserDefinedSQLFunctionVisitor::visit(ASTPtr & ast, ContextPtr context_)
{
chassert(ast);
if (const auto * function = ast->template as<ASTFunction>())
{
std::unordered_set<std::string> udf_in_replace_process;
auto replace_result = tryToReplaceFunction(*function, udf_in_replace_process);
auto replace_result = tryToReplaceFunction(*function, udf_in_replace_process, context_);
if (replace_result)
ast = replace_result;
}
@ -40,7 +42,7 @@ void UserDefinedSQLFunctionVisitor::visit(ASTPtr & ast)
return;
auto * old_ptr = child.get();
visit(child);
visit(child, context_);
auto * new_ptr = child.get();
/// Some AST classes have naked pointers to children elements as members.
@ -50,16 +52,16 @@ void UserDefinedSQLFunctionVisitor::visit(ASTPtr & ast)
}
}
void UserDefinedSQLFunctionVisitor::visit(IAST * ast)
void UserDefinedSQLFunctionVisitor::visit(IAST * ast, ContextPtr context_)
{
if (!ast)
return;
for (auto & child : ast->children)
visit(child);
visit(child, context_);
}
ASTPtr UserDefinedSQLFunctionVisitor::tryToReplaceFunction(const ASTFunction & function, std::unordered_set<std::string> & udf_in_replace_process)
ASTPtr UserDefinedSQLFunctionVisitor::tryToReplaceFunction(const ASTFunction & function, std::unordered_set<std::string> & udf_in_replace_process, ContextPtr context_)
{
if (udf_in_replace_process.find(function.name) != udf_in_replace_process.end())
throw Exception(ErrorCodes::UNSUPPORTED_METHOD,
@ -101,6 +103,17 @@ ASTPtr UserDefinedSQLFunctionVisitor::tryToReplaceFunction(const ASTFunction & f
auto function_body_to_update = function_core_expression->children.at(1)->clone();
Aliases aliases;
QueryAliasesVisitor(aliases).visit(function_body_to_update);
/// Mark table ASTIdentifiers with not a column marker
MarkTableIdentifiersVisitor::Data identifiers_data{aliases};
MarkTableIdentifiersVisitor(identifiers_data).visit(function_body_to_update);
/// Common subexpression elimination. Rewrite rules.
QueryNormalizer::Data normalizer_data(aliases, {}, true, context_->getSettingsRef(), true, false);
QueryNormalizer(normalizer_data).visit(function_body_to_update);
auto expression_list = std::make_shared<ASTExpressionList>();
expression_list->children.emplace_back(std::move(function_body_to_update));
@ -116,7 +129,7 @@ ASTPtr UserDefinedSQLFunctionVisitor::tryToReplaceFunction(const ASTFunction & f
{
if (auto * inner_function = child->as<ASTFunction>())
{
auto replace_result = tryToReplaceFunction(*inner_function, udf_in_replace_process);
auto replace_result = tryToReplaceFunction(*inner_function, udf_in_replace_process, context_);
if (replace_result)
child = replace_result;
}

View File

@ -19,13 +19,13 @@ class ASTFunction;
* After applying visitor:
* SELECT number + 1 FROM system.numbers LIMIT 10;
*/
class UserDefinedSQLFunctionVisitor
class UserDefinedSQLFunctionVisitor : WithContext
{
public:
static void visit(ASTPtr & ast);
static void visit(ASTPtr & ast, ContextPtr context_);
private:
static void visit(IAST *);
static ASTPtr tryToReplaceFunction(const ASTFunction & function, std::unordered_set<std::string> & udf_in_replace_process);
static void visit(IAST *, ContextPtr context_);
static ASTPtr tryToReplaceFunction(const ASTFunction & function, std::unordered_set<std::string> & udf_in_replace_process, ContextPtr context_);
};

View File

@ -99,7 +99,7 @@ BlockIO InterpreterAlterQuery::executeToTable(const ASTAlterQuery & alter)
BlockIO res;
if (!UserDefinedSQLFunctionFactory::instance().empty())
UserDefinedSQLFunctionVisitor::visit(query_ptr);
UserDefinedSQLFunctionVisitor::visit(query_ptr, getContext());
auto table_id = getContext()->tryResolveStorageID(alter);
StoragePtr table;

View File

@ -1599,7 +1599,7 @@ BlockIO InterpreterCreateQuery::createTable(ASTCreateQuery & create)
// substitute possible UDFs with their definitions
if (!UserDefinedSQLFunctionFactory::instance().empty())
UserDefinedSQLFunctionVisitor::visit(query_ptr);
UserDefinedSQLFunctionVisitor::visit(query_ptr, getContext());
/// Set and retrieve list of columns, indices and constraints. Set table engine if needed. Rewrite query in canonical way.
TableProperties properties = getTablePropertiesAndNormalizeCreateQuery(create, mode);

View File

@ -38,7 +38,6 @@ private:
static void visit(const ASTSubquery & subquery, const ASTPtr & ast, Data & data);
static void visit(const ASTArrayJoin &, const ASTPtr & ast, Data & data);
static void visitOther(const ASTPtr & ast, Data & data);
static bool checkIfNamesAreSame(Data & data, const ASTPtr & ast);
};
/// Visits AST nodes and collect their aliases in one map (with links to source nodes).

View File

@ -1571,7 +1571,7 @@ void TreeRewriter::normalize(
ASTPtr & query, Aliases & aliases, const NameSet & source_columns_set, bool ignore_alias, const Settings & settings, bool allow_self_aliases, ContextPtr context_, bool is_create_parameterized_view)
{
if (!UserDefinedSQLFunctionFactory::instance().empty())
UserDefinedSQLFunctionVisitor::visit(query);
UserDefinedSQLFunctionVisitor::visit(query, context_);
CustomizeCountDistinctVisitor::Data data_count_distinct{settings[Setting::count_distinct_implementation]};
CustomizeCountDistinctVisitor(data_count_distinct).visit(query);