1 | #include "config_functions.h" |
---|---|
2 | #if USE_H3 |
3 | #include <array> |
4 | #include <math.h> |
5 | #include <Columns/ColumnConst.h> |
6 | #include <Columns/ColumnsNumber.h> |
7 | #include <DataTypes/DataTypesNumber.h> |
8 | #include <Functions/FunctionFactory.h> |
9 | #include <Functions/IFunctionImpl.h> |
10 | #include <Common/typeid_cast.h> |
11 | #include <ext/range.h> |
12 | |
13 | #include <h3api.h> |
14 | |
15 | |
16 | namespace DB |
17 | { |
18 | namespace ErrorCodes |
19 | { |
20 | extern const int ILLEGAL_COLUMN; |
21 | } |
22 | |
23 | /// Implements the function geoToH3 which takes 3 arguments (latitude, longitude and h3 resolution) |
24 | /// and returns h3 index of this point |
25 | class FunctionGeoToH3 : public IFunction |
26 | { |
27 | public: |
28 | static constexpr auto name = "geoToH3"; |
29 | |
30 | static FunctionPtr create(const Context &) { return std::make_shared<FunctionGeoToH3>(); } |
31 | |
32 | std::string getName() const override { return name; } |
33 | |
34 | size_t getNumberOfArguments() const override { return 3; } |
35 | bool useDefaultImplementationForConstants() const override { return true; } |
36 | |
37 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
38 | { |
39 | auto arg = arguments[0].get(); |
40 | if (!WhichDataType(arg).isFloat64()) |
41 | throw Exception( |
42 | "Illegal type "+ arg->getName() + " of argument "+ std::to_string(1) + " of function "+ getName() + ". Must be Float64", |
43 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
44 | |
45 | arg = arguments[1].get(); |
46 | if (!WhichDataType(arg).isFloat64()) |
47 | throw Exception( |
48 | "Illegal type "+ arg->getName() + " of argument "+ std::to_string(2) + " of function "+ getName() + ". Must be Float64", |
49 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
50 | |
51 | arg = arguments[2].get(); |
52 | if (!WhichDataType(arg).isUInt8()) |
53 | throw Exception( |
54 | "Illegal type "+ arg->getName() + " of argument "+ std::to_string(3) + " of function "+ getName() + ". Must be UInt8", |
55 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
56 | |
57 | return std::make_shared<DataTypeUInt64>(); |
58 | } |
59 | |
60 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override |
61 | { |
62 | const auto col_lon = block.getByPosition(arguments[0]).column.get(); |
63 | const auto col_lat = block.getByPosition(arguments[1]).column.get(); |
64 | const auto col_res = block.getByPosition(arguments[2]).column.get(); |
65 | |
66 | auto dst = ColumnVector<UInt64>::create(); |
67 | auto & dst_data = dst->getData(); |
68 | dst_data.resize(input_rows_count); |
69 | |
70 | for (const auto row : ext::range(0, input_rows_count)) |
71 | { |
72 | const double lon = col_lon->getFloat64(row); |
73 | const double lat = col_lat->getFloat64(row); |
74 | const UInt8 res = col_res->getUInt(row); |
75 | |
76 | GeoCoord coord; |
77 | coord.lon = degsToRads(lon); |
78 | coord.lat = degsToRads(lat); |
79 | |
80 | H3Index hindex = geoToH3(&coord, res); |
81 | |
82 | dst_data[row] = hindex; |
83 | } |
84 | |
85 | block.getByPosition(result).column = std::move(dst); |
86 | } |
87 | }; |
88 | |
89 | |
90 | void registerFunctionGeoToH3(FunctionFactory & factory) |
91 | { |
92 | factory.registerFunction<FunctionGeoToH3>(); |
93 | } |
94 | |
95 | } |
96 | #endif |
97 |