ClickHouse

1
#include <DataTypes/DataTypeFactory.h>
2
#include <DataTypes/DataTypeNullable.h>
3
#include <DataTypes/DataTypeString.h>
4
#include <DataTypes/DataTypesNumber.h>
5
#include <DataTypes/DataTypesDecimal.h>
6
#include <DataTypes/DataTypeDate.h>
7
#include <DataTypes/DataTypeDate32.h>
8
#include <DataTypes/DataTypeDateTime.h>
9
#include <DataTypes/DataTypeDateTime64.h>
10
#include <DataTypes/DataTypeUUID.h>
11
#include <DataTypes/DataTypeIPv4andIPv6.h>
12
#include <Columns/ColumnString.h>
13
#include <Columns/ColumnConst.h>
14
#include <Columns/ColumnNullable.h>
15

16
#include <Interpreters/Context.h>
17
#include <Interpreters/castColumn.h>
18

19
#include <Functions/IFunction.h>
20
#include <Functions/FunctionHelpers.h>
21
#include <Functions/FunctionFactory.h>
22
#include <Functions/extractTimeZoneFromFunctionArguments.h>
23

24
namespace DB
25
{
26

27
namespace ErrorCodes
28
{
29
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
30
    extern const int BAD_ARGUMENTS;
31
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
32
}
33

34
class FunctionCastOrDefault final : public IFunction
35
{
36
public:
37
    static constexpr auto name = "accurateCastOrDefault";
38

39
    static FunctionPtr create(ContextPtr context)
40
    {
41
        return std::make_shared<FunctionCastOrDefault>(context);
42
    }
43

44
    explicit FunctionCastOrDefault(ContextPtr context_)
45
        : keep_nullable(context_->getSettingsRef().cast_keep_nullable)
46
    {
47
    }
48

49
    String getName() const override { return name; }
50

51
    size_t getNumberOfArguments() const override { return 0; }
52
    bool isVariadic() const override { return true; }
53
    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {1}; }
54

55
    bool useDefaultImplementationForNulls() const override { return false; }
56
    bool useDefaultImplementationForNothing() const override { return false; }
57
    bool useDefaultImplementationForConstants() const override { return false; }
58
    bool useDefaultImplementationForLowCardinalityColumns() const override { return true; }
59
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
60

61
    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
62
    {
63
        size_t arguments_size = arguments.size();
64
        if (arguments_size != 2 && arguments_size != 3)
65
            throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
66
                "Function {} expected 2 or 3 arguments. Actual {}",
67
                getName(),
68
                arguments_size);
69

70
        const auto & type_column = arguments[1].column;
71
        const auto * type_column_typed = checkAndGetColumnConst<ColumnString>(type_column.get());
72

73
        if (!type_column_typed)
74
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
75
                "Second argument to {} must be a constant string describing type. Actual {}",
76
                getName(),
77
                arguments[1].type->getName());
78

79
        DataTypePtr result_type = DataTypeFactory::instance().get(type_column_typed->getValue<String>());
80

81
        if (keep_nullable && arguments.front().type->isNullable())
82
            result_type = makeNullable(result_type);
83

84
        if (arguments.size() == 3)
85
        {
86
            auto default_value_type = arguments[2].type;
87

88
            if (!result_type->equals(*default_value_type))
89
            {
90
                throw Exception(ErrorCodes::BAD_ARGUMENTS,
91
                    "Default value type should be same as cast type. Expected {}. Actual {}",
92
                    result_type->getName(),
93
                    default_value_type->getName());
94
            }
95
        }
96

97
        return result_type;
98
    }
99

100
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & return_type, size_t) const override
101
    {
102
        const ColumnWithTypeAndName & column_to_cast = arguments[0];
103
        auto non_const_column_to_cast = column_to_cast.column->convertToFullColumnIfConst();
104
        ColumnWithTypeAndName column_to_cast_non_const { non_const_column_to_cast, column_to_cast.type, column_to_cast.name };
105

106
        auto cast_result = castColumnAccurateOrNull(column_to_cast_non_const, return_type);
107

108
        const auto & cast_result_nullable = assert_cast<const ColumnNullable &>(*cast_result);
109
        const auto & null_map_data = cast_result_nullable.getNullMapData();
110
        size_t null_map_data_size = null_map_data.size();
111
        const auto & nested_column = cast_result_nullable.getNestedColumn();
112
        auto result = return_type->createColumn();
113
        result->reserve(null_map_data_size);
114

115
        ColumnNullable * result_nullable = nullptr;
116
        if (result->isNullable())
117
            result_nullable = assert_cast<ColumnNullable *>(&*result);
118

119
        size_t start_insert_index = 0;
120

121
        Field default_value;
122
        ColumnPtr default_column;
123

124
        if (arguments.size() == 3)
125
        {
126
            auto default_values_column = arguments[2].column;
127

128
            if (isColumnConst(*default_values_column))
129
                default_value = (*default_values_column)[0];
130
            else
131
                default_column = default_values_column->convertToFullColumnIfConst();
132
        }
133
        else
134
        {
135
            default_value = return_type->getDefault();
136
        }
137

138
        for (size_t i = 0; i < null_map_data_size; ++i)
139
        {
140
            bool is_current_index_null = null_map_data[i];
141
            if (!is_current_index_null)
142
                continue;
143

144
            if (i != start_insert_index)
145
            {
146
                if (result_nullable)
147
                    result_nullable->insertRangeFromNotNullable(nested_column, start_insert_index, i - start_insert_index);
148
                else
149
                    result->insertRangeFrom(nested_column, start_insert_index, i - start_insert_index);
150
            }
151

152
            if (default_column)
153
                result->insertFrom(*default_column, i);
154
            else
155
                result->insert(default_value);
156

157
            start_insert_index = i + 1;
158
        }
159

160
        if (null_map_data_size != start_insert_index)
161
        {
162
            if (result_nullable)
163
                result_nullable->insertRangeFromNotNullable(nested_column, start_insert_index, null_map_data_size - start_insert_index);
164
            else
165
                result->insertRangeFrom(nested_column, start_insert_index, null_map_data_size - start_insert_index);
166
        }
167

168
        return result;
169
    }
170

171
private:
172

173
    bool keep_nullable;
174
};
175

176
class FunctionCastOrDefaultTyped final : public IFunction
177
{
178
public:
179
    explicit FunctionCastOrDefaultTyped(ContextPtr context_, String name_, DataTypePtr type_)
180
        : impl(context_), name(std::move(name_)), type(std::move(type_)), which(type)
181
    {
182
    }
183

184
    String getName() const override { return name; }
185

186
private:
187
    FunctionCastOrDefault impl;
188
    String name;
189
    DataTypePtr type;
190
    WhichDataType which;
191

192
    size_t getNumberOfArguments() const override { return 0; }
193
    bool isVariadic() const override { return true; }
194

195
    bool useDefaultImplementationForNulls() const override { return impl.useDefaultImplementationForNulls(); }
196
    bool useDefaultImplementationForNothing() const override { return impl.useDefaultImplementationForNothing(); }
197
    bool useDefaultImplementationForLowCardinalityColumns() const override { return impl.useDefaultImplementationForLowCardinalityColumns();}
198
    bool useDefaultImplementationForConstants() const override { return impl.useDefaultImplementationForConstants();}
199
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & arguments) const override
200
    {
201
        return impl.isSuitableForShortCircuitArgumentsExecution(arguments);
202
    }
203

204
    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
205
    {
206
        FunctionArgumentDescriptors mandatory_args = {{"Value", nullptr, nullptr, nullptr}};
207
        FunctionArgumentDescriptors optional_args;
208

209
        if (isDecimal(type) || isDateTime64(type))
210
            mandatory_args.push_back({"scale", static_cast<FunctionArgumentDescriptor::TypeValidator>(&isNativeInteger), &isColumnConst, "const Integer"});
211

212
        if (isDateTimeOrDateTime64(type))
213
            optional_args.push_back({"timezone", static_cast<FunctionArgumentDescriptor::TypeValidator>(&isString), isColumnConst, "const String"});
214

215
        optional_args.push_back({"default_value", nullptr, nullptr, nullptr});
216

217
        validateFunctionArgumentTypes(*this, arguments, mandatory_args, optional_args);
218

219
        size_t additional_argument_index = 1;
220

221
        size_t scale = 0;
222
        std::string time_zone;
223

224
        if (isDecimal(type) || isDateTime64(type))
225
        {
226
            const auto & scale_argument = arguments[additional_argument_index];
227

228
            WhichDataType scale_argument_type(scale_argument.type);
229

230
            if (!scale_argument_type.isNativeUInt())
231
            {
232
                throw Exception(ErrorCodes::BAD_ARGUMENTS,
233
                    "Function {} decimal scale should have native UInt type. Actual {}",
234
                    getName(), scale_argument.type->getName());
235
            }
236

237
            scale = arguments[additional_argument_index].column->getUInt(0);
238
            ++additional_argument_index;
239
        }
240

241
        if (isDateTimeOrDateTime64(type))
242
        {
243
            if (additional_argument_index < arguments.size())
244
            {
245
                time_zone = extractTimeZoneNameFromColumn(arguments[additional_argument_index].column.get(),
246
                                                          arguments[additional_argument_index].name);
247
                ++additional_argument_index;
248
            }
249
        }
250

251
        DataTypePtr cast_type;
252

253
        if (which.isDateTime64())
254
            cast_type = std::make_shared<DataTypeDateTime64>(scale, time_zone);
255
        else if (which.isDateTime())
256
            cast_type = std::make_shared<DataTypeDateTime>(time_zone);
257
        else if (which.isDecimal32())
258
            cast_type = createDecimalMaxPrecision<Decimal32>(scale);
259
        else if (which.isDecimal64())
260
            cast_type = createDecimalMaxPrecision<Decimal64>(scale);
261
        else if (which.isDecimal128())
262
            cast_type = createDecimalMaxPrecision<Decimal128>(scale);
263
        else if (which.isDecimal256())
264
            cast_type = createDecimalMaxPrecision<Decimal256>(scale);
265
        else
266
            cast_type = type;
267

268
        ColumnWithTypeAndName type_argument =
269
        {
270
            DataTypeString().createColumnConst(1, cast_type->getName()),
271
            std::make_shared<DataTypeString>(),
272
            ""
273
        };
274

275
        ColumnsWithTypeAndName arguments_with_cast_type;
276
        arguments_with_cast_type.reserve(arguments.size());
277

278
        arguments_with_cast_type.emplace_back(arguments[0]);
279
        arguments_with_cast_type.emplace_back(type_argument);
280

281
        if (additional_argument_index < arguments.size())
282
        {
283
            arguments_with_cast_type.emplace_back(arguments[additional_argument_index]);
284
            ++additional_argument_index;
285
        }
286

287
        if (additional_argument_index < arguments.size())
288
            throw Exception(ErrorCodes::BAD_ARGUMENTS, "{} wrong arguments size", getName());
289

290
        return impl.getReturnTypeImpl(arguments_with_cast_type);
291
    }
292

293
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_size) const override
294
    {
295
        /// Scale and time zone
296
        size_t additional_arguments_size = (which.isDecimal() || which.isDateTime64()) + which.isDateTimeOrDateTime64();
297

298
        ColumnWithTypeAndName second_argument =
299
        {
300
            DataTypeString().createColumnConst(arguments.begin()->column->size(), result_type->getName()),
301
            std::make_shared<DataTypeString>(),
302
            ""
303
        };
304

305
        ColumnsWithTypeAndName arguments_with_cast_type;
306
        arguments_with_cast_type.reserve(arguments.size() + 1);
307

308
        arguments_with_cast_type.emplace_back(arguments[0]);
309
        arguments_with_cast_type.emplace_back(second_argument);
310

311
        size_t default_column_argument = 1 + additional_arguments_size;
312
        if (default_column_argument < arguments.size())
313
            arguments_with_cast_type.emplace_back(arguments[default_column_argument]);
314

315
        return impl.executeImpl(arguments_with_cast_type, result_type, input_rows_size);
316
    }
317
};
318

319
REGISTER_FUNCTION(CastOrDefault)
320
{
321
    factory.registerFunction<FunctionCastOrDefault>();
322

323
    factory.registerFunction("toUInt8OrDefault", [](ContextPtr context)
324
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toUInt8OrDefault", std::make_shared<DataTypeUInt8>()); });
325
    factory.registerFunction("toUInt16OrDefault", [](ContextPtr context)
326
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toUInt16OrDefault", std::make_shared<DataTypeUInt16>()); });
327
    factory.registerFunction("toUInt32OrDefault", [](ContextPtr context)
328
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toUInt32OrDefault", std::make_shared<DataTypeUInt32>()); });
329
    factory.registerFunction("toUInt64OrDefault", [](ContextPtr context)
330
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toUInt64OrDefault", std::make_shared<DataTypeUInt64>()); });
331
    factory.registerFunction("toUInt128OrDefault", [](ContextPtr context)
332
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toUInt128OrDefault", std::make_shared<DataTypeUInt128>()); },
333
        FunctionDocumentation{
334
            .description=R"(
335
Converts a string in the first argument of the function to UInt128 by parsing it.
336
If it cannot parse the value, returns the default value, which can be provided as the second function argument, and if provided, must be of UInt128 type.
337
If the default value is not provided in the second argument, it is assumed to be zero.
338
)",
339
            .examples{
340
                {"Successful conversion", "SELECT toUInt128OrDefault('1', 2::UInt128)", "1"},
341
                {"Default value", "SELECT toUInt128OrDefault('upyachka', 123456789012345678901234567890::UInt128)", "123456789012345678901234567890"},
342
                {"Implicit default value", "SELECT toUInt128OrDefault('upyachka')", "0"}},
343
            .categories{"ConversionFunctions"}
344
        });
345
    factory.registerFunction("toUInt256OrDefault", [](ContextPtr context)
346
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toUInt256OrDefault", std::make_shared<DataTypeUInt256>()); });
347

348
    factory.registerFunction("toInt8OrDefault", [](ContextPtr context)
349
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toInt8OrDefault", std::make_shared<DataTypeInt8>()); });
350
    factory.registerFunction("toInt16OrDefault", [](ContextPtr context)
351
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toInt16OrDefault", std::make_shared<DataTypeInt16>()); });
352
    factory.registerFunction("toInt32OrDefault", [](ContextPtr context)
353
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toInt32OrDefault", std::make_shared<DataTypeInt32>()); });
354
    factory.registerFunction("toInt64OrDefault", [](ContextPtr context)
355
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toInt64OrDefault", std::make_shared<DataTypeInt64>()); });
356
    factory.registerFunction("toInt128OrDefault", [](ContextPtr context)
357
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toInt128OrDefault", std::make_shared<DataTypeInt128>()); });
358
    factory.registerFunction("toInt256OrDefault", [](ContextPtr context)
359
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toInt256OrDefault", std::make_shared<DataTypeInt256>()); });
360

361
    factory.registerFunction("toFloat32OrDefault", [](ContextPtr context)
362
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toFloat32OrDefault", std::make_shared<DataTypeFloat32>()); });
363
    factory.registerFunction("toFloat64OrDefault", [](ContextPtr context)
364
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toFloat64OrDefault", std::make_shared<DataTypeFloat64>()); });
365

366
    factory.registerFunction("toDateOrDefault", [](ContextPtr context)
367
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toDateOrDefault", std::make_shared<DataTypeDate>()); });
368
    factory.registerFunction("toDate32OrDefault", [](ContextPtr context)
369
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toDate32OrDefault", std::make_shared<DataTypeDate32>()); });
370
    factory.registerFunction("toDateTimeOrDefault", [](ContextPtr context)
371
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toDateTimeOrDefault", std::make_shared<DataTypeDateTime>()); });
372
    factory.registerFunction("toDateTime64OrDefault", [](ContextPtr context)
373
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toDateTime64OrDefault", std::make_shared<DataTypeDateTime64>(3 /* default scale */)); });
374

375
    factory.registerFunction("toDecimal32OrDefault", [](ContextPtr context)
376
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toDecimal32OrDefault", createDecimalMaxPrecision<Decimal32>(0)); });
377
    factory.registerFunction("toDecimal64OrDefault", [](ContextPtr context)
378
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toDecimal64OrDefault", createDecimalMaxPrecision<Decimal64>(0)); });
379
    factory.registerFunction("toDecimal128OrDefault", [](ContextPtr context)
380
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toDecimal128OrDefault", createDecimalMaxPrecision<Decimal128>(0)); });
381
    factory.registerFunction("toDecimal256OrDefault", [](ContextPtr context)
382
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toDecimal256OrDefault", createDecimalMaxPrecision<Decimal256>(0)); });
383

384
    factory.registerFunction("toUUIDOrDefault", [](ContextPtr context)
385
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toUUIDOrDefault", std::make_shared<DataTypeUUID>()); });
386
    factory.registerFunction("toIPv4OrDefault", [](ContextPtr context)
387
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toIPv4OrDefault", std::make_shared<DataTypeIPv4>()); });
388
    factory.registerFunction("toIPv6OrDefault", [](ContextPtr context)
389
        { return std::make_shared<FunctionCastOrDefaultTyped>(context, "toIPv6OrDefault", std::make_shared<DataTypeIPv6>()); });
390
}
391

392
}
393