Merge pull request ClickHouse#78703 from ClickHouse/parallel_uniq_exact_merging_on_initiator

Parallel `uniqExact` merging on initiator

Author: nickitat

src/AggregateFunctions/Combinators/AggregateFunctionIf.h

@@ -183,9 +183,11 @@ class AggregateFunctionIf final : public IAggregateFunctionHelper<AggregateFunct
         AggregateDataPtr * places,
         size_t place_offset,
         const AggregateDataPtr * rhs,
+        ThreadPool & thread_pool,
+        std::atomic<bool> & is_cancelled,
         Arena * arena) const override
     {
-        nested_func->mergeBatch(row_begin, row_end, places, place_offset, rhs, arena);
+        nested_func->mergeBatch(row_begin, row_end, places, place_offset, rhs, thread_pool, is_cancelled, arena);
     }
 
     void serialize(ConstAggregateDataPtr __restrict place, WriteBuffer & buf, std::optional<size_t> version) const override

src/AggregateFunctions/Combinators/AggregateFunctionOrFill.h

@@ -223,9 +223,11 @@ class AggregateFunctionOrFill final : public IAggregateFunctionHelper<AggregateF
         AggregateDataPtr * places,
         size_t place_offset,
         const AggregateDataPtr * rhs,
+        ThreadPool & thread_pool,
+        std::atomic<bool> & is_cancelled,
         Arena * arena) const override
     {
-        nested_function->mergeBatch(row_begin, row_end, places, place_offset, rhs, arena);
+        nested_function->mergeBatch(row_begin, row_end, places, place_offset, rhs, thread_pool, is_cancelled, arena);
         for (size_t i = row_begin; i < row_end; ++i)
             (places[i] + place_offset)[size_of_data] |= rhs[i][size_of_data];
     }

src/AggregateFunctions/IAggregateFunction.h

@@ -245,6 +245,8 @@ class IAggregateFunction : public std::enable_shared_from_this<IAggregateFunctio
         AggregateDataPtr * places,
         size_t place_offset,
         const AggregateDataPtr * rhs,
+        ThreadPool & thread_pool,
+        std::atomic<bool> & is_cancelled,
         Arena * arena) const = 0;
 
     /** The same for single place.
@@ -509,11 +511,20 @@ class IAggregateFunctionHelper : public IAggregateFunction
         AggregateDataPtr * places,
         size_t place_offset,
         const AggregateDataPtr * rhs,
+        ThreadPool & thread_pool,
+        std::atomic<bool> & is_cancelled,
         Arena * arena) const override
     {
         for (size_t i = row_begin; i < row_end; ++i)
+        {
             if (places[i])
-                static_cast<const Derived *>(this)->merge(places[i] + place_offset, rhs[i], arena);
+            {
+                if constexpr (Derived::parallelizeMergeWithKey())
+                    static_cast<const Derived *>(this)->merge(places[i] + place_offset, rhs[i], thread_pool, is_cancelled, arena);
+                else
+                    static_cast<const Derived *>(this)->merge(places[i] + place_offset, rhs[i], arena);
+            }
+        }
     }
 
     void mergeAndDestroyBatch(AggregateDataPtr * dst_places, AggregateDataPtr * rhs_places, size_t size, size_t offset, ThreadPool & thread_pool, std::atomic<bool> & is_cancelled, Arena * arena) const override

src/AggregateFunctions/UniqExactSet.h

@@ -1,12 +1,11 @@
 #pragma once
 
-#include <exception>
 #include <Common/CurrentThread.h>
 #include <Common/HashTable/HashSet.h>
 #include <Common/ThreadPool.h>
 #include <Common/scope_guard_safe.h>
 #include <Common/setThreadName.h>
-
+#include <Common/threadPoolCallbackRunner.h>
 
 namespace DB
 {
@@ -94,13 +93,6 @@ class UniqExactSet
 
     auto merge(const UniqExactSet & other, ThreadPool * thread_pool = nullptr, std::atomic<bool> * is_cancelled = nullptr)
     {
-        /// If the size is large, we may convert the singleLevelHash to twoLevelHash and merge in parallel.
-        if (other.size() > 40000)
-        {
-            if (isSingleLevel())
-                convertToTwoLevel();
-        }
-
         if (isSingleLevel() && other.isTwoLevel())
             convertToTwoLevel();
 
@@ -122,17 +114,16 @@ class UniqExactSet
             }
             else
             {
+                ThreadPoolCallbackRunnerLocal<void> runner(*thread_pool, "UniqExactMerger");
                 try
                 {
                     auto next_bucket_to_merge = std::make_shared<std::atomic_uint32_t>(0);
 
                     auto thread_func = [&lhs, &rhs, next_bucket_to_merge, is_cancelled, thread_group = CurrentThread::getGroup()]()
                     {
-                        ThreadGroupSwitcher switcher(thread_group, "UniqExactMerger");
-
                         while (true)
                         {
-                            if (is_cancelled->load(std::memory_order_seq_cst))
+                            if (is_cancelled->load())
                                 return;
 
                             const auto bucket = next_bucket_to_merge->fetch_add(1);
@@ -143,13 +134,13 @@ class UniqExactSet
                     };
 
                     for (size_t i = 0; i < std::min<size_t>(thread_pool->getMaxThreads(), rhs.NUM_BUCKETS); ++i)
-                        thread_pool->scheduleOrThrowOnError(thread_func);
-                    thread_pool->wait();
+                        runner(thread_func, Priority{});
+                    runner.waitForAllToFinishAndRethrowFirstError();
                 }
                 catch (...)
                 {
-                    thread_pool->wait();
-                    throw;
+                    is_cancelled->store(true);
+                    runner.waitForAllToFinishAndRethrowFirstError();
                 }
             }
         }