1 | #include <Functions/IFunctionImpl.h> |
2 | #include <Functions/FunctionFactory.h> |
3 | #include <DataTypes/DataTypeArray.h> |
4 | #include <DataTypes/getLeastSupertype.h> |
5 | #include <Columns/ColumnArray.h> |
6 | #include <Interpreters/castColumn.h> |
7 | |
8 | |
9 | namespace DB |
10 | { |
11 | |
12 | /// array(c1, c2, ...) - create an array. |
13 | class FunctionArray : public IFunction |
14 | { |
15 | public: |
16 | static constexpr auto name = "array" ; |
17 | static FunctionPtr create(const Context & context) |
18 | { |
19 | return std::make_shared<FunctionArray>(context); |
20 | } |
21 | |
22 | FunctionArray(const Context & context_) |
23 | : context(context_) |
24 | { |
25 | } |
26 | |
27 | bool useDefaultImplementationForNulls() const override { return false; } |
28 | bool useDefaultImplementationForConstants() const override { return true; } |
29 | |
30 | bool isVariadic() const override { return true; } |
31 | size_t getNumberOfArguments() const override { return 0; } |
32 | |
33 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
34 | { |
35 | return std::make_shared<DataTypeArray>(getLeastSupertype(arguments)); |
36 | } |
37 | |
38 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override |
39 | { |
40 | size_t num_elements = arguments.size(); |
41 | |
42 | if (num_elements == 0) |
43 | { |
44 | /// We should return constant empty array. |
45 | block.getByPosition(result).column = block.getByPosition(result).type->createColumnConstWithDefaultValue(input_rows_count); |
46 | return; |
47 | } |
48 | |
49 | const DataTypePtr & return_type = block.getByPosition(result).type; |
50 | const DataTypePtr & elem_type = static_cast<const DataTypeArray &>(*return_type).getNestedType(); |
51 | |
52 | size_t block_size = input_rows_count; |
53 | |
54 | /** If part of columns have not same type as common type of all elements of array, |
55 | * then convert them to common type. |
56 | * If part of columns are constants, |
57 | * then convert them to full columns. |
58 | */ |
59 | |
60 | Columns columns_holder(num_elements); |
61 | ColumnRawPtrs columns(num_elements); |
62 | |
63 | for (size_t i = 0; i < num_elements; ++i) |
64 | { |
65 | const auto & arg = block.getByPosition(arguments[i]); |
66 | |
67 | ColumnPtr preprocessed_column = arg.column; |
68 | |
69 | if (!arg.type->equals(*elem_type)) |
70 | preprocessed_column = castColumn(arg, elem_type, context); |
71 | |
72 | preprocessed_column = preprocessed_column->convertToFullColumnIfConst(); |
73 | |
74 | columns_holder[i] = std::move(preprocessed_column); |
75 | columns[i] = columns_holder[i].get(); |
76 | } |
77 | |
78 | /// Create and fill the result array. |
79 | |
80 | auto out = ColumnArray::create(elem_type->createColumn()); |
81 | IColumn & out_data = out->getData(); |
82 | IColumn::Offsets & out_offsets = out->getOffsets(); |
83 | |
84 | out_data.reserve(block_size * num_elements); |
85 | out_offsets.resize(block_size); |
86 | |
87 | IColumn::Offset current_offset = 0; |
88 | for (size_t i = 0; i < block_size; ++i) |
89 | { |
90 | for (size_t j = 0; j < num_elements; ++j) |
91 | out_data.insertFrom(*columns[j], i); |
92 | |
93 | current_offset += num_elements; |
94 | out_offsets[i] = current_offset; |
95 | } |
96 | |
97 | block.getByPosition(result).column = std::move(out); |
98 | } |
99 | |
100 | private: |
101 | String getName() const override |
102 | { |
103 | return name; |
104 | } |
105 | |
106 | bool addField(DataTypePtr type_res, const Field & f, Array & arr) const; |
107 | |
108 | private: |
109 | const Context & context; |
110 | }; |
111 | |
112 | |
113 | void registerFunctionArray(FunctionFactory & factory) |
114 | { |
115 | factory.registerFunction<FunctionArray>(); |
116 | } |
117 | |
118 | } |
119 | |