This commit is contained in:
Nikolai Kochetov 2024-04-08 16:54:43 +02:00 committed by GitHub
commit b97e16f2b4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 49 additions and 25 deletions

View File

@ -24,6 +24,34 @@ namespace ErrorCodes
extern const int UNSUPPORTED_METHOD;
}
QueryTreeNodePtr makeExecutableSubqueryForIn(const QueryTreeNodePtr & in_second_argument, const ContextPtr & context)
{
auto subquery_to_execute = in_second_argument;
if (auto * table_node = in_second_argument->as<TableNode>())
{
auto storage_snapshot = table_node->getStorageSnapshot();
auto columns_to_select = storage_snapshot->getColumns(GetColumnsOptions(GetColumnsOptions::Ordinary));
size_t columns_to_select_size = columns_to_select.size();
auto column_nodes_to_select = std::make_shared<ListNode>();
column_nodes_to_select->getNodes().reserve(columns_to_select_size);
NamesAndTypes projection_columns;
projection_columns.reserve(columns_to_select_size);
for (auto & column : columns_to_select)
{
column_nodes_to_select->getNodes().emplace_back(std::make_shared<ColumnNode>(column, subquery_to_execute));
projection_columns.emplace_back(column.name, column.type);
}
auto subquery_for_table = std::make_shared<QueryNode>(Context::createCopy(context));
subquery_for_table->setIsSubquery(true);
subquery_for_table->getProjectionNode() = std::move(column_nodes_to_select);
subquery_for_table->getJoinTree() = std::move(subquery_to_execute);
subquery_for_table->resolveProjectionColumns(std::move(projection_columns));
subquery_to_execute = std::move(subquery_for_table);
}
return subquery_to_execute;
}
namespace
{
@ -88,29 +116,7 @@ public:
if (sets.findSubquery(set_key))
return;
auto subquery_to_execute = in_second_argument;
if (auto * table_node = in_second_argument->as<TableNode>())
{
auto storage_snapshot = table_node->getStorageSnapshot();
auto columns_to_select = storage_snapshot->getColumns(GetColumnsOptions(GetColumnsOptions::Ordinary));
size_t columns_to_select_size = columns_to_select.size();
auto column_nodes_to_select = std::make_shared<ListNode>();
column_nodes_to_select->getNodes().reserve(columns_to_select_size);
NamesAndTypes projection_columns;
projection_columns.reserve(columns_to_select_size);
for (auto & column : columns_to_select)
{
column_nodes_to_select->getNodes().emplace_back(std::make_shared<ColumnNode>(column, subquery_to_execute));
projection_columns.emplace_back(column.name, column.type);
}
auto subquery_for_table = std::make_shared<QueryNode>(Context::createCopy(planner_context.getQueryContext()));
subquery_for_table->setIsSubquery(true);
subquery_for_table->getProjectionNode() = std::move(column_nodes_to_select);
subquery_for_table->getJoinTree() = std::move(subquery_to_execute);
subquery_for_table->resolveProjectionColumns(std::move(projection_columns));
subquery_to_execute = std::move(subquery_for_table);
}
auto subquery_to_execute = makeExecutableSubqueryForIn(in_second_argument, planner_context.getQueryContext());
sets.addFromSubquery(set_key, std::move(subquery_to_execute), settings);
}
else

View File

@ -14,4 +14,8 @@ struct SelectQueryOptions;
*/
void collectSets(const QueryTreeNodePtr & node, PlannerContext & planner_context);
/// Build subquery which we execute for IN function.
/// It is needed to support `IN table` case.
QueryTreeNodePtr makeExecutableSubqueryForIn(const QueryTreeNodePtr & in_second_argument, const ContextPtr & context);
}

View File

@ -13,6 +13,7 @@
#include <Functions/FunctionFactory.h>
#include <Interpreters/DatabaseCatalog.h>
#include <Interpreters/InterpreterSelectQueryAnalyzer.h>
#include <Planner/CollectSets.h>
#include <Planner/Utils.h>
#include <Processors/Executors/CompletedPipelineExecutor.h>
#include <Processors/QueryPlan/ExpressionStep.h>
@ -361,10 +362,12 @@ QueryTreeNodePtr buildQueryTreeForShard(const PlannerContextPtr & planner_contex
{
auto & in_function_subquery_node = in_function_node->getArguments().getNodes().at(1);
auto in_function_node_type = in_function_subquery_node->getNodeType();
if (in_function_node_type != QueryTreeNodeType::QUERY && in_function_node_type != QueryTreeNodeType::UNION)
if (in_function_node_type != QueryTreeNodeType::QUERY && in_function_node_type != QueryTreeNodeType::UNION && in_function_node_type != QueryTreeNodeType::TABLE)
continue;
auto temporary_table_expression_node = executeSubqueryNode(in_function_subquery_node,
auto subquery_to_execute = makeExecutableSubqueryForIn(in_function_subquery_node, planner_context->getQueryContext());
auto temporary_table_expression_node = executeSubqueryNode(subquery_to_execute,
planner_context->getMutableQueryContext(),
global_in_or_join_node.subquery_depth);

View File

@ -42,6 +42,17 @@ def test_cluster(start_cluster):
)
def test_global_in(start_cluster):
node1.query("CREATE TABLE u(uid Int16) ENGINE=Log as select 0");
assert set(
node1.query(
"""SELECT hostName(), * FROM clusterAllReplicas("one_shard_two_nodes", system.one) where dummy GLOBAL IN u"""
).splitlines()
) == {"node1\t0", "node2\t0"}
@pytest.mark.parametrize(
"cluster",
[