Backport #94095 to 25.12: Fix accuracy of uniqTheta when using UInt8 aggregation keys in parallel

robot-clickhouse · robot-clickhouse · commit 857ae92541ee · 2026-01-14T16:16:47.000Z
diff --git a/src/AggregateFunctions/ThetaSketchData.h b/src/AggregateFunctions/ThetaSketchData.h
@@ -20,7 +20,9 @@ template <typename Key>
 class ThetaSketchData : private boost::noncopyable
 {
 private:
+    /// Used for insertions
     std::unique_ptr<datasketches::update_theta_sketch> sk_update;
+    /// Used for merging
     std::unique_ptr<datasketches::theta_union> sk_union;
 
     datasketches::update_theta_sketch * getSkUpdate()
@@ -47,12 +49,26 @@ class ThetaSketchData : private boost::noncopyable
     void insertOriginal(std::string_view value)
     {
         getSkUpdate()->update(value.data(), value.size());
+        /// In case of optimization for u8 keys (see addBatchLookupTable()) it is possible to have few calls of insert() after merge(),
+        /// and we should update sk_union as well, note, that there should not be too many, so performance wise it should be OK
+        if (sk_union)
+        {
+            sk_union->update(*sk_update);
+            sk_update.reset(nullptr);
+        }
     }
 
     /// Note that `datasketches::update_theta_sketch.update` will do the hash again.
     void insert(Key value)
     {
         getSkUpdate()->update(value);
+        /// In case of optimization for u8 keys (see addBatchLookupTable()) it is possible to have few calls of insert() after merge(),
+        /// and we should update sk_union as well, note, that there should not be too many, so performance wise it should be OK
+        if (sk_union)
+        {
+            sk_union->update(*sk_update);
+            sk_update.reset(nullptr);
+        }
     }
 
     UInt64 size() const
diff --git a/tests/queries/0_stateless/01798_uniq_theta_sketch.oldanalyzer.reference b/tests/queries/0_stateless/01798_uniq_theta_sketch.oldanalyzer.reference
@@ -2,7 +2,7 @@ uniqTheta many agrs
 10	10	100	100	1000	1000
 17	10	10	100	100	610	610	766
 52	10	10	100	100	608	608	766
-5	10	10	100	100	608	608	765
+5	10	10	100	100	609	609	765
 9	10	10	100	100	608	608	765
 13	10	10	100	100	607	607	765
 46	10	10	100	100	607	607	765
diff --git a/tests/queries/0_stateless/01798_uniq_theta_sketch.reference b/tests/queries/0_stateless/01798_uniq_theta_sketch.reference
@@ -2,7 +2,7 @@ uniqTheta many agrs
 10	10	100	100	1000	1000
 17	10	10	100	100	610	610	766
 52	10	10	100	100	608	608	766
-5	10	10	100	100	608	608	765
+5	10	10	100	100	609	609	765
 9	10	10	100	100	608	608	765
 13	10	10	100	100	607	607	765
 46	10	10	100	100	607	607	765
diff --git a/tests/queries/0_stateless/03790_uniqTheta_error.reference b/tests/queries/0_stateless/03790_uniqTheta_error.reference
diff --git a/tests/queries/0_stateless/03790_uniqTheta_error.sql b/tests/queries/0_stateless/03790_uniqTheta_error.sql
@@ -0,0 +1,5 @@
+-- Tags: no-fasttest
+-- - no-fasttest -- compiled w/o datasketches
+
+-- Regression for very high error in uniqTheta() due to optimization for u8 keys
+select throwIf(stddevSampStable(theta)/avg(theta)>0.1) from (select number%3 key, uniqTheta(generateUUIDv4(number)) theta FROM numbers_mt(9160000) GROUP BY key with totals SETTINGS max_threads = 8) format Null;
