Skip to content

Analyzer: compute ALIAS columns right after reading #59595

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vdimir merged 3 commits into from
Feb 29, 2024
Merged

Analyzer: compute ALIAS columns right after reading #59595

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5a9026f
Analyzer: compute ALIAS columns right after reading
vdimir Feb 5, 2024
9f1a073
review fix
vdimir Feb 21, 2024
973c253
fix 02911_support_alias_column_in_indices
vdimir Feb 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions src/Analyzer/Passes/QueryAnalysisPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6651,7 +6651,6 @@ void QueryAnalyzer::initializeTableExpressionData(const QueryTreeNodePtr & table
if (column_default && column_default->kind == ColumnDefaultKind::Alias)
{
auto alias_expression = buildQueryTree(column_default->expression, scope.context);
alias_expression = buildCastFunction(alias_expression, column_name_and_type.type, scope.context, false /*resolve*/);
auto column_node = std::make_shared<ColumnNode>(column_name_and_type, std::move(alias_expression), table_expression_node);
column_name_to_column_node.emplace(column_name_and_type.name, column_node);
alias_columns_to_resolve.emplace_back(column_name_and_type.name, column_node);
Expand Down Expand Up @@ -6684,7 +6683,9 @@ void QueryAnalyzer::initializeTableExpressionData(const QueryTreeNodePtr & table
alias_column_resolve_scope,
false /*allow_lambda_expression*/,
false /*allow_table_expression*/);

auto & resolved_expression = alias_column_to_resolve->getExpression();
if (!resolved_expression->getResultType()->equals(*alias_column_to_resolve->getResultType()))
resolved_expression = buildCastFunction(resolved_expression, alias_column_to_resolve->getResultType(), scope.context, true);
column_name_to_column_node = std::move(alias_column_resolve_scope.column_name_to_column_node);
column_name_to_column_node[alias_column_to_resolve_name] = alias_column_to_resolve;
}
Expand Down
7 changes: 6 additions & 1 deletion src/Interpreters/getHeaderForProcessingStage.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,12 @@ Block getHeaderForProcessingStage(

auto & table_expression_data = query_info.planner_context->getTableExpressionDataOrThrow(left_table_expression);
const auto & query_context = query_info.planner_context->getQueryContext();
auto columns = table_expression_data.getColumns();

NamesAndTypes columns;
const auto & column_name_to_column = table_expression_data.getColumnNameToColumn();
for (const auto & column_name : table_expression_data.getSelectedColumnsNames())
columns.push_back(column_name_to_column.at(column_name));

auto new_query_node = buildSubqueryToReadColumnsFromTableExpression(columns, left_table_expression, query_context);
query = new_query_node->toAST();
}
Expand Down
160 changes: 95 additions & 65 deletions src/Planner/CollectTableExpressionData.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,34 +29,13 @@ namespace
class CollectSourceColumnsVisitor : public InDepthQueryTreeVisitor<CollectSourceColumnsVisitor>
{
public:
explicit CollectSourceColumnsVisitor(PlannerContext & planner_context_)
explicit CollectSourceColumnsVisitor(PlannerContextPtr & planner_context_, bool keep_alias_columns_ = true)
: planner_context(planner_context_)
, keep_alias_columns(keep_alias_columns_)
{}

void visitImpl(QueryTreeNodePtr & node)
{
/// Special case for USING clause which contains references to ALIAS columns.
/// We can not modify such ColumnNode.
if (auto * join_node = node->as<JoinNode>())
{
if (!join_node->isUsingJoinExpression())
return;

auto & using_list = join_node->getJoinExpression()->as<ListNode&>();
for (auto & using_element : using_list)
{
auto & column_node = using_element->as<ColumnNode&>();
/// This list contains column nodes from left and right tables.
auto & columns_from_subtrees = column_node.getExpressionOrThrow()->as<ListNode&>().getNodes();

/// Visit left table column node.
visitUsingColumn(columns_from_subtrees[0]);
/// Visit right table column node.
visitUsingColumn(columns_from_subtrees[1]);
}
return;
}

auto * column_node = node->as<ColumnNode>();
if (!column_node)
return;
Expand All @@ -72,22 +51,55 @@ class CollectSourceColumnsVisitor : public InDepthQueryTreeVisitor<CollectSource

/// JOIN using expression
if (column_node->hasExpression() && column_source_node_type == QueryTreeNodeType::JOIN)
{
auto & columns_from_subtrees = column_node->getExpression()->as<ListNode &>().getNodes();
if (columns_from_subtrees.size() != 2)
throw Exception(ErrorCodes::LOGICAL_ERROR,
"Expected two columns in JOIN using expression for column {}", column_node->dumpTree());

visit(columns_from_subtrees[0]);
visit(columns_from_subtrees[1]);
return;
}

auto & table_expression_data = planner_context.getOrCreateTableExpressionData(column_source_node);
auto & table_expression_data = planner_context->getOrCreateTableExpressionData(column_source_node);

if (column_node->hasExpression() && column_source_node_type != QueryTreeNodeType::ARRAY_JOIN)
if (isAliasColumn(node))
{
/// Replace ALIAS column with expression
/// Column is an ALIAS column with expression
bool column_already_exists = table_expression_data.hasColumn(column_node->getColumnName());
if (!column_already_exists)
{
auto column_identifier = planner_context.getGlobalPlannerContext()->createColumnIdentifier(node);
table_expression_data.addAliasColumnName(column_node->getColumnName(), column_identifier);
CollectSourceColumnsVisitor visitor_for_alias_column(planner_context);
/// While we are processing expression of ALIAS columns we should not add source columns to selected.
/// See also comment for `select_added_columns`
visitor_for_alias_column.select_added_columns = false;
visitor_for_alias_column.keep_alias_columns = keep_alias_columns;
visitor_for_alias_column.visit(column_node->getExpression());

if (!keep_alias_columns)
{
/// For PREWHERE we can just replace ALIAS column with it's expression,
/// because ActionsDAG for PREWHERE applied right on top of table expression
/// and cannot affect subqueries or other table expressions.
node = column_node->getExpression();
return;
}

auto column_identifier = planner_context->getGlobalPlannerContext()->createColumnIdentifier(node);

ActionsDAGPtr alias_column_actions_dag = std::make_shared<ActionsDAG>();
PlannerActionsVisitor actions_visitor(planner_context, false);
auto outputs = actions_visitor.visit(alias_column_actions_dag, column_node->getExpression());
if (outputs.size() != 1)
throw Exception(ErrorCodes::LOGICAL_ERROR,
"Expected single output in actions dag for alias column {}. Actual {}", column_node->dumpTree(), outputs.size());
const auto & column_name = column_node->getColumnName();
const auto & alias_node = alias_column_actions_dag->addAlias(*outputs[0], column_name);
alias_column_actions_dag->addOrReplaceInOutputs(alias_node);
table_expression_data.addAliasColumn(column_node->getColumn(), column_identifier, alias_column_actions_dag, select_added_columns);
}

node = column_node->getExpression();
visitImpl(node);
return;
}

Expand All @@ -102,45 +114,58 @@ class CollectSourceColumnsVisitor : public InDepthQueryTreeVisitor<CollectSource

bool column_already_exists = table_expression_data.hasColumn(column_node->getColumnName());
if (column_already_exists)
{
/// Column may be added when we collected data for ALIAS column
/// But now we see it directly in the query, so make sure it's marked as selected
if (select_added_columns)
table_expression_data.markSelectedColumn(column_node->getColumnName());
return;
}

auto column_identifier = planner_context.getGlobalPlannerContext()->createColumnIdentifier(node);
table_expression_data.addColumn(column_node->getColumn(), column_identifier);
auto column_identifier = planner_context->getGlobalPlannerContext()->createColumnIdentifier(node);
table_expression_data.addColumn(column_node->getColumn(), column_identifier, select_added_columns);
}

static bool needChildVisit(const QueryTreeNodePtr & parent, const QueryTreeNodePtr & child_node)
static bool isAliasColumn(const QueryTreeNodePtr & node)
{
if (auto * join_node = parent->as<JoinNode>())
{
if (join_node->getJoinExpression() == child_node && join_node->isUsingJoinExpression())
return false;
}
auto child_node_type = child_node->getNodeType();
return !(child_node_type == QueryTreeNodeType::QUERY || child_node_type == QueryTreeNodeType::UNION);
const auto * column_node = node->as<ColumnNode>();
if (!column_node || !column_node->hasExpression())
return false;
const auto & column_source = column_node->getColumnSourceOrNull();
if (!column_source)
return false;
return column_source->getNodeType() != QueryTreeNodeType::JOIN &&
column_source->getNodeType() != QueryTreeNodeType::ARRAY_JOIN;
}

private:

void visitUsingColumn(QueryTreeNodePtr & node)
static bool needChildVisit(const QueryTreeNodePtr & parent_node, const QueryTreeNodePtr & child_node)
{
auto & column_node = node->as<ColumnNode&>();
if (column_node.hasExpression())
{
auto & table_expression_data = planner_context.getOrCreateTableExpressionData(column_node.getColumnSource());
bool column_already_exists = table_expression_data.hasColumn(column_node.getColumnName());
if (column_already_exists)
return;

auto column_identifier = planner_context.getGlobalPlannerContext()->createColumnIdentifier(node);
table_expression_data.addAliasColumnName(column_node.getColumnName(), column_identifier);
auto child_node_type = child_node->getNodeType();
return !(child_node_type == QueryTreeNodeType::QUERY ||
child_node_type == QueryTreeNodeType::UNION ||
isAliasColumn(parent_node));
}

visitImpl(column_node.getExpressionOrThrow());
}
else
visitImpl(node);
void setKeepAliasColumns(bool keep_alias_columns_)
{
keep_alias_columns = keep_alias_columns_;
}

PlannerContext & planner_context;
private:
PlannerContextPtr & planner_context;

/// Replace ALIAS columns with their expressions or register them in table expression data.
/// Usually we can replace them when we build some "local" actions DAG
/// (for example Row Policy or PREWHERE) that is applied on top of the table expression.
/// In other cases, we keep ALIAS columns as ColumnNode with an expression child node,
/// and handle them in the Planner by inserting ActionsDAG to compute them after reading from storage.
bool keep_alias_columns = true;
Copy link
Member

@novikd novikd Feb 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When can not we do it?

Copy link
Member Author

@vdimir vdimir Feb 21, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In other cases, when we keep ColumnNode with expression and compute after read. Added it to comment


/// Flag `select_added_columns` indicates if we should mark column as explicitly selected.
/// For example, for table with columns (a Int32, b ALIAS a+1) and query SELECT b FROM table
/// Column `b` is selected explicitly by user, but not `a` (that is also read though).
/// Distinguishing such columns is important for checking access rights for ALIAS columns.
bool select_added_columns = true;
};

class CollectPrewhereTableExpressionVisitor : public ConstInDepthQueryTreeVisitor<CollectPrewhereTableExpressionVisitor>
Expand Down Expand Up @@ -274,7 +299,7 @@ void collectTableExpressionData(QueryTreeNodePtr & query_node, PlannerContextPtr
}
}

CollectSourceColumnsVisitor collect_source_columns_visitor(*planner_context);
CollectSourceColumnsVisitor collect_source_columns_visitor(planner_context);
for (auto & node : query_node_typed.getChildren())
{
if (!node || node == query_node_typed.getPrewhere())
Expand All @@ -300,21 +325,26 @@ void collectTableExpressionData(QueryTreeNodePtr & query_node, PlannerContextPtr
}

auto & table_expression_data = planner_context->getOrCreateTableExpressionData(prewhere_table_expression);
const auto & column_names = table_expression_data.getColumnNames();
NameSet required_column_names_without_prewhere(column_names.begin(), column_names.end());
const auto & read_column_names = table_expression_data.getColumnNames();
NameSet required_column_names_without_prewhere(read_column_names.begin(), read_column_names.end());
const auto & selected_column_names = table_expression_data.getSelectedColumnsNames();
required_column_names_without_prewhere.insert(selected_column_names.begin(), selected_column_names.end());

collect_source_columns_visitor.setKeepAliasColumns(false);
collect_source_columns_visitor.visit(query_node_typed.getPrewhere());

auto prewhere_actions_dag = std::make_shared<ActionsDAG>();

QueryTreeNodePtr query_tree_node = query_node_typed.getPrewhere();

PlannerActionsVisitor visitor(planner_context, false /*use_column_identifier_as_action_node_name*/);
auto expression_nodes = visitor.visit(prewhere_actions_dag, query_node_typed.getPrewhere());
auto expression_nodes = visitor.visit(prewhere_actions_dag, query_tree_node);
if (expression_nodes.size() != 1)
throw Exception(ErrorCodes::ILLEGAL_PREWHERE,
"Invalid PREWHERE. Expected single boolean expression. In query {}",
query_node->formatASTForErrorMessage());

prewhere_actions_dag->getOutputs().push_back(expression_nodes[0]);
prewhere_actions_dag->getOutputs().push_back(expression_nodes.back());

for (const auto & prewhere_input_node : prewhere_actions_dag->getInputs())
if (required_column_names_without_prewhere.contains(prewhere_input_node->result_name))
Expand All @@ -324,9 +354,9 @@ void collectTableExpressionData(QueryTreeNodePtr & query_node, PlannerContextPtr
}
}

void collectSourceColumns(QueryTreeNodePtr & expression_node, PlannerContextPtr & planner_context)
void collectSourceColumns(QueryTreeNodePtr & expression_node, PlannerContextPtr & planner_context, bool keep_alias_columns)
{
CollectSourceColumnsVisitor collect_source_columns_visitor(*planner_context);
CollectSourceColumnsVisitor collect_source_columns_visitor(planner_context, keep_alias_columns);
collect_source_columns_visitor.visit(expression_node);
}

Expand Down
2 changes: 1 addition & 1 deletion src/Planner/CollectTableExpressionData.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,6 @@ void collectTableExpressionData(QueryTreeNodePtr & query_node, PlannerContextPtr
*
* ALIAS table column nodes are registered in table expression data and replaced in query tree with inner alias expression.
*/
void collectSourceColumns(QueryTreeNodePtr & expression_node, PlannerContextPtr & planner_context);
void collectSourceColumns(QueryTreeNodePtr & expression_node, PlannerContextPtr & planner_context, bool keep_alias_columns = true);

}
5 changes: 4 additions & 1 deletion src/Planner/PlannerActionsVisitor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -451,13 +451,15 @@ class PlannerActionsVisitorImpl
std::unordered_map<QueryTreeNodePtr, std::string> node_to_node_name;
const PlannerContextPtr planner_context;
ActionNodeNameHelper action_node_name_helper;
bool use_column_identifier_as_action_node_name;
};

PlannerActionsVisitorImpl::PlannerActionsVisitorImpl(ActionsDAGPtr actions_dag,
const PlannerContextPtr & planner_context_,
bool use_column_identifier_as_action_node_name_)
: planner_context(planner_context_)
, action_node_name_helper(node_to_node_name, *planner_context, use_column_identifier_as_action_node_name_)
, use_column_identifier_as_action_node_name(use_column_identifier_as_action_node_name_)
{
actions_stack.emplace_back(std::move(actions_dag), nullptr);
}
Expand Down Expand Up @@ -503,7 +505,8 @@ PlannerActionsVisitorImpl::NodeNameAndNodeMinLevel PlannerActionsVisitorImpl::vi
{
auto column_node_name = action_node_name_helper.calculateActionNodeName(node);
const auto & column_node = node->as<ColumnNode &>();

if (column_node.hasExpression() && !use_column_identifier_as_action_node_name)
return visitImpl(column_node.getExpression());
Int64 actions_stack_size = static_cast<Int64>(actions_stack.size() - 1);
for (Int64 i = actions_stack_size; i >= 0; --i)
{
Expand Down
Loading