Merge pull request ClickHouse#88224 from ClickHouse/backport/25.8/86560

Backport ClickHouse#86560 to 25.8: Fix stack overflow in quantileDD merge

Author: clickhouse-gh[bot]

src/AggregateFunctions/DDSketch.h

@@ -1,49 +1,49 @@
 #pragma once
 
-#include <memory> // for std::unique_ptr
-#include <cmath>
-#include <stdexcept>
 #include <limits>
-#include <iostream>
+#include <memory> // for std::unique_ptr
 #include <base/types.h>
 
-#include <IO/ReadBuffer.h>
-#include <IO/WriteBuffer.h>
-
+#include <AggregateFunctions/DDSketch/DDSketchEncoding.h>
 #include <AggregateFunctions/DDSketch/Mapping.h>
 #include <AggregateFunctions/DDSketch/Store.h>
-#include <AggregateFunctions/DDSketch/DDSketchEncoding.h>
+#include <IO/ReadBuffer.h>
+#include <IO/ReadHelpers.h>
+#include <IO/WriteBuffer.h>
+#include <IO/WriteHelpers.h>
+#include <Common/Exception.h>
 
 namespace DB
 {
 
 namespace ErrorCodes
 {
-    extern const int BAD_ARGUMENTS;
-    extern const int INCORRECT_DATA;
+extern const int BAD_ARGUMENTS;
+extern const int INCORRECT_DATA;
 }
 
 class DDSketchDenseLogarithmic
 {
 public:
     explicit DDSketchDenseLogarithmic(Float64 relative_accuracy = 0.01)
-        : mapping(std::make_unique<DDSketchLogarithmicMapping>(relative_accuracy)),
-          store(std::make_unique<DDSketchDenseStore>()),
-          negative_store(std::make_unique<DDSketchDenseStore>()),
-          zero_count(0.0),
-          count(0.0)
+        : mapping(std::make_unique<DDSketchLogarithmicMapping>(relative_accuracy))
+        , store(std::make_unique<DDSketchDenseStore>())
+        , negative_store(std::make_unique<DDSketchDenseStore>())
+        , zero_count(0.0)
+        , count(0.0)
     {
     }
 
-    DDSketchDenseLogarithmic(std::unique_ptr<DDSketchLogarithmicMapping> mapping_,
-             std::unique_ptr<DDSketchDenseStore> store_,
-             std::unique_ptr<DDSketchDenseStore> negative_store_,
-             Float64 zero_count_)
-        : mapping(std::move(mapping_)),
-          store(std::move(store_)),
-          negative_store(std::move(negative_store_)),
-          zero_count(zero_count_),
-          count(store->count + negative_store->count + zero_count_)
+    DDSketchDenseLogarithmic(
+        std::unique_ptr<DDSketchLogarithmicMapping> mapping_,
+        std::unique_ptr<DDSketchDenseStore> store_,
+        std::unique_ptr<DDSketchDenseStore> negative_store_,
+        Float64 zero_count_)
+        : mapping(std::move(mapping_))
+        , store(std::move(store_))
+        , negative_store(std::move(negative_store_))
+        , zero_count(zero_count_)
+        , count(store->count + negative_store->count + zero_count_)
     {
     }
 
@@ -97,7 +97,11 @@ class DDSketchDenseLogarithmic
         return quantile_value;
     }
 
-    void copy(const DDSketchDenseLogarithmic& other)
+    Float64 getGamma() const { return mapping->getGamma(); }
+
+    Float64 getCount() const { return count; }
+
+    void copy(const DDSketchDenseLogarithmic & other)
     {
         Float64 rel_acc = (other.mapping->getGamma() - 1) / (other.mapping->getGamma() + 1);
         mapping = std::make_unique<DDSketchLogarithmicMapping>(rel_acc);
@@ -109,9 +113,9 @@ class DDSketchDenseLogarithmic
         count = other.count;
     }
 
-    void merge(const DDSketchDenseLogarithmic& other)
+    void merge(const DDSketchDenseLogarithmic & other)
     {
-        if (mapping->getGamma() != other.mapping->getGamma())
+        if (*mapping != *other.mapping)
         {
             // modify the one with higher precision to match the one with lower precision
             if (mapping->getGamma() > other.mapping->getGamma())
@@ -147,7 +151,7 @@ class DDSketchDenseLogarithmic
 
     /// NOLINTBEGIN(readability-static-accessed-through-instance)
 
-    void serialize(WriteBuffer& buf) const
+    void serialize(WriteBuffer & buf) const
     {
         // Write the mapping
         writeBinary(enc.FlagIndexMappingBaseLogarithmic.byte, buf);
@@ -165,7 +169,7 @@ class DDSketchDenseLogarithmic
         writeBinary(zero_count, buf);
     }
 
-    void deserialize(ReadBuffer& buf)
+    void deserialize(ReadBuffer & buf)
     {
         // Read the mapping
         UInt8 flag = 0;
@@ -219,7 +223,7 @@ class DDSketchDenseLogarithmic
         auto new_positive_store = std::make_unique<DDSketchDenseStore>();
         auto new_negative_store = std::make_unique<DDSketchDenseStore>();
 
-        auto remap_store = [this, &new_mapping](DDSketchDenseStore& old_store, std::unique_ptr<DDSketchDenseStore>& target_store)
+        auto remap_store = [this, &new_mapping](DDSketchDenseStore & old_store, std::unique_ptr<DDSketchDenseStore> & target_store)
         {
             for (int i = 0; i < old_store.length(); ++i)
             {

src/AggregateFunctions/DDSketch/DDSketchEncoding.h

@@ -1,7 +1,6 @@
 #pragma once
 
-#include <vector>
-#include <stdexcept>
+#include <base/types.h>
 
 /**
   * An encoded DDSketch comprises multiple contiguous blocks (sequences of bytes).
@@ -36,7 +35,10 @@ class DDSketchEncoding
     {
     public:
         UInt8 byte;
-        Flag(UInt8 t, UInt8 s) : byte(t | s) { }
+        Flag(UInt8 t, UInt8 s)
+            : byte(t | s)
+        {
+        }
         [[maybe_unused]] UInt8 Type() const { return byte & flagTypeMask; }
         [[maybe_unused]] UInt8 SubFlag() const { return byte & subFlagMask; }
     };

src/AggregateFunctions/DDSketch/Mapping.h

@@ -1,26 +1,29 @@
 #pragma once
 
-#include <base/types.h>
 #include <cmath>
-#include <stdexcept>
 #include <limits>
+#include <base/types.h>
 
 #include <IO/ReadBuffer.h>
+#include <IO/ReadHelpers.h>
 #include <IO/WriteBuffer.h>
+#include <IO/WriteHelpers.h>
+#include <Common/Exception.h>
 
 namespace DB
 {
 
 namespace ErrorCodes
 {
-    extern const int BAD_ARGUMENTS;
+extern const int BAD_ARGUMENTS;
 }
 
 class DDSketchLogarithmicMapping
 {
 public:
     explicit DDSketchLogarithmicMapping(Float64 relative_accuracy_, Float64 offset_ = 0.0)
-        : relative_accuracy(relative_accuracy_), offset(offset_)
+        : relative_accuracy(relative_accuracy_)
+        , offset(offset_)
     {
         if (relative_accuracy <= 0 || relative_accuracy >= 1)
         {
@@ -44,48 +47,40 @@ class DDSketchLogarithmicMapping
         return static_cast<int>(logGamma(value) + offset);
     }
 
-    Float64 value(int key) const
-    {
-        return lowerBound(key) * (1 + relative_accuracy);
-    }
+    Float64 value(int key) const { return lowerBound(key) * (1 + relative_accuracy); }
 
-    Float64 logGamma(Float64 value) const
-    {
-        return std::log(value) * multiplier;
-    }
+    Float64 logGamma(Float64 value) const { return std::log(value) * multiplier; }
 
-    Float64 powGamma(Float64 value) const
-    {
-        return std::exp(value / multiplier);
-    }
+    Float64 powGamma(Float64 value) const { return std::exp(value / multiplier); }
 
-    Float64 lowerBound(int index) const
-    {
-        return powGamma(static_cast<Float64>(index) - offset);
-    }
+    Float64 lowerBound(int index) const { return powGamma(static_cast<Float64>(index) - offset); }
 
-    Float64 getGamma() const
-    {
-        return gamma;
-    }
+    Float64 getGamma() const { return gamma; }
 
-    Float64 getMinPossible() const
-    {
-        return min_possible;
-    }
+    Float64 getMinPossible() const { return min_possible; }
 
-    [[maybe_unused]] Float64 getMaxPossible() const
+    [[maybe_unused]] Float64 getMaxPossible() const { return max_possible; }
+
+    bool operator==(const DDSketchLogarithmicMapping & other) const
     {
-        return max_possible;
+        if (gamma == other.gamma)
+        {
+            return true;
+        }
+
+        auto [min, max] = std::minmax(gamma, other.gamma);
+        const Float64 difference = max - min;
+        const Float64 acceptable = (std::nextafter(min, max) - min) * min;
+        return difference <= acceptable;
     }
 
-    void serialize(WriteBuffer& buf) const
+    void serialize(WriteBuffer & buf) const
     {
         writeBinary(gamma, buf);
         writeBinary(offset, buf);
     }
 
-    void deserialize(ReadBuffer& buf)
+    void deserialize(ReadBuffer & buf)
     {
         readBinary(gamma, buf);
         readBinary(offset, buf);

src/AggregateFunctions/DDSketch/Store.h

@@ -1,13 +1,15 @@
 #pragma once
 
-#include <base/types.h>
-#include <vector>
 #include <cmath>
 #include <limits>
+#include <vector>
+#include <base/types.h>
 
+#include <AggregateFunctions/DDSketch/DDSketchEncoding.h>
 #include <IO/ReadBuffer.h>
+#include <IO/ReadHelpers.h>
 #include <IO/WriteBuffer.h>
-#include <AggregateFunctions/DDSketch/DDSketchEncoding.h>
+#include <IO/WriteHelpers.h>
 
 
 // We start with 128 bins and grow the number of bins by 128
@@ -18,6 +20,11 @@ constexpr UInt32 CHUNK_SIZE = 128;
 namespace DB
 {
 
+namespace ErrorCodes
+{
+extern const int INCORRECT_DATA;
+}
+
 class DDSketchDenseStore
 {
 public:
@@ -27,9 +34,12 @@ class DDSketchDenseStore
     int offset = 0;
     std::vector<Float64> bins;
 
-    explicit DDSketchDenseStore(UInt32 chunk_size_ = CHUNK_SIZE) : chunk_size(chunk_size_) {}
+    explicit DDSketchDenseStore(UInt32 chunk_size_ = CHUNK_SIZE)
+        : chunk_size(chunk_size_)
+    {
+    }
 
-    void copy(DDSketchDenseStore* other)
+    void copy(DDSketchDenseStore * other)
     {
         bins = other->bins;
         count = other->count;
@@ -38,10 +48,7 @@ class DDSketchDenseStore
         offset = other->offset;
     }
 
-    int length() const
-    {
-        return static_cast<int>(bins.size());
-    }
+    int length() const { return static_cast<int>(bins.size()); }
 
     void add(int key, Float64 weight)
     {
@@ -64,9 +71,10 @@ class DDSketchDenseStore
         return max_key;
     }
 
-    void merge(DDSketchDenseStore* other)
+    void merge(DDSketchDenseStore * other)
     {
-        if (other->count == 0) return;
+        if (other->count == 0)
+            return;
 
         if (count == 0)
         {
@@ -89,9 +97,8 @@ class DDSketchDenseStore
 
     /// NOLINTBEGIN(readability-static-accessed-through-instance)
 
-    void serialize(WriteBuffer& buf) const
+    void serialize(WriteBuffer & buf) const
     {
-
         // Calculate the size of the dense and sparse encodings to choose the smallest one
         UInt64 num_bins = 0;
         UInt64 num_non_empty_bins = 0;
@@ -144,8 +151,10 @@ class DDSketchDenseStore
         }
     }
 
-    void deserialize(ReadBuffer& buf)
+    void deserialize(ReadBuffer & buf)
     {
+        count = 0;
+
         UInt8 encoding_mode;
         readBinary(encoding_mode, buf);
         if (encoding_mode == enc.BinEncodingContiguousCounts)
@@ -165,7 +174,7 @@ class DDSketchDenseStore
                 start_key += index_delta;
             }
         }
-        else
+        else if (encoding_mode == enc.BinEncodingIndexDeltasAndCounts)
         {
             UInt64 num_non_empty_bins;
             readVarUInt(num_non_empty_bins, buf);
@@ -180,6 +189,10 @@ class DDSketchDenseStore
                 add(previous_index, bin_count);
             }
         }
+        else
+        {
+            throw Exception(ErrorCodes::INCORRECT_DATA, "Invalid flag for encoding mode");
+        }
     }
 
     /// NOLINTEND(readability-static-accessed-through-instance)