1 | #include <Functions/IFunctionImpl.h> |
---|---|
2 | #include <Functions/FunctionFactory.h> |
3 | #include <Functions/FunctionHelpers.h> |
4 | #include <DataTypes/DataTypeArray.h> |
5 | #include <Columns/ColumnArray.h> |
6 | |
7 | |
8 | namespace DB |
9 | { |
10 | |
11 | namespace ErrorCodes |
12 | { |
13 | extern const int ILLEGAL_COLUMN; |
14 | } |
15 | |
16 | /** Creates an array, multiplying the column (the first argument) by the number of elements in the array (the second argument). |
17 | */ |
18 | class FunctionReplicate : public IFunction |
19 | { |
20 | public: |
21 | static constexpr auto name = "replicate"; |
22 | |
23 | static FunctionPtr create(const Context &) |
24 | { |
25 | return std::make_shared<FunctionReplicate>(); |
26 | } |
27 | |
28 | String getName() const override |
29 | { |
30 | return name; |
31 | } |
32 | |
33 | size_t getNumberOfArguments() const override |
34 | { |
35 | return 2; |
36 | } |
37 | |
38 | bool useDefaultImplementationForNulls() const override { return false; } |
39 | |
40 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
41 | { |
42 | const DataTypeArray * array_type = checkAndGetDataType<DataTypeArray>(arguments[1].get()); |
43 | if (!array_type) |
44 | throw Exception("Second argument for function "+ getName() + " must be array.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
45 | return std::make_shared<DataTypeArray>(arguments[0]); |
46 | } |
47 | |
48 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t) override |
49 | { |
50 | ColumnPtr first_column = block.getByPosition(arguments[0]).column; |
51 | const ColumnArray * array_column = checkAndGetColumn<ColumnArray>(block.getByPosition(arguments[1]).column.get()); |
52 | ColumnPtr temp_column; |
53 | if (!array_column) |
54 | { |
55 | auto const_array_column = checkAndGetColumnConst<ColumnArray>(block.getByPosition(arguments[1]).column.get()); |
56 | if (!const_array_column) |
57 | throw Exception("Unexpected column for replicate", ErrorCodes::ILLEGAL_COLUMN); |
58 | temp_column = const_array_column->convertToFullColumn(); |
59 | array_column = checkAndGetColumn<ColumnArray>(temp_column.get()); |
60 | } |
61 | block.getByPosition(result).column |
62 | = ColumnArray::create(first_column->replicate(array_column->getOffsets())->convertToFullColumnIfConst(), array_column->getOffsetsPtr()); |
63 | } |
64 | }; |
65 | |
66 | |
67 | void registerFunctionReplicate(FunctionFactory & factory) |
68 | { |
69 | factory.registerFunction<FunctionReplicate>(); |
70 | } |
71 | |
72 | } |
73 |