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
9namespace DB
10{
11namespace 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 |
28class FunctionNeighbor : public IFunction
29{
30public:
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
185private:
186 const Context & context;
187};
188
189void registerFunctionNeighbor(FunctionFactory & factory)
190{
191 factory.registerFunction<FunctionNeighbor>();
192}
193
194}
195