1 | #include <Functions/IFunctionImpl.h> |
2 | #include <Functions/FunctionHelpers.h> |
3 | #include <Functions/FunctionFactory.h> |
4 | #include <DataTypes/DataTypeArray.h> |
5 | #include <DataTypes/DataTypesNumber.h> |
6 | #include <Columns/ColumnArray.h> |
7 | |
8 | |
9 | namespace DB |
10 | { |
11 | |
12 | namespace ErrorCodes |
13 | { |
14 | extern const int ILLEGAL_TYPE_OF_ARGUMENT; |
15 | extern const int TOO_LARGE_ARRAY_SIZE; |
16 | } |
17 | |
18 | /// Reasonable threshold. |
19 | static constexpr size_t max_arrays_size_in_block = 1000000000; |
20 | |
21 | |
22 | /* arrayWithConstant(num, const) - make array of constants with length num. |
23 | * arrayWithConstant(3, 'hello') = ['hello', 'hello', 'hello'] |
24 | * arrayWithConstant(1, 'hello') = ['hello'] |
25 | * arrayWithConstant(0, 'hello') = [] |
26 | */ |
27 | |
28 | class FunctionArrayWithConstant : public IFunction |
29 | { |
30 | public: |
31 | static constexpr auto name = "arrayWithConstant" ; |
32 | |
33 | static FunctionPtr create(const Context &) { return std::make_shared<FunctionArrayWithConstant>(); } |
34 | |
35 | String getName() const override { return name; } |
36 | size_t getNumberOfArguments() const override { return 2; } |
37 | |
38 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
39 | { |
40 | if (!isNativeNumber(arguments[0])) |
41 | throw Exception("Illegal type " + arguments[0]->getName() + |
42 | " of argument of function " + getName() + |
43 | ", expected Integer" , ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
44 | return std::make_shared<DataTypeArray>(arguments[1]); |
45 | } |
46 | |
47 | bool useDefaultImplementationForConstants() const override { return true; } |
48 | bool useDefaultImplementationForNulls() const override { return false; } |
49 | |
50 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t num_rows) override |
51 | { |
52 | const auto * col_num = block.getByPosition(arguments[0]).column.get(); |
53 | const auto * col_value = block.getByPosition(arguments[1]).column.get(); |
54 | |
55 | auto offsets_col = ColumnArray::ColumnOffsets::create(); |
56 | ColumnArray::Offsets & offsets = offsets_col->getData(); |
57 | offsets.reserve(num_rows); |
58 | |
59 | ColumnArray::Offset offset = 0; |
60 | for (size_t i = 0; i < num_rows; ++i) |
61 | { |
62 | auto array_size = col_num->getInt(i); |
63 | |
64 | if (unlikely(array_size < 0)) |
65 | throw Exception("Array size cannot be negative: while executing function " + getName(), ErrorCodes::TOO_LARGE_ARRAY_SIZE); |
66 | |
67 | offset += array_size; |
68 | |
69 | if (unlikely(offset > max_arrays_size_in_block)) |
70 | throw Exception("Too large array size while executing function " + getName(), ErrorCodes::TOO_LARGE_ARRAY_SIZE); |
71 | |
72 | offsets.push_back(offset); |
73 | } |
74 | |
75 | block.getByPosition(result).column = ColumnArray::create(col_value->replicate(offsets)->convertToFullColumnIfConst(), std::move(offsets_col)); |
76 | } |
77 | }; |
78 | |
79 | void registerFunctionArrayWithConstant(FunctionFactory & factory) |
80 | { |
81 | factory.registerFunction<FunctionArrayWithConstant>(); |
82 | } |
83 | |
84 | } |
85 | |