1 | #include <Columns/ColumnConst.h> |
2 | #include <DataTypes/DataTypesNumber.h> |
3 | #include <DataTypes/getLeastSupertype.h> |
4 | #include <Functions/FunctionFactory.h> |
5 | #include <Functions/IFunctionImpl.h> |
6 | #include <IO/WriteHelpers.h> |
7 | #include <Interpreters/castColumn.h> |
8 | |
9 | namespace DB |
10 | { |
11 | namespace ErrorCodes |
12 | { |
13 | extern const int ILLEGAL_COLUMN; |
14 | extern const int ILLEGAL_TYPE_OF_ARGUMENT; |
15 | extern const int ARGUMENT_OUT_OF_BOUND; |
16 | extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; |
17 | } |
18 | |
19 | // Implements function, giving value for column within range of given |
20 | // Example: |
21 | // | c1 | |
22 | // | 10 | |
23 | // | 20 | |
24 | // SELECT c1, neighbor(c1, 1) as c2: |
25 | // | c1 | c2 | |
26 | // | 10 | 20 | |
27 | // | 20 | 0 | |
28 | class FunctionNeighbor : public IFunction |
29 | { |
30 | public: |
31 | static constexpr auto name = "neighbor" ; |
32 | static FunctionPtr create(const Context & context) { return std::make_shared<FunctionNeighbor>(context); } |
33 | |
34 | FunctionNeighbor(const Context & context_) : context(context_) {} |
35 | |
36 | /// Get the name of the function. |
37 | String getName() const override { return name; } |
38 | |
39 | size_t getNumberOfArguments() const override { return 0; } |
40 | |
41 | bool isVariadic() const override { return true; } |
42 | |
43 | bool isDeterministic() const override { return false; } |
44 | |
45 | bool isDeterministicInScopeOfQuery() const override { return false; } |
46 | |
47 | bool useDefaultImplementationForNulls() const override { return false; } |
48 | |
49 | bool useDefaultImplementationForConstants() const override { return false; } |
50 | |
51 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
52 | { |
53 | size_t number_of_arguments = arguments.size(); |
54 | |
55 | if (number_of_arguments < 2 || number_of_arguments > 3) |
56 | throw Exception( |
57 | "Number of arguments for function " + getName() + " doesn't match: passed " + toString(number_of_arguments) |
58 | + ", should be from 2 to 3" , |
59 | ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH); |
60 | |
61 | // second argument must be an integer |
62 | if (!isInteger(arguments[1])) |
63 | throw Exception( |
64 | "Illegal type " + arguments[1]->getName() + " of second argument of function " + getName() + " - should be an integer" , |
65 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
66 | else if (arguments[1]->isNullable()) |
67 | throw Exception( |
68 | "Illegal type " + arguments[1]->getName() + " of second argument of function " + getName() + " - can not be Nullable" , |
69 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
70 | |
71 | // check that default value column has supertype with first argument |
72 | if (number_of_arguments == 3) |
73 | return getLeastSupertype({arguments[0], arguments[2]}); |
74 | |
75 | return arguments[0]; |
76 | } |
77 | |
78 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override |
79 | { |
80 | const DataTypePtr & result_type = block.getByPosition(result).type; |
81 | |
82 | const ColumnWithTypeAndName & source_elem = block.getByPosition(arguments[0]); |
83 | const ColumnWithTypeAndName & offset_elem = block.getByPosition(arguments[1]); |
84 | bool has_defaults = arguments.size() == 3; |
85 | |
86 | ColumnPtr source_column_casted = castColumn(source_elem, result_type, context); |
87 | ColumnPtr offset_column = offset_elem.column; |
88 | |
89 | ColumnPtr default_column_casted; |
90 | if (has_defaults) |
91 | { |
92 | const ColumnWithTypeAndName & default_elem = block.getByPosition(arguments[2]); |
93 | default_column_casted = castColumn(default_elem, result_type, context); |
94 | } |
95 | |
96 | bool source_is_constant = isColumnConst(*source_column_casted); |
97 | bool offset_is_constant = isColumnConst(*offset_column); |
98 | |
99 | bool default_is_constant = false; |
100 | if (has_defaults) |
101 | default_is_constant = isColumnConst(*default_column_casted); |
102 | |
103 | if (source_is_constant) |
104 | source_column_casted = assert_cast<const ColumnConst &>(*source_column_casted).getDataColumnPtr(); |
105 | if (offset_is_constant) |
106 | offset_column = assert_cast<const ColumnConst &>(*offset_column).getDataColumnPtr(); |
107 | if (default_is_constant) |
108 | default_column_casted = assert_cast<const ColumnConst &>(*default_column_casted).getDataColumnPtr(); |
109 | |
110 | if (offset_is_constant) |
111 | { |
112 | /// Optimization for the case when we can copy many values at once. |
113 | |
114 | Int64 offset = offset_column->getInt(0); |
115 | |
116 | auto result_column = result_type->createColumn(); |
117 | |
118 | auto insert_range_from = [&](bool is_const, const ColumnPtr & src, Int64 begin, Int64 size) |
119 | { |
120 | /// Saturation of bounds. |
121 | if (begin < 0) |
122 | { |
123 | size += begin; |
124 | begin = 0; |
125 | } |
126 | if (size <= 0) |
127 | return; |
128 | if (size > Int64(input_rows_count)) |
129 | size = input_rows_count; |
130 | |
131 | if (!src) |
132 | { |
133 | for (Int64 i = 0; i < size; ++i) |
134 | result_column->insertDefault(); |
135 | } |
136 | else if (is_const) |
137 | { |
138 | for (Int64 i = 0; i < size; ++i) |
139 | result_column->insertFrom(*src, 0); |
140 | } |
141 | else |
142 | { |
143 | result_column->insertRangeFrom(*src, begin, size); |
144 | } |
145 | }; |
146 | |
147 | if (offset == 0) |
148 | { |
149 | /// Degenerate case, just copy source column as is. |
150 | block.getByPosition(result).column = source_is_constant ? ColumnConst::create(source_column_casted, input_rows_count) : source_column_casted; |
151 | } |
152 | else if (offset > 0) |
153 | { |
154 | insert_range_from(source_is_constant, source_column_casted, offset, Int64(input_rows_count) - offset); |
155 | insert_range_from(default_is_constant, default_column_casted, Int64(input_rows_count) - offset, offset); |
156 | block.getByPosition(result).column = std::move(result_column); |
157 | } |
158 | else |
159 | { |
160 | insert_range_from(default_is_constant, default_column_casted, 0, -offset); |
161 | insert_range_from(source_is_constant, source_column_casted, 0, Int64(input_rows_count) + offset); |
162 | block.getByPosition(result).column = std::move(result_column); |
163 | } |
164 | } |
165 | else |
166 | { |
167 | auto result_column = result_type->createColumn(); |
168 | |
169 | for (size_t row = 0; row < input_rows_count; ++row) |
170 | { |
171 | Int64 src_idx = row + offset_column->getInt(offset_is_constant ? 0 : row); |
172 | |
173 | if (src_idx >= 0 && src_idx < Int64(input_rows_count)) |
174 | result_column->insertFrom(*source_column_casted, source_is_constant ? 0 : src_idx); |
175 | else if (has_defaults) |
176 | result_column->insertFrom(*default_column_casted, default_is_constant ? 0 : row); |
177 | else |
178 | result_column->insertDefault(); |
179 | } |
180 | |
181 | block.getByPosition(result).column = std::move(result_column); |
182 | } |
183 | } |
184 | |
185 | private: |
186 | const Context & context; |
187 | }; |
188 | |
189 | void registerFunctionNeighbor(FunctionFactory & factory) |
190 | { |
191 | factory.registerFunction<FunctionNeighbor>(); |
192 | } |
193 | |
194 | } |
195 | |