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
12namespace DB
13{
14
15namespace 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 */
35class FunctionPointInEllipses : public IFunction
36{
37public:
38 static constexpr auto name = "pointInEllipses";
39 static FunctionPtr create(const Context &) { return std::make_shared<FunctionPointInEllipses>(); }
40
41private:
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
191void registerFunctionPointInEllipses(FunctionFactory & factory)
192{
193 factory.registerFunction<FunctionPointInEllipses>();
194}
195
196}
197