1 | #include "config_functions.h" |
---|---|
2 | #if USE_H3 |
3 | # include <Columns/ColumnsNumber.h> |
4 | # include <DataTypes/DataTypesNumber.h> |
5 | # include <Functions/FunctionFactory.h> |
6 | # include <Functions/IFunction.h> |
7 | # include <Common/typeid_cast.h> |
8 | # include <ext/range.h> |
9 | |
10 | # include <h3api.h> |
11 | |
12 | |
13 | namespace DB |
14 | { |
15 | class FunctionH3IsValid : public IFunction |
16 | { |
17 | public: |
18 | static constexpr auto name = "h3IsValid"; |
19 | |
20 | static FunctionPtr create(const Context &) { return std::make_shared<FunctionH3IsValid>(); } |
21 | |
22 | std::string getName() const override { return name; } |
23 | |
24 | size_t getNumberOfArguments() const override { return 1; } |
25 | bool useDefaultImplementationForConstants() const override { return true; } |
26 | |
27 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
28 | { |
29 | auto arg = arguments[0].get(); |
30 | if (!WhichDataType(arg).isUInt64()) |
31 | throw Exception( |
32 | "Illegal type "+ arg->getName() + " of argument "+ std::to_string(1) + " of function "+ getName() + ". Must be UInt64", |
33 | ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); |
34 | |
35 | return std::make_shared<DataTypeUInt8>(); |
36 | } |
37 | |
38 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override |
39 | { |
40 | const auto col_hindex = block.getByPosition(arguments[0]).column.get(); |
41 | |
42 | auto dst = ColumnVector<UInt8>::create(); |
43 | auto & dst_data = dst->getData(); |
44 | dst_data.resize(input_rows_count); |
45 | |
46 | for (const auto row : ext::range(0, input_rows_count)) |
47 | { |
48 | const UInt64 hindex = col_hindex->getUInt(row); |
49 | |
50 | UInt8 is_valid = h3IsValid(hindex) == 0 ? 0 : 1; |
51 | |
52 | dst_data[row] = is_valid; |
53 | } |
54 | |
55 | block.getByPosition(result).column = std::move(dst); |
56 | } |
57 | }; |
58 | |
59 | |
60 | void registerFunctionH3IsValid(FunctionFactory & factory) |
61 | { |
62 | factory.registerFunction<FunctionH3IsValid>(); |
63 | } |
64 | |
65 | } |
66 | #endif |
67 |