ClickHouse

1
#include <Columns/ColumnsDateTime.h>
2
#include <Columns/ColumnsNumber.h>
3
#include <Common/DateLUTImpl.h>
4
#include <Common/IntervalKind.h>
5
#include <DataTypes/DataTypeDate.h>
6
#include <DataTypes/DataTypeDateTime.h>
7
#include <DataTypes/DataTypeDateTime64.h>
8
#include <DataTypes/DataTypeInterval.h>
9
#include <Functions/DateTimeTransforms.h>
10
#include <Functions/FunctionFactory.h>
11
#include <Functions/IFunction.h>
12
#include <IO/WriteHelpers.h>
13

14

15
namespace DB
16
{
17
namespace ErrorCodes
18
{
19
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
20
    extern const int ILLEGAL_COLUMN;
21
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
22
    extern const int ARGUMENT_OUT_OF_BOUND;
23
}
24

25

26
class FunctionToStartOfInterval : public IFunction
27
{
28
public:
29
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionToStartOfInterval>(); }
30

31
    static constexpr auto name = "toStartOfInterval";
32
    String getName() const override { return name; }
33
    bool isVariadic() const override { return true; }
34
    size_t getNumberOfArguments() const override { return 0; }
35
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
36
    bool useDefaultImplementationForConstants() const override { return true; }
37
    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {1, 2}; }
38
    bool hasInformationAboutMonotonicity() const override { return true; }
39
    Monotonicity getMonotonicityForRange(const IDataType &, const Field &, const Field &) const override { return { .is_monotonic = true, .is_always_monotonic = true }; }
40

41
    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
42
    {
43
        bool value_is_date = false;
44
        auto check_first_argument = [&]
45
        {
46
            const DataTypePtr & type_arg1 = arguments[0].type;
47
            if (!isDate(type_arg1) && !isDateTime(type_arg1) && !isDateTime64(type_arg1))
48
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
49
                    "Illegal type {} of 1st argument of function {}, expected a Date, DateTime or DateTime64",
50
                    type_arg1->getName(), getName());
51
            value_is_date = isDate(type_arg1);
52
        };
53

54
        const DataTypeInterval * interval_type = nullptr;
55
        enum class ResultType
56
        {
57
            Date,
58
            DateTime,
59
            DateTime64
60
        };
61
        ResultType result_type;
62
        auto check_second_argument = [&]
63
        {
64
            const DataTypePtr & type_arg2 = arguments[1].type;
65

66
            interval_type = checkAndGetDataType<DataTypeInterval>(type_arg2.get());
67
            if (!interval_type)
68
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
69
                    "Illegal type {} of 2nd argument of function {}, expected a time interval",
70
                    type_arg2->getName(), getName());
71

72
            switch (interval_type->getKind()) // NOLINT(bugprone-switch-missing-default-case)
73
            {
74
                case IntervalKind::Kind::Nanosecond:
75
                case IntervalKind::Kind::Microsecond:
76
                case IntervalKind::Kind::Millisecond:
77
                    result_type = ResultType::DateTime64;
78
                    break;
79
                case IntervalKind::Kind::Second:
80
                case IntervalKind::Kind::Minute:
81
                case IntervalKind::Kind::Hour:
82
                case IntervalKind::Kind::Day: /// weird why Day leads to DateTime but too afraid to change it
83
                    result_type = ResultType::DateTime;
84
                    break;
85
                case IntervalKind::Kind::Week:
86
                case IntervalKind::Kind::Month:
87
                case IntervalKind::Kind::Quarter:
88
                case IntervalKind::Kind::Year:
89
                    result_type = ResultType::Date;
90
                    break;
91
            }
92
        };
93

94
        auto check_third_argument = [&]
95
        {
96
            const DataTypePtr & type_arg3 = arguments[2].type;
97
            if (!isString(type_arg3))
98
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
99
                    "Illegal type {} of 3rd argument of function {}, expected a constant timezone string",
100
                    type_arg3->getName(), getName());
101
            if (value_is_date && result_type == ResultType::Date) /// weird why this is && instead of || but too afraid to change it
102
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
103
                    "The timezone argument of function {} with interval type {} is allowed only when the 1st argument has type DateTime or DateTimt64",
104
                    getName(), interval_type->getKind().toString());
105
        };
106

107
        if (arguments.size() == 2)
108
        {
109
            check_first_argument();
110
            check_second_argument();
111
        }
112
        else if (arguments.size() == 3)
113
        {
114
            check_first_argument();
115
            check_second_argument();
116
            check_third_argument();
117
        }
118
        else
119
        {
120
            throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
121
                "Number of arguments for function {} doesn't match: passed {}, should be 2 or 3",
122
                getName(), arguments.size());
123
        }
124

125
        switch (result_type)
126
        {
127
            case ResultType::Date:
128
                return std::make_shared<DataTypeDate>();
129
            case ResultType::DateTime:
130
                return std::make_shared<DataTypeDateTime>(extractTimeZoneNameFromFunctionArguments(arguments, 2, 0, false));
131
            case ResultType::DateTime64:
132
            {
133
                UInt32 scale = 0;
134
                if (interval_type->getKind() == IntervalKind::Kind::Nanosecond)
135
                    scale = 9;
136
                else if (interval_type->getKind() == IntervalKind::Kind::Microsecond)
137
                    scale = 6;
138
                else if (interval_type->getKind() == IntervalKind::Kind::Millisecond)
139
                    scale = 3;
140

141
                return std::make_shared<DataTypeDateTime64>(scale, extractTimeZoneNameFromFunctionArguments(arguments, 2, 0, false));
142
            }
143
        }
144

145
        std::unreachable();
146
    }
147

148
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t /* input_rows_count */) const override
149
    {
150
        const auto & time_column = arguments[0];
151
        const auto & interval_column = arguments[1];
152
        const auto & time_zone = extractTimeZoneFromFunctionArguments(arguments, 2, 0);
153
        auto result_column = dispatchForTimeColumn(time_column, interval_column, result_type, time_zone);
154
        return result_column;
155
    }
156

157
private:
158
    ColumnPtr dispatchForTimeColumn(
159
        const ColumnWithTypeAndName & time_column, const ColumnWithTypeAndName & interval_column,
160
        const DataTypePtr & result_type, const DateLUTImpl & time_zone) const
161
    {
162
        const auto & time_column_type = *time_column.type.get();
163
        const auto & time_column_col = *time_column.column.get();
164

165
        if (isDateTime64(time_column_type))
166
        {
167
            const auto * time_column_vec = checkAndGetColumn<ColumnDateTime64>(time_column_col);
168
            auto scale = assert_cast<const DataTypeDateTime64 &>(time_column_type).getScale();
169

170
            if (time_column_vec)
171
                return dispatchForIntervalColumn(assert_cast<const DataTypeDateTime64 &>(time_column_type), *time_column_vec, interval_column, result_type, time_zone, scale);
172
        }
173
        else if (isDateTime(time_column_type))
174
        {
175
            const auto * time_column_vec = checkAndGetColumn<ColumnDateTime>(time_column_col);
176
            if (time_column_vec)
177
                return dispatchForIntervalColumn(assert_cast<const DataTypeDateTime &>(time_column_type), *time_column_vec, interval_column, result_type, time_zone);
178
        }
179
        else if (isDate(time_column_type))
180
        {
181
            const auto * time_column_vec = checkAndGetColumn<ColumnDate>(time_column_col);
182
            if (time_column_vec)
183
                return dispatchForIntervalColumn(assert_cast<const DataTypeDate &>(time_column_type), *time_column_vec, interval_column, result_type, time_zone);
184
        }
185
        throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal column for 1st argument of function {}, expected a Date, DateTime or DateTime64", getName());
186
    }
187

188
    template <typename TimeDataType, typename TimeColumnType>
189
    ColumnPtr dispatchForIntervalColumn(
190
        const TimeDataType & time_data_type, const TimeColumnType & time_column, const ColumnWithTypeAndName & interval_column,
191
        const DataTypePtr & result_type, const DateLUTImpl & time_zone, UInt16 scale = 1) const
192
    {
193
        const auto * interval_type = checkAndGetDataType<DataTypeInterval>(interval_column.type.get());
194
        if (!interval_type)
195
            throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Illegal column for 2nd argument of function {}, must be a time interval", getName());
196

197
        const auto * interval_column_const_int64 = checkAndGetColumnConst<ColumnInt64>(interval_column.column.get());
198
        if (!interval_column_const_int64)
199
            throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Illegal column for 2nd argument of function {}, must be a const time interval", getName());
200

201
        const Int64 num_units = interval_column_const_int64->getValue<Int64>();
202
        if (num_units <= 0)
203
            throw Exception(ErrorCodes::ARGUMENT_OUT_OF_BOUND, "Value for 2nd argument of function {} must be positive", getName());
204

205
        switch (interval_type->getKind()) // NOLINT(bugprone-switch-missing-default-case)
206
        {
207
            case IntervalKind::Kind::Nanosecond:
208
                return execute<TimeDataType, TimeColumnType, DataTypeDateTime64, IntervalKind::Kind::Nanosecond>(time_data_type, time_column, num_units, result_type, time_zone, scale);
209
            case IntervalKind::Kind::Microsecond:
210
                return execute<TimeDataType, TimeColumnType, DataTypeDateTime64, IntervalKind::Kind::Microsecond>(time_data_type, time_column, num_units, result_type, time_zone, scale);
211
            case IntervalKind::Kind::Millisecond:
212
                return execute<TimeDataType, TimeColumnType, DataTypeDateTime64, IntervalKind::Kind::Millisecond>(time_data_type, time_column, num_units, result_type, time_zone, scale);
213
            case IntervalKind::Kind::Second:
214
                return execute<TimeDataType, TimeColumnType, DataTypeDateTime, IntervalKind::Kind::Second>(time_data_type, time_column, num_units, result_type, time_zone, scale);
215
            case IntervalKind::Kind::Minute:
216
                return execute<TimeDataType, TimeColumnType, DataTypeDateTime, IntervalKind::Kind::Minute>(time_data_type, time_column, num_units, result_type, time_zone, scale);
217
            case IntervalKind::Kind::Hour:
218
                return execute<TimeDataType, TimeColumnType, DataTypeDateTime, IntervalKind::Kind::Hour>(time_data_type, time_column, num_units, result_type, time_zone, scale);
219
            case IntervalKind::Kind::Day:
220
                return execute<TimeDataType, TimeColumnType, DataTypeDateTime, IntervalKind::Kind::Day>(time_data_type, time_column, num_units, result_type, time_zone, scale);
221
            case IntervalKind::Kind::Week:
222
                return execute<TimeDataType, TimeColumnType, DataTypeDate, IntervalKind::Kind::Week>(time_data_type, time_column, num_units, result_type, time_zone, scale);
223
            case IntervalKind::Kind::Month:
224
                return execute<TimeDataType, TimeColumnType, DataTypeDate, IntervalKind::Kind::Month>(time_data_type, time_column, num_units, result_type, time_zone, scale);
225
            case IntervalKind::Kind::Quarter:
226
                return execute<TimeDataType, TimeColumnType, DataTypeDate, IntervalKind::Kind::Quarter>(time_data_type, time_column, num_units, result_type, time_zone, scale);
227
            case IntervalKind::Kind::Year:
228
                return execute<TimeDataType, TimeColumnType, DataTypeDate, IntervalKind::Kind::Year>(time_data_type, time_column, num_units, result_type, time_zone, scale);
229
        }
230

231
        std::unreachable();
232
    }
233

234
    template <typename TimeDataType, typename TimeColumnType, typename ResultDataType, IntervalKind::Kind unit>
235
    ColumnPtr execute(
236
        const TimeDataType &, const TimeColumnType & time_column_type, Int64 num_units,
237
        const DataTypePtr & result_type, const DateLUTImpl & time_zone, UInt16 scale) const
238
    {
239
        using ResultColumnType = typename ResultDataType::ColumnType;
240
        using ResultFieldType = typename ResultDataType::FieldType;
241

242
        const auto & time_data = time_column_type.getData();
243
        size_t size = time_data.size();
244

245
        auto result_col = result_type->createColumn();
246
        auto * col_to = assert_cast<ResultColumnType *>(result_col.get());
247
        auto & result_data = col_to->getData();
248
        result_data.resize(size);
249

250
        Int64 scale_multiplier = DecimalUtils::scaleMultiplier<DateTime64>(scale);
251

252
        for (size_t i = 0; i != size; ++i)
253
            result_data[i] = static_cast<ResultFieldType>(ToStartOfInterval<unit>::execute(time_data[i], num_units, time_zone, scale_multiplier));
254

255
        return result_col;
256
    }
257
};
258

259
REGISTER_FUNCTION(ToStartOfInterval)
260
{
261
    factory.registerFunction<FunctionToStartOfInterval>();
262
}
263

264
}
265