ClickHouse

1
#include <memory>
2
#include <random>
3

4
#include <DataTypes/DataTypesNumber.h>
5
#include <Common/thread_local_rng.h>
6
#include <IO/ReadBuffer.h>
7
#include <IO/WriteBuffer.h>
8
#include <AggregateFunctions/IAggregateFunction.h>
9
#include <AggregateFunctions/AggregateFunctionFactory.h>
10

11

12
namespace DB
13
{
14
struct Settings;
15

16
namespace ErrorCodes
17
{
18
    extern const int AGGREGATE_FUNCTION_THROW;
19
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
20
}
21

22
namespace
23
{
24

25
struct AggregateFunctionThrowData
26
{
27
    bool allocated;
28

29
    AggregateFunctionThrowData() : allocated(true) {}
30
    ~AggregateFunctionThrowData()
31
    {
32
        volatile bool * allocated_ptr = &allocated;
33

34
        if (*allocated_ptr)
35
            *allocated_ptr = false;
36
        else
37
            abort();
38
    }
39
};
40

41
/** Throw on creation with probability specified in parameter.
42
  * It will check correct destruction of the state.
43
  * This is intended to check for exception safety.
44
  */
45
class AggregateFunctionThrow final : public IAggregateFunctionDataHelper<AggregateFunctionThrowData, AggregateFunctionThrow>
46
{
47
private:
48
    Float64 throw_probability;
49

50
public:
51
    AggregateFunctionThrow(const DataTypes & argument_types_, const Array & parameters_, Float64 throw_probability_)
52
        : IAggregateFunctionDataHelper(argument_types_, parameters_, createResultType())
53
        , throw_probability(throw_probability_)
54
    {}
55

56
    String getName() const override
57
    {
58
        return "aggThrow";
59
    }
60

61
    static DataTypePtr createResultType()
62
    {
63
        return std::make_shared<DataTypeUInt8>();
64
    }
65

66
    bool allocatesMemoryInArena() const override { return false; }
67

68
    void create(AggregateDataPtr __restrict place) const override
69
    {
70
        if (std::uniform_real_distribution<>(0.0, 1.0)(thread_local_rng) <= throw_probability)
71
            throw Exception(ErrorCodes::AGGREGATE_FUNCTION_THROW, "Aggregate function {} has thrown exception successfully", getName());
72

73
        new (place) Data;
74
    }
75

76
    void destroy(AggregateDataPtr __restrict place) const noexcept override
77
    {
78
        data(place).~Data();
79
    }
80

81
    bool hasTrivialDestructor() const override { return std::is_trivially_destructible_v<Data>; }
82

83
    void add(AggregateDataPtr __restrict, const IColumn **, size_t, Arena *) const override
84
    {
85
    }
86

87
    void merge(AggregateDataPtr __restrict, ConstAggregateDataPtr, Arena *) const override
88
    {
89
    }
90

91
    void serialize(ConstAggregateDataPtr __restrict, WriteBuffer & buf, std::optional<size_t> /* version */) const override
92
    {
93
        char c = 0;
94
        buf.write(c);
95
    }
96

97
    void deserialize(AggregateDataPtr __restrict /* place */, ReadBuffer & buf, std::optional<size_t> /* version */, Arena *) const override
98
    {
99
        char c = 0;
100
        buf.readStrict(c);
101
    }
102

103
    void insertResultInto(AggregateDataPtr __restrict, IColumn & to, Arena *) const override
104
    {
105
        to.insertDefault();
106
    }
107
};
108

109
}
110

111
void registerAggregateFunctionAggThrow(AggregateFunctionFactory & factory)
112
{
113
    factory.registerFunction("aggThrow", [](const std::string & name, const DataTypes & argument_types, const Array & parameters, const Settings *)
114
    {
115
        Float64 throw_probability = 1.0;
116
        if (parameters.size() == 1)
117
            throw_probability = parameters[0].safeGet<Float64>();
118
        else if (parameters.size() > 1)
119
            throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Aggregate function {} cannot have more than one parameter", name);
120

121
        return std::make_shared<AggregateFunctionThrow>(argument_types, parameters, throw_probability);
122
    });
123
}
124

125
}
126