From f5cfe597f2ae562ea8e4e7b89e24e0e769408d91 Mon Sep 17 00:00:00 2001 From: Maksim Kita Date: Tue, 23 Apr 2024 13:00:02 +0300 Subject: [PATCH] Fixed tests --- src/Analyzer/Utils.cpp | 30 +++++-------------- src/Analyzer/Utils.h | 9 ++++-- src/Processors/Sources/RecursiveCTESource.cpp | 25 +++++++++------- src/Storages/buildQueryTreeForShard.cpp | 2 +- 4 files changed, 29 insertions(+), 37 deletions(-) diff --git a/src/Analyzer/Utils.cpp b/src/Analyzer/Utils.cpp index d813ddf345b..3c3489681f6 100644 --- a/src/Analyzer/Utils.cpp +++ b/src/Analyzer/Utils.cpp @@ -841,28 +841,6 @@ QueryTreeNodePtr getExpressionSource(const QueryTreeNodePtr & node) return source; } -QueryTreeNodePtr buildSubqueryToReadColumnsFromTableExpression(QueryTreeNodePtr table_node, const ContextPtr & context) -{ - const auto & storage_snapshot = table_node->as()->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(); - 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(column, table_node)); - projection_columns.emplace_back(column.name, column.type); - } - auto subquery_for_table = std::make_shared(Context::createCopy(context)); - subquery_for_table->setIsSubquery(true); - subquery_for_table->getProjectionNode() = std::move(column_nodes_to_select); - subquery_for_table->getJoinTree() = std::move(table_node); - subquery_for_table->resolveProjectionColumns(std::move(projection_columns)); - return subquery_for_table; -} - /** There are no limits on the maximum size of the result for the subquery. * Since the result of the query is not the result of the entire query. */ @@ -938,4 +916,12 @@ QueryTreeNodePtr buildSubqueryToReadColumnsFromTableExpression(const NamesAndTyp return buildSubqueryToReadColumnsFromTableExpression(columns, table_expression, context_copy); } +QueryTreeNodePtr buildSubqueryToReadColumnsFromTableExpression(const QueryTreeNodePtr & table_node, const ContextPtr & context) +{ + const auto & storage_snapshot = table_node->as()->getStorageSnapshot(); + auto columns_to_select_list = storage_snapshot->getColumns(GetColumnsOptions(GetColumnsOptions::Ordinary)); + NamesAndTypes columns_to_select(columns_to_select_list.begin(), columns_to_select_list.end()); + return buildSubqueryToReadColumnsFromTableExpression(columns_to_select, table_node, context); +} + } diff --git a/src/Analyzer/Utils.h b/src/Analyzer/Utils.h index 77cc7389598..f64b724abeb 100644 --- a/src/Analyzer/Utils.h +++ b/src/Analyzer/Utils.h @@ -115,9 +115,6 @@ QueryTreeNodePtr createCastFunction(QueryTreeNodePtr node, DataTypePtr result_ty /// Checks that node has only one source and returns it QueryTreeNodePtr getExpressionSource(const QueryTreeNodePtr & node); -/// Build subquery which we execute for `IN table` function. -QueryTreeNodePtr buildSubqueryToReadColumnsFromTableExpression(QueryTreeNodePtr table_node, const ContextPtr & context); - /// Update mutable context for subquery execution void updateContextForSubqueryExecution(ContextMutablePtr & mutable_context); @@ -149,4 +146,10 @@ QueryTreeNodePtr buildSubqueryToReadColumnsFromTableExpression(const NamesAndTyp const QueryTreeNodePtr & table_expression, const ContextPtr & context); +/** Build subquery to read all columns from table expression. + * Specified context will be copied and used as query context. + */ +QueryTreeNodePtr buildSubqueryToReadColumnsFromTableExpression(const QueryTreeNodePtr & table_node, const ContextPtr & context); + + } diff --git a/src/Processors/Sources/RecursiveCTESource.cpp b/src/Processors/Sources/RecursiveCTESource.cpp index 55fd192896e..b94cb188086 100644 --- a/src/Processors/Sources/RecursiveCTESource.cpp +++ b/src/Processors/Sources/RecursiveCTESource.cpp @@ -15,6 +15,7 @@ #include #include +#include namespace DB { @@ -135,6 +136,7 @@ public: if (read_rows_during_recursive_step == 0) { finished = true; + truncateTemporaryTable(intermediate_temporary_table_storage); continue; } @@ -143,12 +145,7 @@ public: for (auto & recursive_table_node : recursive_table_nodes) recursive_table_node->updateStorage(intermediate_temporary_table_storage, recursive_query_context); - /// TODO: Support proper locking - TableExclusiveLockHolder table_exclusive_lock; - working_temporary_table_storage->truncate({}, - working_temporary_table_storage->getInMemoryMetadataPtr(), - recursive_query_context, - table_exclusive_lock); + truncateTemporaryTable(working_temporary_table_storage); std::swap(intermediate_temporary_table_storage, working_temporary_table_storage); } @@ -193,11 +190,7 @@ private: return std::make_shared(input_header, convert_to_temporary_tables_header_actions); }); - auto squashing_transform = std::make_shared(pipeline_builder.getHeader(), - recursive_subquery_settings.min_external_table_block_size_rows, - recursive_subquery_settings.min_external_table_block_size_bytes); - pipeline_builder.resize(1); - pipeline_builder.addTransform(std::move(squashing_transform)); + /// TODO: Support squashing transform auto intermediate_temporary_table_storage_sink = intermediate_temporary_table_storage->write( {}, @@ -214,6 +207,16 @@ private: executor.emplace(pipeline); } + void truncateTemporaryTable(StoragePtr & temporary_table) + { + /// TODO: Support proper locking + TableExclusiveLockHolder table_exclusive_lock; + temporary_table->truncate({}, + temporary_table->getInMemoryMetadataPtr(), + recursive_query_context, + table_exclusive_lock); + } + Block header; QueryTreeNodePtr recursive_cte_union_node; std::vector recursive_table_nodes; diff --git a/src/Storages/buildQueryTreeForShard.cpp b/src/Storages/buildQueryTreeForShard.cpp index 5bbdbe487b0..4f655f9b5e8 100644 --- a/src/Storages/buildQueryTreeForShard.cpp +++ b/src/Storages/buildQueryTreeForShard.cpp @@ -366,7 +366,7 @@ QueryTreeNodePtr buildQueryTreeForShard(const PlannerContextPtr & planner_contex auto subquery_to_execute = in_function_subquery_node; if (subquery_to_execute->as()) - subquery_to_execute = buildSubqueryToReadColumnsFromTableExpression(std::move(subquery_to_execute), planner_context->getQueryContext()); + subquery_to_execute = buildSubqueryToReadColumnsFromTableExpression(subquery_to_execute, planner_context->getQueryContext()); auto temporary_table_expression_node = executeSubqueryNode(subquery_to_execute, planner_context->getMutableQueryContext(),