improve hash dispatch blocks

1) extend remote inserters to accelerate data transform
2) make advantage of low cardinality to improve groupby

Author: lgbo-ustc

src/Interpreters/ASTRewriters/CollectRequiredColumnsAction.cpp

@@ -147,6 +147,15 @@ void CollectRequiredColumnsAction::visit(const ASTIdentifier * ident_ast)
                         .alias_name = ident_ast->tryGetAlias(),
                         .type = col.type
                     };
+                    /*
+                    LOG_TRACE(
+                        &Poco::Logger::get("CollectRequiredColumnsAction"),
+                        "add ident @ {}, full name:{}, short name:{}, alias:{}.",
+                        *best_pos,
+                        ident_ast->name(),
+                        ident_ast->shortName(),
+                        ident_ast->tryGetAlias());
+                    */
                     final_result.required_columns[*best_pos].push_back(column_metadta);
                     found = true;
                     added_names.insert(ident_ast->name());

src/Interpreters/ASTRewriters/IdentifierQualiferRemoveAction.cpp

@@ -32,6 +32,8 @@ ASTs IdentifiterQualiferRemoveAction::collectChildren(const ASTPtr & ast)
 
 void IdentifiterQualiferRemoveAction::visit(const ASTPtr & ast)
 {
+    if (!ast)
+        return;
     if (const auto * function_ast = ast->as<ASTFunction>())
     {
         auto frame = frames.getTopFrame();

src/Interpreters/ASTRewriters/NestedJoinQueryRewriteAction.cpp

@@ -160,8 +160,8 @@ void NestedJoinQueryRewriteAction::visit(const ASTSelectQuery * select_ast)
         {
             auto ident = std::make_shared<ASTIdentifier>(col.splitedFullName());
             ident->alias = col.alias_name;
-            if (ident->alias.empty())
-                throw Exception(ErrorCodes::LOGICAL_ERROR, "Alias name is expected for {}", col.full_name);
+            //if (ident->alias.empty())
+            //    throw Exception(ErrorCodes::LOGICAL_ERROR, "Alias is expected for {}", col.full_name);
             nested_select_expr_list->children.emplace_back(ident);
         }
         nested_select_ast->setExpression(ASTSelectQuery::Expression::SELECT, nested_select_expr_list);

src/Interpreters/ASTRewriters/NestedJoinQueryRewriteAction.h

@@ -10,6 +10,7 @@
 #include <Parsers/IAST_fwd.h>
 #include <base/types.h>
 #include <Poco/Logger.h>
+#include <Interpreters/DatabaseAndTableWithAlias.h>
 namespace DB
 {
 /**

src/Interpreters/ASTRewriters/StageQueryDistributedAggregationRewriteAction.cpp

@@ -300,7 +300,12 @@ void StageQueryDistributedAggregationRewriteAction::visitSelectQueryWithAggregat
 
 void StageQueryDistributedAggregationRewriteAction::visitSelectQueryWithGroupby(const ASTSelectQuery * select_ast)
 {
+    LOG_TRACE(logger, "visitSelectQueryWithGroupby:{}", queryToString(*select_ast));
     auto frame = frames.getTopFrame();
+    auto tables = getDatabaseAndTablesWithColumns(getTableExpressions(*select_ast), context, true, true);
+    if (tables.size() != 1)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Tables size should be 1");
+    
     auto * rewrite_table_expr =  frame->children_results[0]->as<ASTTableExpression>();
     if (!rewrite_table_expr)
         throw Exception(ErrorCodes::LOGICAL_ERROR, "ASTTableExpression is expected. return query is : {}", queryToString(frame->children_results[0]));
@@ -313,21 +318,36 @@ void StageQueryDistributedAggregationRewriteAction::visitSelectQueryWithGroupby(
 
         rewrite_table_expr->subquery->as<ASTSubquery>()->setAlias(table_alias);
     }
-
-    auto tables = getDatabaseAndTablesWithColumns(getTableExpressions(*select_ast), context, true, true);
-    if (tables.size() != 1)
-        throw Exception(ErrorCodes::LOGICAL_ERROR, "Tables size should be 1");
+    
+    if (isAllRequiredColumnsLowCardinality(select_ast->groupBy(), tables))
+    {
+        CollectQueryStoragesAction collect_storage_action(context);
+        ASTDepthFirstVisitor<CollectQueryStoragesAction> collect_storage_visitor(collect_storage_action, frame->children_results[0]);
+        auto storages = collect_storage_visitor.visit();
+        /// If all columns used in the groupby clasue are low cardinality, do not shuffle the data and 
+        /// run the groupby in the two-phase way.
+        if (storages.size() > 1)
+            visitSelectQueryWithAggregation(select_ast);
+        return;
+    }
 
     CollectRequiredColumnsAction collect_columns_action(tables);
     ASTDepthFirstVisitor<CollectRequiredColumnsAction> collect_columns_visitor(collect_columns_action, select_ast->clone());
     auto required_columns = collect_columns_visitor.visit().required_columns;
 
-    auto insert_query = createShuffleInsert(
+    ASTPtr insert_query = createShuffleInsert(
         TableFunctionShuffleAggregation::name,
         rewrite_table_expr,
         ColumnWithDetailNameAndType::toNamesAndTypesList(required_columns[0]),
         select_ast->groupBy());
 
+    auto * insert_query_ptr = insert_query->as<ASTInsertQuery>();
+    auto * insert_select_ptr = insert_query_ptr->select->as<ASTSelectWithUnionQuery>()->list_of_selects->children[0]->as<ASTSelectQuery>();
+    IdentifiterQualiferRemoveAction remove_qualifier_action;
+    ASTDepthFirstVisitor<IdentifiterQualiferRemoveAction> remove_qualifier_visitor(remove_qualifier_action, select_ast->where());
+    auto where_expr = remove_qualifier_visitor.visit();
+    insert_select_ptr->setExpression(ASTSelectQuery::Expression::WHERE, std::move(where_expr));
+
     ASTs upstream_queries;
     frame->mergeChildrenUpstreamQueries();
     if (!frame->upstream_queries[0].empty())
@@ -341,6 +361,7 @@ void StageQueryDistributedAggregationRewriteAction::visitSelectQueryWithGroupby(
 
     frame->result_ast = select_ast->clone();
     auto * result_select_ast = frame->result_ast->as<ASTSelectQuery>();
+    result_select_ast->setExpression(ASTSelectQuery::Expression::WHERE, nullptr);
     ASTBuildUtil::updateSelectQueryTables(
         result_select_ast,
         ASTBuildUtil::createTablesInSelectQueryElement(table_function->as<ASTFunction>())->as<ASTTablesInSelectQueryElement>());
@@ -371,4 +392,23 @@ ASTPtr StageQueryDistributedAggregationRewriteAction::createShuffleInsert(
 
     return ASTBuildUtil::createTableFunctionInsertSelectQuery(table_function, ASTBuildUtil::wrapSelectQuery(select_query));
 }
+
+bool StageQueryDistributedAggregationRewriteAction::isAllRequiredColumnsLowCardinality(const ASTPtr & ast, const TablesWithColumns & tables)
+{
+    CollectRequiredColumnsAction collect_columns_action(tables);
+    ASTDepthFirstVisitor<CollectRequiredColumnsAction> collect_columns_visitor(collect_columns_action, ast);
+    auto required_columns = collect_columns_visitor.visit().required_columns;
+    for (auto & cols : required_columns)
+    {
+        for (auto & col : cols)
+        {
+            LOG_TRACE(logger, "check group by col. {} {}", col.full_name, col.type->getName());
+            if (!col.type->lowCardinality())
+                return false;
+        }
+    }
+    LOG_TRACE(logger, "all columns are locaCardinality");
+    return true;
+
+}
 }

src/Interpreters/ASTRewriters/StageQueryDistributedAggregationRewriteAction.h

@@ -7,6 +7,7 @@
 #include <Parsers/ASTSubquery.h>
 #include <Parsers/IAST.h>
 #include <Parsers/IAST_fwd.h>
+#include <Interpreters/DatabaseAndTableWithAlias.h>
 
 namespace DB
 {
@@ -67,5 +68,7 @@ class StageQueryDistributedAggregationRewriteAction : public EmptyASTDepthFirstV
 
     ASTPtr createShuffleInsert(
         const String & table_function_name, ASTTableExpression * table_expr, const NamesAndTypesList & table_desc, ASTPtr groupby_clause);
+
+    bool isAllRequiredColumnsLowCardinality(const ASTPtr & ast, const TablesWithColumns & tables);
 };
 }

src/Processors/Transforms/StageQueryTransform.cpp

@@ -226,15 +226,8 @@ ParallelStageBlockIOsTransform::ParallelStageBlockIOsTransform(
 
 ParallelStageBlockIOsTransform::~ParallelStageBlockIOsTransform()
 {
-#if 0
-    for (auto & task : background_tasks)
-    {
-        task->deactivate();
-    }
-#else
     if (thread_pool)
         thread_pool->wait();
-#endif
     LOG_TRACE(logger, "run query({}) in elapsedMilliseconds:{}", queryToString(output_block_io.query), elapsed);
 }
 
@@ -305,20 +298,11 @@ void ParallelStageBlockIOsTransform::startBackgroundTasks()
             queryToString(block_io.query),
             task_watch.elapsedMilliseconds());
     };
-#if 0
-    auto & thread_pool = context->getSchedulePool();
-    for (auto & block_io : input_block_ios)
-    {
-        background_tasks.emplace_back(thread_pool.createTask("BackgroundBlockIOTask", [build_task, &block_io](){ build_task(block_io);}));
-        background_tasks.back()->activateAndSchedule();
-    }
-#else
     thread_pool = std::make_unique<ThreadPool>(input_block_ios.size());
     for (auto & block : input_block_ios)
     {
         thread_pool->scheduleOrThrowOnError([&]() { build_task(block); });
     }
-#endif
     has_start_background_tasks = true;
 }
 }

src/Storages/DistributedShuffle/ShuffleBlockTable.cpp

@@ -30,11 +30,18 @@ void ShuffleBlockTable::addChunk(Chunk && chunk)
 {
     if (chunk.hasRows()) [[likely]]
     {
-        if (is_sink_finished)[[unlikely]]
-            throw Exception(ErrorCodes::LOGICAL_ERROR, "Try in insert into a sink finished table({}.{})", session_id, table_id);
-        std::unique_lock lock(mutex);
-        rows += chunk.getNumRows();
-        chunks.emplace_back(std::move(chunk));
+        {
+            std::unique_lock lock(mutex);
+            if (is_sink_finished) [[unlikely]]
+                throw Exception(ErrorCodes::LOGICAL_ERROR, "Try in insert into a sink finished table({}.{})", session_id, table_id);
+            while (remained_rows > max_rows_limit)
+                wait_consume_data.wait(lock);
+            rows += chunk.getNumRows();
+            remained_rows += chunk.getNumRows();
+            if (remained_rows > max_rows)
+                max_rows = remained_rows;
+            chunks.emplace_back(std::move(chunk));
+        }
         wait_more_data.notify_one();
     }
     else
@@ -46,30 +53,34 @@ void ShuffleBlockTable::addChunk(Chunk && chunk)
 
 Chunk ShuffleBlockTable::popChunk()
 {
-    std::unique_lock lock(mutex);
-    while (chunks.empty())
+    Chunk res;
     {
-        if (!is_sink_finished)
-        {
-            wait_more_data.wait(lock, [&] { return is_sink_finished || !chunks.empty(); });
-        }
-        else
+        std::unique_lock lock(mutex);
+        while (chunks.empty())
         {
-            break;
+            if (!is_sink_finished)
+            {
+                wait_more_data.wait(lock, [&] { return is_sink_finished || !chunks.empty(); });
+            }
+            else
+            {
+                break;
+            }
         }
-    }
-    //LOG_TRACE(logger, "{}.{} popChunk. isSinkFinished()={}, chunks.size()={}", session_id, table_id, is_sink_finished, chunks.size());
+        //LOG_TRACE(logger, "{}.{} popChunk. isSinkFinished()={}, chunks.size()={}", session_id, table_id, is_sink_finished, chunks.size());
 
-    Chunk res;
-    if (!chunks.empty()) [[likely]]
-    {
-        res.swap(chunks.front());
-        chunks.pop_front();
-        if (unlikely(!res.hasRows()))
+        if (!chunks.empty()) [[likely]]
         {
-            throw Exception(ErrorCodes::LOGICAL_ERROR, "Chunk should not be empty. table({}.{})", session_id, table_id);
+            res.swap(chunks.front());
+            remained_rows -= res.getNumRows();
+            chunks.pop_front();
+            if (unlikely(!res.hasRows()))
+            {
+                throw Exception(ErrorCodes::LOGICAL_ERROR, "Chunk should not be empty. table({}.{})", session_id, table_id);
+            }
         }
     }
+    wait_consume_data.notify_all();
     return res;
 }
 

src/Storages/DistributedShuffle/ShuffleBlockTable.h

@@ -18,19 +18,6 @@
 namespace DB
 {
 ///
-/// How to clear all the data when a query session has finished ?
-/// The following measures were taken at current
-/// 1）Chunks in ShuffleBlockTable are read only once, so we use popChunkWithoutMutex() for loading a chunk.
-///   That ensures that all chunks are released after the loading finish.
-/// 2) When ShuffleBlockTable becomes empty, it will call ShuffleBlockSession::releaseTable() to
-///   release it-self.
-/// 3) When ShuffleBlockSession becomes empty, it will call ShuffleBlockTableManager::tryCloseSession() to
-///   release it-self.
-/// All above will ensure all datas are released in normal processing. But more need be considered, exceptions could
-/// happen during the processing which make the release actions not be called. Some measures may be token.
-/// 1) In TCPHandler, catch all exceptions , and make a session releasing action on all nodes
-/// 2) All sessions have a max TTL, make background routine to check timeout sessions and clear them.
-///
 
 class ShuffleBlockTable
 {
@@ -47,7 +34,7 @@ class ShuffleBlockTable
 
     ~ShuffleBlockTable()
     {
-        LOG_TRACE(logger, "close table {}.{}", session_id, table_id);
+        LOG_TRACE(logger, "close table {}.{}. rows:{}. max_rows:{}", session_id, table_id, rows, max_rows);
     }
 
     inline const Block & getHeader() const
@@ -79,8 +66,12 @@ class ShuffleBlockTable
     std::atomic<bool> is_sink_finished = false;
     std::list<Chunk> chunks;
     std::condition_variable wait_more_data;
+    std::condition_variable wait_consume_data;
     Poco::Logger * logger = &Poco::Logger::get("ShuffleBlockTable");
     size_t rows = 0;
+    size_t remained_rows = 0;
+    size_t max_rows = 0;
+    const static size_t max_rows_limit = 20000000;
 };
 using ShuffleBlockTablePtr = std::shared_ptr<ShuffleBlockTable>;