1 | #include <Functions/IFunctionImpl.h> |
---|---|
2 | #include <Functions/FunctionFactory.h> |
3 | #include <Functions/FunctionHelpers.h> |
4 | #include <DataTypes/DataTypeEnum.h> |
5 | #include <DataTypes/DataTypesNumber.h> |
6 | |
7 | |
8 | namespace DB |
9 | { |
10 | |
11 | namespace ErrorCodes |
12 | { |
13 | extern const int ILLEGAL_TYPE_OF_ARGUMENT; |
14 | } |
15 | |
16 | |
17 | /// Returns number of fields in Enum data type of passed value. |
18 | class FunctionGetSizeOfEnumType : public IFunction |
19 | { |
20 | public: |
21 | static constexpr auto name = "getSizeOfEnumType"; |
22 | static FunctionPtr create(const Context &) |
23 | { |
24 | return std::make_shared<FunctionGetSizeOfEnumType>(); |
25 | } |
26 | |
27 | String getName() const override |
28 | { |
29 | return name; |
30 | } |
31 | |
32 | bool useDefaultImplementationForNulls() const override { return false; } |
33 | |
34 | size_t getNumberOfArguments() const override |
35 | { |
36 | return 1; |
37 | } |
38 | |
39 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
40 | { |
41 | WhichDataType which(arguments[0]); |
42 | |
43 | if (which.isEnum8()) |
44 | return std::make_shared<DataTypeUInt8>(); |
45 | else if (which.isEnum16()) |
46 | return std::make_shared<DataTypeUInt16>(); |
47 | |
48 | throw Exception("The argument for function "+ getName() + " must be Enum", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
49 | } |
50 | |
51 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override |
52 | { |
53 | block.getByPosition(result).column = getResultIfAlwaysReturnsConstantAndHasArguments(block, arguments)->cloneResized(input_rows_count); |
54 | } |
55 | |
56 | ColumnPtr getResultIfAlwaysReturnsConstantAndHasArguments(const Block & block, const ColumnNumbers & arguments) const override |
57 | { |
58 | if (auto type8 = checkAndGetDataType<DataTypeEnum8>(block.getByPosition(arguments[0]).type.get())) |
59 | return DataTypeUInt8().createColumnConst(1, type8->getValues().size()); |
60 | else if (auto type16 = checkAndGetDataType<DataTypeEnum16>(block.getByPosition(arguments[0]).type.get())) |
61 | return DataTypeUInt16().createColumnConst(1, type16->getValues().size()); |
62 | else |
63 | throw Exception("The argument for function "+ getName() + " must be Enum", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
64 | } |
65 | }; |
66 | |
67 | |
68 | void registerFunctionGetSizeOfEnumType(FunctionFactory & factory) |
69 | { |
70 | factory.registerFunction<FunctionGetSizeOfEnumType>(); |
71 | } |
72 | |
73 | } |
74 |