1 | #include <Functions/IFunctionImpl.h> |
2 | #include <Functions/FunctionFactory.h> |
3 | #include <Functions/FunctionHelpers.h> |
4 | #include <DataTypes/DataTypeArray.h> |
5 | #include <DataTypes/DataTypesNumber.h> |
6 | #include <Columns/ColumnArray.h> |
7 | #include <Columns/ColumnsNumber.h> |
8 | |
9 | |
10 | namespace DB |
11 | { |
12 | |
13 | namespace ErrorCodes |
14 | { |
15 | extern const int ILLEGAL_COLUMN; |
16 | extern const int ILLEGAL_TYPE_OF_ARGUMENT; |
17 | } |
18 | |
19 | /// arrayEnumerate(arr) - Returns the array [1,2,3,..., length(arr)] |
20 | class FunctionArrayEnumerate : public IFunction |
21 | { |
22 | public: |
23 | static constexpr auto name = "arrayEnumerate" ; |
24 | |
25 | static FunctionPtr create(const Context &) |
26 | { |
27 | return std::make_shared<FunctionArrayEnumerate>(); |
28 | } |
29 | |
30 | String getName() const override |
31 | { |
32 | return name; |
33 | } |
34 | |
35 | size_t getNumberOfArguments() const override { return 1; } |
36 | bool useDefaultImplementationForConstants() const override { return true; } |
37 | |
38 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
39 | { |
40 | const DataTypeArray * array_type = checkAndGetDataType<DataTypeArray>(arguments[0].get()); |
41 | if (!array_type) |
42 | throw Exception("First argument for function " + getName() + " must be an array but it has type " |
43 | + arguments[0]->getName() + "." , ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
44 | |
45 | return std::make_shared<DataTypeArray>(std::make_shared<DataTypeUInt32>()); |
46 | } |
47 | |
48 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t) override |
49 | { |
50 | if (const ColumnArray * array = checkAndGetColumn<ColumnArray>(block.getByPosition(arguments[0]).column.get())) |
51 | { |
52 | const ColumnArray::Offsets & offsets = array->getOffsets(); |
53 | |
54 | auto res_nested = ColumnUInt32::create(); |
55 | |
56 | ColumnUInt32::Container & res_values = res_nested->getData(); |
57 | res_values.resize(array->getData().size()); |
58 | ColumnArray::Offset prev_off = 0; |
59 | for (ColumnArray::Offset i = 0; i < offsets.size(); ++i) |
60 | { |
61 | ColumnArray::Offset off = offsets[i]; |
62 | for (ColumnArray::Offset j = prev_off; j < off; ++j) |
63 | res_values[j] = j - prev_off + 1; |
64 | |
65 | prev_off = off; |
66 | } |
67 | |
68 | block.getByPosition(result).column = ColumnArray::create(std::move(res_nested), array->getOffsetsPtr()); |
69 | } |
70 | else |
71 | { |
72 | throw Exception("Illegal column " + block.getByPosition(arguments[0]).column->getName() |
73 | + " of first argument of function " + getName(), |
74 | ErrorCodes::ILLEGAL_COLUMN); |
75 | } |
76 | } |
77 | }; |
78 | |
79 | |
80 | void registerFunctionArrayEnumerate(FunctionFactory & factory) |
81 | { |
82 | factory.registerFunction<FunctionArrayEnumerate>(); |
83 | } |
84 | |
85 | } |
86 | |