ClickHouse/dbms/src/Interpreters/RequiredSourceColumnsVisitor.cpp

219 lines
6.2 KiB
C++

#include <Interpreters/RequiredSourceColumnsVisitor.h>
#include <Common/typeid_cast.h>
#include <Core/Names.h>
#include <Parsers/IAST.h>
#include <Parsers/ASTIdentifier.h>
#include <Parsers/ASTFunction.h>
#include <Parsers/ASTSelectQuery.h>
#include <Parsers/ASTSubquery.h>
#include <Parsers/ASTTablesInSelectQuery.h>
namespace DB
{
namespace ErrorCodes
{
extern const int TYPE_MISMATCH;
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
}
static std::vector<String> extractNamesFromLambda(const ASTFunction & node)
{
if (node.arguments->children.size() != 2)
throw Exception("lambda requires two arguments", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
const auto * lambda_args_tuple = node.arguments->children[0]->as<ASTFunction>();
if (!lambda_args_tuple || lambda_args_tuple->name != "tuple")
throw Exception("First argument of lambda must be a tuple", ErrorCodes::TYPE_MISMATCH);
std::vector<String> names;
for (auto & child : lambda_args_tuple->arguments->children)
{
const auto * identifier = child->as<ASTIdentifier>();
if (!identifier)
throw Exception("lambda argument declarations must be identifiers", ErrorCodes::TYPE_MISMATCH);
names.push_back(identifier->name);
}
return names;
}
bool RequiredSourceColumnsMatcher::needChildVisit(ASTPtr & node, const ASTPtr & child)
{
if (child->as<ASTSelectQuery>())
return false;
/// Processed. Do not need children.
if (node->as<ASTTableExpression>() || node->as<ASTArrayJoin>() || node->as<ASTSelectQuery>())
return false;
if (const auto * f = node->as<ASTFunction>())
{
/// "indexHint" is a special function for index analysis. Everything that is inside it is not calculated. @sa KeyCondition
/// "lambda" visit children itself.
if (f->name == "indexHint" || f->name == "lambda")
return false;
}
return true;
}
void RequiredSourceColumnsMatcher::visit(ASTPtr & ast, Data & data)
{
/// results are columns
if (auto * t = ast->as<ASTIdentifier>())
{
visit(*t, ast, data);
return;
}
if (auto * t = ast->as<ASTFunction>())
{
data.addColumnAliasIfAny(*ast);
visit(*t, ast, data);
return;
}
/// results are tables
if (auto * t = ast->as<ASTTablesInSelectQueryElement>())
{
visit(*t, ast, data);
return;
}
if (auto * t = ast->as<ASTTableExpression>())
{
visit(*t, ast, data);
return;
}
if (auto * t = ast->as<ASTSelectQuery>())
{
data.addTableAliasIfAny(*ast);
visit(*t, ast, data);
return;
}
if (ast->as<ASTSubquery>())
{
data.addTableAliasIfAny(*ast);
return;
}
/// other
if (auto * t = ast->as<ASTArrayJoin>())
{
data.has_array_join = true;
visit(*t, ast, data);
return;
}
}
void RequiredSourceColumnsMatcher::visit(ASTSelectQuery & select, const ASTPtr &, Data & data)
{
/// special case for top-level SELECT items: they are publics
for (auto & node : select.select()->children)
{
if (const auto * identifier = node->as<ASTIdentifier>())
data.addColumnIdentifier(*identifier);
else
data.addColumnAliasIfAny(*node);
}
std::vector<ASTPtr *> out;
for (auto & node : select.children)
if (node != select.select())
Visitor(data).visit(node);
/// revisit select_expression_list (with children) when all the aliases are set
Visitor(data).visit(select.refSelect());
}
void RequiredSourceColumnsMatcher::visit(const ASTIdentifier & node, const ASTPtr &, Data & data)
{
if (node.name.empty())
throw Exception("Expected not empty name", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
if (!data.private_aliases.count(node.name))
data.addColumnIdentifier(node);
}
void RequiredSourceColumnsMatcher::visit(const ASTFunction & node, const ASTPtr &, Data & data)
{
/// Do not add formal parameters of the lambda expression
if (node.name == "lambda")
{
Names local_aliases;
for (const auto & name : extractNamesFromLambda(node))
if (data.private_aliases.insert(name).second)
local_aliases.push_back(name);
/// visit child with masked local aliases
RequiredSourceColumnsVisitor(data).visit(node.arguments->children[1]);
for (const auto & name : local_aliases)
data.private_aliases.erase(name);
}
}
void RequiredSourceColumnsMatcher::visit(ASTTablesInSelectQueryElement & node, const ASTPtr &, Data & data)
{
ASTTableExpression * expr = nullptr;
ASTTableJoin * join = nullptr;
for (auto & child : node.children)
{
if (auto * e = child->as<ASTTableExpression>())
expr = e;
if (auto * j = child->as<ASTTableJoin>())
join = j;
}
if (join)
data.has_table_join = true;
data.tables.emplace_back(ColumnNamesContext::JoinedTable{expr, join});
}
/// ASTIdentifiers here are tables. Do not visit them as generic ones.
void RequiredSourceColumnsMatcher::visit(ASTTableExpression & node, const ASTPtr &, Data & data)
{
if (node.database_and_table_name)
data.addTableAliasIfAny(*node.database_and_table_name);
if (node.table_function)
data.addTableAliasIfAny(*node.table_function);
if (node.subquery)
data.addTableAliasIfAny(*node.subquery);
}
void RequiredSourceColumnsMatcher::visit(const ASTArrayJoin & node, const ASTPtr &, Data & data)
{
ASTPtr expression_list = node.expression_list;
if (!expression_list || expression_list->children.empty())
throw Exception("Expected not empty expression_list", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
std::vector<ASTPtr *> out;
/// Tech debt. Ignore ARRAY JOIN top-level identifiers and aliases. There's its own logic for them.
for (auto & expr : expression_list->children)
{
data.addArrayJoinAliasIfAny(*expr);
if (const auto * identifier = expr->as<ASTIdentifier>())
{
data.addArrayJoinIdentifier(*identifier);
continue;
}
out.push_back(&expr);
}
for (ASTPtr * add_node : out)
Visitor(data).visit(*add_node);
}
}