impl

nickitat · nickitat · commit d7e3defaf643 · 2025-04-05T19:26:49.000Z
diff --git a/src/AggregateFunctions/AggregateFunctionUniq.h b/src/AggregateFunctions/AggregateFunctionUniq.h
@@ -276,13 +276,13 @@ struct Adder
 {
     /// We have to introduce this template parameter (and a bunch of ugly code dealing with it), because we cannot
     /// add runtime branches in whatever_hash_set::insert - it will immediately pop up in the perf top.
-    template <bool use_single_level_hash_table = true>
+    template <SetLevelHint hint = SetLevelHint::unknown>
     static void ALWAYS_INLINE add(Data & data, const IColumn ** columns, size_t num_args, size_t row_num)
     {
         if constexpr (Data::is_variadic)
         {
             if constexpr (IsUniqExactSet<typename Data::Set>::value)
-                data.set.template insert<T, use_single_level_hash_table>(
+                data.set.template insert<T, hint>(
                     UniqVariadicHash<Data::is_exact, Data::argument_is_tuple>::apply(num_args, columns, row_num));
             else
                 data.set.insert(T{UniqVariadicHash<Data::is_exact, Data::argument_is_tuple>::apply(num_args, columns, row_num)});
@@ -316,12 +316,11 @@ struct Adder
                 hash.update(value.data, value.size);
                 const auto key = hash.get128();
 
-                data.set.template insert<const UInt128 &, use_single_level_hash_table>(key);
+                data.set.template insert<const UInt128 &, hint>(key);
             }
             else
             {
-                data.set.template insert<const T &, use_single_level_hash_table>(
-                    assert_cast<const ColumnVector<T> &>(column).getData()[row_num]);
+                data.set.template insert<const T &, hint>(assert_cast<const ColumnVector<T> &>(column).getData()[row_num]);
             }
         }
 #if USE_DATASKETCHES
@@ -341,9 +340,9 @@ struct Adder
             use_single_level_hash_table = data.set.isSingleLevel();
 
         if (use_single_level_hash_table)
-            addImpl<true>(data, columns, num_args, row_begin, row_end, flags, null_map);
+            addImpl<SetLevelHint::singleLevel>(data, columns, num_args, row_begin, row_end, flags, null_map);
         else
-            addImpl<false>(data, columns, num_args, row_begin, row_end, flags, null_map);
+            addImpl<SetLevelHint::twoLevel>(data, columns, num_args, row_begin, row_end, flags, null_map);
 
         if constexpr (Data::is_able_to_parallelize_merge)
         {
@@ -353,22 +352,28 @@ struct Adder
     }
 
 private:
-    template <bool use_single_level_hash_table>
-    static void ALWAYS_INLINE
-    addImpl(Data & data, const IColumn ** columns, size_t num_args, size_t row_begin, size_t row_end, const char8_t * flags, const UInt8 * null_map)
+    template <SetLevelHint hint>
+    static void ALWAYS_INLINE addImpl(
+        Data & data,
+        const IColumn ** columns,
+        size_t num_args,
+        size_t row_begin,
+        size_t row_end,
+        const char8_t * flags,
+        const UInt8 * null_map)
     {
         if (!flags)
         {
             if (!null_map)
             {
                 for (size_t row = row_begin; row < row_end; ++row)
-                    add<use_single_level_hash_table>(data, columns, num_args, row);
+                    add<hint>(data, columns, num_args, row);
             }
             else
             {
                 for (size_t row = row_begin; row < row_end; ++row)
                     if (!null_map[row])
-                        add<use_single_level_hash_table>(data, columns, num_args, row);
+                        add<hint>(data, columns, num_args, row);
             }
         }
         else
@@ -377,13 +382,13 @@ struct Adder
             {
                 for (size_t row = row_begin; row < row_end; ++row)
                     if (flags[row])
-                        add<use_single_level_hash_table>(data, columns, num_args, row);
+                        add<hint>(data, columns, num_args, row);
             }
             else
             {
                 for (size_t row = row_begin; row < row_end; ++row)
                     if (!null_map[row] && flags[row])
-                        add<use_single_level_hash_table>(data, columns, num_args, row);
+                        add<hint>(data, columns, num_args, row);
             }
         }
     }
diff --git a/src/AggregateFunctions/UniqExactSet.h b/src/AggregateFunctions/UniqExactSet.h
@@ -16,6 +16,13 @@ namespace ErrorCodes
 extern const int TOO_LARGE_ARRAY_SIZE;
 }
 
+enum class SetLevelHint
+{
+    singleLevel,
+    twoLevel,
+    unknown,
+};
+
 template <typename SingleLevelSet, typename TwoLevelSet>
 class UniqExactSet
 {
@@ -25,19 +32,30 @@ class UniqExactSet
 public:
     using value_type = typename SingleLevelSet::value_type;
 
-    template <typename Arg, bool use_single_level_hash_table = true>
+    template <typename Arg, SetLevelHint hint>
     auto ALWAYS_INLINE insert(Arg && arg)
     {
-        if (isSingleLevel())
+        if constexpr (hint == SetLevelHint::singleLevel)
         {
-            auto && [_, inserted] = asSingleLevel().insert(std::forward<Arg>(arg));
-            if (inserted && worthConvertingToTwoLevel(asSingleLevel().size()))
-                convertToTwoLevel();
+            asSingleLevel().insert(std::forward<Arg>(arg));
         }
-        else
+        else if constexpr (hint == SetLevelHint::twoLevel)
         {
             asTwoLevel().insert(std::forward<Arg>(arg));
         }
+        else
+        {
+            if (isSingleLevel())
+            {
+                auto && [_, inserted] = asSingleLevel().insert(std::forward<Arg>(arg));
+                if (inserted && worthConvertingToTwoLevel(asSingleLevel().size()))
+                    convertToTwoLevel();
+            }
+            else
+            {
+                asTwoLevel().insert(std::forward<Arg>(arg));
+            }
+        }
     }
 
     /// In merge, if one of the lhs and rhs is twolevelset and the other is singlelevelset, then the singlelevelset will need to convertToTwoLevel().
diff --git a/tests/performance/uniqExact_with_key_parallel_merging.xml b/tests/performance/uniqExact_with_key_parallel_merging.xml
@@ -0,0 +1,30 @@
+<test>
+  <create_query>
+    CREATE TABLE test (a UInt64, b UInt64) ENGINE = MergeTree() ORDER BY () AS SELECT number, intDiv(number, 10) FROM numbers_mt(1e8);
+  </create_query>
+
+  <substitutions>
+    <substitution>
+        <name>threads</name>
+        <values>
+            <value>1</value>
+            <value>2</value>
+            <value>4</value>
+            <value>8</value>
+        </values>
+    </substitution>
+  </substitutions>
+
+  <query>
+    SELECT
+      multiIf(a >= 100000 AND a &lt; 10000000, 0, 1) AS index,
+      COUNT(DISTINCT b) AS total
+    FROM test
+    GROUP BY 1
+    SETTINGS max_threads = {threads};
+  </query>
+
+  <drop_query>
+    DROP TABLE test;
+  </drop_query>
+</test>

Original file line number	Diff line number	Diff line change
`@@ -276,13 +276,13 @@ struct Adder`
`276`	`276`	`{`
`277`	`277`	`/// We have to introduce this template parameter (and a bunch of ugly code dealing with it), because we cannot`
`278`	`278`	`/// add runtime branches in whatever_hash_set::insert - it will immediately pop up in the perf top.`
`279`		`- template <bool use_single_level_hash_table = true>`
	`279`	`+ template <SetLevelHint hint = SetLevelHint::unknown>`
`280`	`280`	`static void ALWAYS_INLINE add(Data & data, const IColumn ** columns, size_t num_args, size_t row_num)`
`281`	`281`	`{`
`282`	`282`	`if constexpr (Data::is_variadic)`
`283`	`283`	`{`
`284`	`284`	`if constexpr (IsUniqExactSet<typename Data::Set>::value)`
`285`		`- data.set.template insert<T, use_single_level_hash_table>(`
	`285`	`+ data.set.template insert<T, hint>(`
`286`	`286`	`UniqVariadicHash<Data::is_exact, Data::argument_is_tuple>::apply(num_args, columns, row_num));`
`287`	`287`	`else`
`288`	`288`	`data.set.insert(T{UniqVariadicHash<Data::is_exact, Data::argument_is_tuple>::apply(num_args, columns, row_num)});`
`@@ -316,12 +316,11 @@ struct Adder`
`316`	`316`	`hash.update(value.data, value.size);`
`317`	`317`	`const auto key = hash.get128();`
`318`	`318`
`319`		`- data.set.template insert<const UInt128 &, use_single_level_hash_table>(key);`
	`319`	`+ data.set.template insert<const UInt128 &, hint>(key);`
`320`	`320`	`}`
`321`	`321`	`else`
`322`	`322`	`{`
`323`		`- data.set.template insert<const T &, use_single_level_hash_table>(`
`324`		`- assert_cast<const ColumnVector<T> &>(column).getData()[row_num]);`
	`323`	`+ data.set.template insert<const T &, hint>(assert_cast<const ColumnVector<T> &>(column).getData()[row_num]);`
`325`	`324`	`}`
`326`	`325`	`}`
`327`	`326`	`#if USE_DATASKETCHES`
`@@ -341,9 +340,9 @@ struct Adder`
`341`	`340`	`use_single_level_hash_table = data.set.isSingleLevel();`
`342`	`341`
`343`	`342`	`if (use_single_level_hash_table)`
`344`		`- addImpl<true>(data, columns, num_args, row_begin, row_end, flags, null_map);`
	`343`	`+ addImpl<SetLevelHint::singleLevel>(data, columns, num_args, row_begin, row_end, flags, null_map);`
`345`	`344`	`else`
`346`		`- addImpl<false>(data, columns, num_args, row_begin, row_end, flags, null_map);`
	`345`	`+ addImpl<SetLevelHint::twoLevel>(data, columns, num_args, row_begin, row_end, flags, null_map);`
`347`	`346`
`348`	`347`	`if constexpr (Data::is_able_to_parallelize_merge)`
`349`	`348`	`{`
`@@ -353,22 +352,28 @@ struct Adder`
`353`	`352`	`}`
`354`	`353`
`355`	`354`	`private:`
`356`		`- template <bool use_single_level_hash_table>`
`357`		`- static void ALWAYS_INLINE`
`358`		`- addImpl(Data & data, const IColumn ** columns, size_t num_args, size_t row_begin, size_t row_end, const char8_t * flags, const UInt8 * null_map)`
	`355`	`+ template <SetLevelHint hint>`
	`356`	`+ static void ALWAYS_INLINE addImpl(`
	`357`	`+ Data & data,`
	`358`	`+ const IColumn ** columns,`
	`359`	`+ size_t num_args,`
	`360`	`+ size_t row_begin,`
	`361`	`+ size_t row_end,`
	`362`	`+ const char8_t * flags,`
	`363`	`+ const UInt8 * null_map)`
`359`	`364`	`{`
`360`	`365`	`if (!flags)`
`361`	`366`	`{`
`362`	`367`	`if (!null_map)`
`363`	`368`	`{`
`364`	`369`	`for (size_t row = row_begin; row < row_end; ++row)`
`365`		`- add<use_single_level_hash_table>(data, columns, num_args, row);`
	`370`	`+ add<hint>(data, columns, num_args, row);`
`366`	`371`	`}`
`367`	`372`	`else`
`368`	`373`	`{`
`369`	`374`	`for (size_t row = row_begin; row < row_end; ++row)`
`370`	`375`	`if (!null_map[row])`
`371`		`- add<use_single_level_hash_table>(data, columns, num_args, row);`
	`376`	`+ add<hint>(data, columns, num_args, row);`
`372`	`377`	`}`
`373`	`378`	`}`
`374`	`379`	`else`
`@@ -377,13 +382,13 @@ struct Adder`
`377`	`382`	`{`
`378`	`383`	`for (size_t row = row_begin; row < row_end; ++row)`
`379`	`384`	`if (flags[row])`
`380`		`- add<use_single_level_hash_table>(data, columns, num_args, row);`
	`385`	`+ add<hint>(data, columns, num_args, row);`
`381`	`386`	`}`
`382`	`387`	`else`
`383`	`388`	`{`
`384`	`389`	`for (size_t row = row_begin; row < row_end; ++row)`
`385`	`390`	`if (!null_map[row] && flags[row])`
`386`		`- add<use_single_level_hash_table>(data, columns, num_args, row);`
	`391`	`+ add<hint>(data, columns, num_args, row);`
`387`	`392`	`}`
`388`	`393`	`}`
`389`	`394`	`}`