aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/clickhouse/src/AggregateFunctions/QuantileBFloat16Histogram.h
blob: de9f61e01a208af544ae283fea69e3181068d3cf (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
#pragma once

#include <base/types.h>
#include <base/sort.h>
#include <Common/HashTable/HashMap.h>

#include <IO/ReadBuffer.h>
#include <IO/WriteBuffer.h>


namespace DB
{

/** `bfloat16` is a 16-bit floating point data type that is the same as the corresponding most significant 16 bits of the `float`.
  * https://en.wikipedia.org/wiki/Bfloat16_floating-point_format
  *
  * To calculate quantile, simply convert input value to 16 bit (convert to float, then take the most significant 16 bits),
  * and calculate the histogram of these values.
  *
  * Hash table is the preferred way to store histogram, because the number of distinct values is small:
  * ```
  * SELECT uniq(bfloat)
  * FROM
  * (
  *     SELECT
  *         number,
  *         toFloat32(number) AS f,
  *         bitShiftRight(bitAnd(reinterpretAsUInt32(reinterpretAsFixedString(f)), 4294901760) AS cut, 16),
  *         reinterpretAsFloat32(reinterpretAsFixedString(cut)) AS bfloat
  *     FROM numbers(100000000)
  * )
  *
  * ┌─uniq(bfloat)─┐
  * │         2623 │
  * └──────────────┘
  * ```
  * (when increasing the range of values 1000 times, the number of distinct bfloat16 values increases just by 1280).
  *
  * Then calculate quantile from the histogram.
  *
  * This sketch is very simple and rough. Its relative precision is constant 1 / 256 = 0.390625%.
  */
template <typename Value>
struct QuantileBFloat16Histogram
{
    using BFloat16 = UInt16;
    using Weight = UInt64;

    /// Make automatic memory for 16 elements to avoid allocations for small states.
    /// The usage of trivial hash is ok, because we effectively take logarithm of the values and pathological cases are unlikely.
    using Data = HashMapWithStackMemory<BFloat16, Weight, TrivialHash, 4>;

    Data data;

    void add(const Value & x)
    {
        add(x, 1);
    }

    void add(const Value & x, Weight w)
    {
        if (!isNaN(x))
            data[toBFloat16(x)] += w;
    }

    void merge(const QuantileBFloat16Histogram & rhs)
    {
        for (const auto & pair : rhs.data)
            data[pair.getKey()] += pair.getMapped();
    }

    void serialize(WriteBuffer & buf) const
    {
        data.write(buf);
    }

    void deserialize(ReadBuffer & buf)
    {
        data.read(buf);
    }

    Value get(Float64 level) const
    {
        return getImpl<Value>(level);
    }

    void getMany(const Float64 * levels, const size_t * indices, size_t size, Value * result) const
    {
        getManyImpl(levels, indices, size, result);
    }

    Float64 getFloat(Float64 level) const
    {
        return getImpl<Float64>(level);
    }

    void getManyFloat(const Float64 * levels, const size_t * indices, size_t size, Float64 * result) const
    {
        getManyImpl(levels, indices, size, result);
    }

private:
    /// Take the most significant 16 bits of the floating point number.
    BFloat16 toBFloat16(const Value & x) const
    {
        return std::bit_cast<UInt32>(static_cast<Float32>(x)) >> 16;
    }

    /// Put the bits into most significant 16 bits of the floating point number and fill other bits with zeros.
    Float32 toFloat32(const BFloat16 & x) const
    {
        return std::bit_cast<Float32>(x << 16);
    }

    using Pair = PairNoInit<Float32, Weight>;

    template <typename T>
    T getImpl(Float64 level) const
    {
        size_t size = data.size();

        if (0 == size)
            return std::numeric_limits<T>::quiet_NaN();

        std::unique_ptr<Pair[]> array_holder(new Pair[size]);
        Pair * array = array_holder.get();

        Float64 sum_weight = 0;
        Pair * arr_it = array;
        for (const auto & pair : data)
        {
            sum_weight += pair.getMapped();
            *arr_it = {toFloat32(pair.getKey()), pair.getMapped()};
            ++arr_it;
        }

        ::sort(array, array + size, [](const Pair & a, const Pair & b) { return a.first < b.first; });

        Float64 threshold = std::ceil(sum_weight * level);
        Float64 accumulated = 0;

        for (const Pair * p = array; p != (array + size); ++p)
        {
            accumulated += p->second;

            if (accumulated >= threshold)
                return static_cast<T>(p->first);
        }

        return static_cast<T>(array[size - 1].first);
    }

    template <typename T>
    void getManyImpl(const Float64 * levels, const size_t * indices, size_t num_levels, T * result) const
    {
        size_t size = data.size();

        if (0 == size)
        {
            for (size_t i = 0; i < num_levels; ++i)
                result[i] = std::numeric_limits<T>::quiet_NaN();

            return;
        }

        std::unique_ptr<Pair[]> array_holder(new Pair[size]);
        Pair * array = array_holder.get();

        Float64 sum_weight = 0;
        Pair * arr_it = array;
        for (const auto & pair : data)
        {
            sum_weight += pair.getMapped();
            *arr_it = {toFloat32(pair.getKey()), pair.getMapped()};
            ++arr_it;
        }

        ::sort(array, array + size, [](const Pair & a, const Pair & b) { return a.first < b.first; });

        size_t level_index = 0;
        Float64 accumulated = 0;
        Float64 threshold = std::ceil(sum_weight * levels[indices[level_index]]);

        for (const Pair * p = array; p != (array + size); ++p)
        {
            accumulated += p->second;

            while (accumulated >= threshold)
            {
                result[indices[level_index]] = static_cast<T>(p->first);
                ++level_index;

                if (level_index == num_levels)
                    return;

                threshold = std::ceil(sum_weight * levels[indices[level_index]]);
            }
        }

        while (level_index < num_levels)
        {
            result[indices[level_index]] = static_cast<T>(array[size - 1].first);
            ++level_index;
        }
    }
};

}