1 | #include <Functions/IFunctionImpl.h> |
2 | #include <Functions/FunctionFactory.h> |
3 | #include <Functions/GatherUtils/GatherUtils.h> |
4 | #include <DataTypes/DataTypeArray.h> |
5 | #include <DataTypes/DataTypeNullable.h> |
6 | #include <Columns/ColumnArray.h> |
7 | #include <Columns/ColumnConst.h> |
8 | #include <Common/typeid_cast.h> |
9 | #include <IO/WriteHelpers.h> |
10 | |
11 | |
12 | namespace DB |
13 | { |
14 | |
15 | namespace ErrorCodes |
16 | { |
17 | extern const int LOGICAL_ERROR; |
18 | extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; |
19 | extern const int ILLEGAL_TYPE_OF_ARGUMENT; |
20 | } |
21 | |
22 | /** arraySlice(arr, offset, length) - make slice of array. Offsets and length may be < 0 or Null |
23 | * - if offset < 0, indexation from right element |
24 | * - if length < 0, length = len(array) - (positive_index(offset) - 1) + length |
25 | * indexation: |
26 | * [ 1, 2, 3, 4, 5, 6] |
27 | * [-6, -5, -4, -3, -2, -1] |
28 | * examples: |
29 | * arraySlice([1, 2, 3, 4, 5, 6], -4, 2) -> [3, 4] |
30 | * arraySlice([1, 2, 3, 4, 5, 6], 2, -1) -> [2, 3, 4, 5] (6 - (2 - 1) + (-1) = 4) |
31 | * arraySlice([1, 2, 3, 4, 5, 6], -5, -1) = arraySlice([1, 2, 3, 4, 5, 6], 2, -1) -> [2, 3, 4, 5] |
32 | */ |
33 | class FunctionArraySlice : public IFunction |
34 | { |
35 | public: |
36 | static constexpr auto name = "arraySlice" ; |
37 | static FunctionPtr create(const Context &) { return std::make_shared<FunctionArraySlice>(); } |
38 | |
39 | String getName() const override { return name; } |
40 | |
41 | bool isVariadic() const override { return true; } |
42 | size_t getNumberOfArguments() const override { return 0; } |
43 | |
44 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
45 | { |
46 | const size_t number_of_arguments = arguments.size(); |
47 | |
48 | if (number_of_arguments < 2 || number_of_arguments > 3) |
49 | throw Exception("Number of arguments for function " + getName() + " doesn't match: passed " |
50 | + toString(number_of_arguments) + ", should be 2 or 3" , |
51 | ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH); |
52 | |
53 | if (arguments[0]->onlyNull()) |
54 | return arguments[0]; |
55 | |
56 | auto array_type = typeid_cast<const DataTypeArray *>(arguments[0].get()); |
57 | if (!array_type) |
58 | throw Exception("First argument for function " + getName() + " must be an array but it has type " |
59 | + arguments[0]->getName() + "." , ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
60 | |
61 | for (size_t i = 1; i < number_of_arguments; ++i) |
62 | { |
63 | if (!isInteger(removeNullable(arguments[i])) && !arguments[i]->onlyNull()) |
64 | throw Exception( |
65 | "Argument " + toString(i) + " for function " + getName() + " must be integer but it has type " |
66 | + arguments[i]->getName() + "." , ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
67 | } |
68 | |
69 | return arguments[0]; |
70 | } |
71 | |
72 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override |
73 | { |
74 | const auto & return_type = block.getByPosition(result).type; |
75 | |
76 | if (return_type->onlyNull()) |
77 | { |
78 | block.getByPosition(result).column = return_type->createColumnConstWithDefaultValue(input_rows_count); |
79 | return; |
80 | } |
81 | |
82 | auto result_column = return_type->createColumn(); |
83 | |
84 | auto & array_column = block.getByPosition(arguments[0]).column; |
85 | const auto & offset_column = block.getByPosition(arguments[1]).column; |
86 | const auto & length_column = arguments.size() > 2 ? block.getByPosition(arguments[2]).column : nullptr; |
87 | |
88 | std::unique_ptr<GatherUtils::IArraySource> source; |
89 | |
90 | size_t size = array_column->size(); |
91 | bool is_const = false; |
92 | |
93 | if (auto const_array_column = typeid_cast<const ColumnConst *>(array_column.get())) |
94 | { |
95 | is_const = true; |
96 | array_column = const_array_column->getDataColumnPtr(); |
97 | } |
98 | |
99 | if (auto argument_column_array = typeid_cast<const ColumnArray *>(array_column.get())) |
100 | source = GatherUtils::createArraySource(*argument_column_array, is_const, size); |
101 | else |
102 | throw Exception{"First arguments for function " + getName() + " must be array." , ErrorCodes::LOGICAL_ERROR}; |
103 | |
104 | auto sink = GatherUtils::createArraySink(typeid_cast<ColumnArray &>(*result_column), size); |
105 | |
106 | if (offset_column->onlyNull()) |
107 | { |
108 | if (!length_column || length_column->onlyNull()) |
109 | { |
110 | block.getByPosition(result).column = array_column; |
111 | return; |
112 | } |
113 | else if (isColumnConst(*length_column)) |
114 | GatherUtils::sliceFromLeftConstantOffsetBounded(*source, *sink, 0, length_column->getInt(0)); |
115 | else |
116 | { |
117 | auto const_offset_column = ColumnConst::create(ColumnInt8::create(1, 1), size); |
118 | GatherUtils::sliceDynamicOffsetBounded(*source, *sink, *const_offset_column, *length_column); |
119 | } |
120 | } |
121 | else if (isColumnConst(*offset_column)) |
122 | { |
123 | ssize_t offset = offset_column->getUInt(0); |
124 | |
125 | if (!length_column || length_column->onlyNull()) |
126 | { |
127 | if (offset > 0) |
128 | GatherUtils::sliceFromLeftConstantOffsetUnbounded(*source, *sink, static_cast<size_t>(offset - 1)); |
129 | else |
130 | GatherUtils::sliceFromRightConstantOffsetUnbounded(*source, *sink, static_cast<size_t>(-offset)); |
131 | } |
132 | else if (isColumnConst(*length_column)) |
133 | { |
134 | ssize_t length = length_column->getInt(0); |
135 | if (offset > 0) |
136 | GatherUtils::sliceFromLeftConstantOffsetBounded(*source, *sink, static_cast<size_t>(offset - 1), length); |
137 | else |
138 | GatherUtils::sliceFromRightConstantOffsetBounded(*source, *sink, static_cast<size_t>(-offset), length); |
139 | } |
140 | else |
141 | GatherUtils::sliceDynamicOffsetBounded(*source, *sink, *offset_column, *length_column); |
142 | } |
143 | else |
144 | { |
145 | if (!length_column || length_column->onlyNull()) |
146 | GatherUtils::sliceDynamicOffsetUnbounded(*source, *sink, *offset_column); |
147 | else |
148 | GatherUtils::sliceDynamicOffsetBounded(*source, *sink, *offset_column, *length_column); |
149 | } |
150 | |
151 | block.getByPosition(result).column = std::move(result_column); |
152 | } |
153 | |
154 | bool useDefaultImplementationForConstants() const override { return true; } |
155 | bool useDefaultImplementationForNulls() const override { return false; } |
156 | }; |
157 | |
158 | |
159 | void registerFunctionArraySlice(FunctionFactory & factory) |
160 | { |
161 | factory.registerFunction<FunctionArraySlice>(); |
162 | } |
163 | |
164 | |
165 | } |
166 | |