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
8namespace DB
9{
10
11namespace 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.
18class FunctionGetSizeOfEnumType : public IFunction
19{
20public:
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
68void registerFunctionGetSizeOfEnumType(FunctionFactory & factory)
69{
70 factory.registerFunction<FunctionGetSizeOfEnumType>();
71}
72
73}
74