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
10namespace DB
11{
12
13namespace 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)]
20class FunctionArrayEnumerate : public IFunction
21{
22public:
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
80void registerFunctionArrayEnumerate(FunctionFactory & factory)
81{
82 factory.registerFunction<FunctionArrayEnumerate>();
83}
84
85}
86