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
12namespace DB
13{
14
15namespace 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 */
33class FunctionArraySlice : public IFunction
34{
35public:
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
159void registerFunctionArraySlice(FunctionFactory & factory)
160{
161 factory.registerFunction<FunctionArraySlice>();
162}
163
164
165}
166