ClickHouse
119 строк · 4.4 Кб
1#include <DataTypes/DataTypeArray.h>
2#include <DataTypes/getLeastSupertype.h>
3#include <Functions/FunctionFactory.h>
4
5
6namespace DB
7{
8
9namespace ErrorCodes
10{
11extern const int TOO_FEW_ARGUMENTS_FOR_FUNCTION;
12}
13
14namespace
15{
16
17/// Implements the CASE construction when it is
18/// provided an expression. Users should not call this function.
19class FunctionCaseWithExpression : public IFunction
20{
21public:
22static constexpr auto name = "caseWithExpression";
23static FunctionPtr create(ContextPtr context_) { return std::make_shared<FunctionCaseWithExpression>(context_); }
24
25explicit FunctionCaseWithExpression(ContextPtr context_) : context(context_) {}
26bool isVariadic() const override { return true; }
27bool useDefaultImplementationForConstants() const override { return false; }
28bool useDefaultImplementationForNulls() const override { return false; }
29bool useDefaultImplementationForNothing() const override { return false; }
30bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
31size_t getNumberOfArguments() const override { return 0; }
32String getName() const override { return name; }
33
34DataTypePtr getReturnTypeImpl(const DataTypes & args) const override
35{
36if (args.empty())
37throw Exception(ErrorCodes::TOO_FEW_ARGUMENTS_FOR_FUNCTION, "Function {} expects at least 1 arguments", getName());
38
39/// See the comments in executeImpl() to understand why we actually have to
40/// get the return type of a transform function.
41
42/// Get the types of the arrays that we pass to the transform function.
43DataTypes dst_array_types;
44
45for (size_t i = 2; i < args.size() - 1; i += 2)
46dst_array_types.push_back(args[i]);
47
48// Type of the ELSE branch
49dst_array_types.push_back(args.back());
50
51return getLeastSupertype(dst_array_types);
52}
53
54ColumnPtr executeImpl(const ColumnsWithTypeAndName & args, const DataTypePtr & result_type, size_t input_rows_count) const override
55{
56if (args.empty())
57throw Exception(ErrorCodes::TOO_FEW_ARGUMENTS_FOR_FUNCTION, "Function {} expects at least 1 argument", getName());
58
59/// In the following code, we turn the construction:
60/// CASE expr WHEN val[0] THEN branch[0] ... WHEN val[N-1] then branch[N-1] ELSE branchN
61/// into the construction transform(expr, src, dest, branchN)
62/// where:
63/// src = [val[0], val[1], ..., val[N-1]]
64/// dst = [branch[0], ..., branch[N-1]]
65/// then we perform it.
66
67/// Create the arrays required by the transform function.
68ColumnsWithTypeAndName src_array_elems;
69DataTypes src_array_types;
70
71ColumnsWithTypeAndName dst_array_elems;
72DataTypes dst_array_types;
73
74for (size_t i = 1; i < (args.size() - 1); ++i)
75{
76if (i % 2)
77{
78src_array_elems.push_back(args[i]);
79src_array_types.push_back(args[i].type);
80}
81else
82{
83dst_array_elems.push_back(args[i]);
84dst_array_types.push_back(args[i].type);
85}
86}
87
88DataTypePtr src_array_type = std::make_shared<DataTypeArray>(getLeastSupertype(src_array_types));
89DataTypePtr dst_array_type = std::make_shared<DataTypeArray>(getLeastSupertype(dst_array_types));
90
91ColumnWithTypeAndName src_array_col{nullptr, src_array_type, ""};
92ColumnWithTypeAndName dst_array_col{nullptr, dst_array_type, ""};
93
94auto fun_array = FunctionFactory::instance().get("array", context);
95
96src_array_col.column = fun_array->build(src_array_elems)->execute(src_array_elems, src_array_type, input_rows_count);
97dst_array_col.column = fun_array->build(dst_array_elems)->execute(dst_array_elems, dst_array_type, input_rows_count);
98
99/// Execute transform.
100ColumnsWithTypeAndName transform_args{args.front(), src_array_col, dst_array_col, args.back()};
101return FunctionFactory::instance().get("transform", context)->build(transform_args)
102->execute(transform_args, result_type, input_rows_count);
103}
104
105private:
106ContextPtr context;
107};
108
109}
110
111REGISTER_FUNCTION(CaseWithExpression)
112{
113factory.registerFunction<FunctionCaseWithExpression>();
114
115/// These are obsolete function names.
116factory.registerAlias("caseWithExpr", "caseWithExpression");
117}
118
119}
120