1 | #include <DataTypes/DataTypesNumber.h> |
2 | #include <Columns/ColumnsNumber.h> |
3 | #include <Columns/ColumnConst.h> |
4 | #include <Common/typeid_cast.h> |
5 | #include <Common/assert_cast.h> |
6 | #include <Functions/IFunctionImpl.h> |
7 | #include <Functions/FunctionHelpers.h> |
8 | #include <Functions/FunctionFactory.h> |
9 | #include <ext/range.h> |
10 | |
11 | |
12 | namespace DB |
13 | { |
14 | |
15 | namespace ErrorCodes |
16 | { |
17 | extern const int TOO_MANY_ARGUMENTS_FOR_FUNCTION; |
18 | extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; |
19 | extern const int ILLEGAL_COLUMN; |
20 | } |
21 | |
22 | /** |
23 | * The function checks if a point is in one of ellipses in set. |
24 | * The number of arguments must be 2 + 4*N where N is the number of ellipses. |
25 | * The arguments must be arranged as follows: (x, y, x_0, y_0, a_0, b_0, ..., x_i, y_i, a_i, b_i) |
26 | * All ellipses parameters must be const values; |
27 | * |
28 | * The function first checks bounding box condition. |
29 | * If a point is inside an ellipse's bounding box, the quadratic condition is evaluated. |
30 | * |
31 | * Here we assume that points in one block are close and are likely to fit in one ellipse, |
32 | * so the last success ellipse index is remembered to check this ellipse first for next point. |
33 | * |
34 | */ |
35 | class FunctionPointInEllipses : public IFunction |
36 | { |
37 | public: |
38 | static constexpr auto name = "pointInEllipses" ; |
39 | static FunctionPtr create(const Context &) { return std::make_shared<FunctionPointInEllipses>(); } |
40 | |
41 | private: |
42 | |
43 | struct Ellipse |
44 | { |
45 | Float64 x; |
46 | Float64 y; |
47 | Float64 a; |
48 | Float64 b; |
49 | }; |
50 | |
51 | String getName() const override { return name; } |
52 | |
53 | bool isVariadic() const override { return true; } |
54 | |
55 | size_t getNumberOfArguments() const override { return 0; } |
56 | |
57 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
58 | { |
59 | if (arguments.size() < 6 || arguments.size() % 4 != 2) |
60 | { |
61 | throw Exception( |
62 | "Incorrect number of arguments of function " + getName() + ". Must be 2 for your point plus 4 * N for ellipses (x_i, y_i, a_i, b_i)." , |
63 | ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH); |
64 | } |
65 | |
66 | /// For array on stack, see below. |
67 | if (arguments.size() > 10000) |
68 | { |
69 | throw Exception( |
70 | "Number of arguments of function " + getName() + " is too large." , ErrorCodes::TOO_MANY_ARGUMENTS_FOR_FUNCTION); |
71 | } |
72 | |
73 | for (const auto arg_idx : ext::range(0, arguments.size())) |
74 | { |
75 | const auto arg = arguments[arg_idx].get(); |
76 | if (!WhichDataType(arg).isFloat64()) |
77 | { |
78 | throw Exception( |
79 | "Illegal type " + arg->getName() + " of argument " + std::to_string(arg_idx + 1) + " of function " + getName() + ". Must be Float64" , |
80 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
81 | } |
82 | } |
83 | |
84 | return std::make_shared<DataTypeUInt8>(); |
85 | } |
86 | |
87 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override |
88 | { |
89 | const auto size = input_rows_count; |
90 | |
91 | /// Prepare array of ellipses. |
92 | size_t ellipses_count = (arguments.size() - 2) / 4; |
93 | std::vector<Ellipse> ellipses(ellipses_count); |
94 | |
95 | for (const auto ellipse_idx : ext::range(0, ellipses_count)) |
96 | { |
97 | Float64 ellipse_data[4]; |
98 | for (const auto idx : ext::range(0, 4)) |
99 | { |
100 | int arg_idx = 2 + 4 * ellipse_idx + idx; |
101 | const auto column = block.getByPosition(arguments[arg_idx]).column.get(); |
102 | if (const auto col = checkAndGetColumnConst<ColumnVector<Float64>>(column)) |
103 | { |
104 | ellipse_data[idx] = col->getValue<Float64>(); |
105 | } |
106 | else |
107 | { |
108 | throw Exception( |
109 | "Illegal type " + column->getName() + " of argument " + std::to_string(arg_idx + 1) + " of function " + getName() + ". Must be const Float64" , |
110 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
111 | } |
112 | } |
113 | ellipses[ellipse_idx] = Ellipse{ellipse_data[0], ellipse_data[1], ellipse_data[2], ellipse_data[3]}; |
114 | } |
115 | |
116 | int const_cnt = 0; |
117 | for (const auto idx : ext::range(0, 2)) |
118 | { |
119 | const auto column = block.getByPosition(arguments[idx]).column.get(); |
120 | if (typeid_cast<const ColumnConst *> (column)) |
121 | { |
122 | ++const_cnt; |
123 | } |
124 | else if (!typeid_cast<const ColumnVector<Float64> *> (column)) |
125 | { |
126 | throw Exception("Illegal column " + column->getName() + " of argument of function " + getName(), |
127 | ErrorCodes::ILLEGAL_COLUMN); |
128 | } |
129 | } |
130 | |
131 | const auto col_x = block.getByPosition(arguments[0]).column.get(); |
132 | const auto col_y = block.getByPosition(arguments[1]).column.get(); |
133 | if (const_cnt == 0) |
134 | { |
135 | const auto col_vec_x = assert_cast<const ColumnVector<Float64> *> (col_x); |
136 | const auto col_vec_y = assert_cast<const ColumnVector<Float64> *> (col_y); |
137 | |
138 | auto dst = ColumnVector<UInt8>::create(); |
139 | auto & dst_data = dst->getData(); |
140 | dst_data.resize(size); |
141 | |
142 | size_t start_index = 0; |
143 | for (const auto row : ext::range(0, size)) |
144 | { |
145 | dst_data[row] = isPointInEllipses(col_vec_x->getData()[row], col_vec_y->getData()[row], ellipses.data(), ellipses_count, start_index); |
146 | } |
147 | |
148 | block.getByPosition(result).column = std::move(dst); |
149 | } |
150 | else if (const_cnt == 2) |
151 | { |
152 | const auto col_const_x = assert_cast<const ColumnConst *> (col_x); |
153 | const auto col_const_y = assert_cast<const ColumnConst *> (col_y); |
154 | size_t start_index = 0; |
155 | UInt8 res = isPointInEllipses(col_const_x->getValue<Float64>(), col_const_y->getValue<Float64>(), ellipses.data(), ellipses_count, start_index); |
156 | block.getByPosition(result).column = DataTypeUInt8().createColumnConst(size, res); |
157 | } |
158 | else |
159 | { |
160 | throw Exception( |
161 | "Illegal types " + col_x->getName() + ", " + col_y->getName() + " of arguments 1, 2 of function " + getName() + ". Both must be either const or vector" , |
162 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
163 | } |
164 | } |
165 | |
166 | static bool isPointInEllipses(Float64 x, Float64 y, const Ellipse * ellipses, size_t ellipses_count, size_t & start_index) |
167 | { |
168 | size_t index = 0 + start_index; |
169 | for (size_t i = 0; i < ellipses_count; ++i) |
170 | { |
171 | Ellipse el = ellipses[index]; |
172 | double p1 = ((x - el.x) / el.a); |
173 | double p2 = ((y - el.y) / el.b); |
174 | if (x <= el.x + el.a && x >= el.x - el.a && y <= el.y + el.b && y >= el.y - el.b /// Bounding box check |
175 | && p1 * p1 + p2 * p2 <= 1.0) /// Precise check |
176 | { |
177 | start_index = index; |
178 | return true; |
179 | } |
180 | ++index; |
181 | if (index == ellipses_count) |
182 | { |
183 | index = 0; |
184 | } |
185 | } |
186 | return false; |
187 | } |
188 | }; |
189 | |
190 | |
191 | void registerFunctionPointInEllipses(FunctionFactory & factory) |
192 | { |
193 | factory.registerFunction<FunctionPointInEllipses>(); |
194 | } |
195 | |
196 | } |
197 | |