ClickHouse

1
#include "config.h"
2

3
#if USE_H3
4

5
#include <Columns/ColumnsNumber.h>
6
#include <DataTypes/DataTypesNumber.h>
7
#include <Functions/FunctionFactory.h>
8
#include <Functions/IFunction.h>
9
#include <IO/WriteHelpers.h>
10
#include <Common/typeid_cast.h>
11
#include <base/range.h>
12

13
#include <constants.h>
14
#include <h3api.h>
15

16

17
namespace DB
18
{
19
namespace ErrorCodes
20
{
21
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
22
extern const int ARGUMENT_OUT_OF_BOUND;
23
extern const int ILLEGAL_COLUMN;
24
}
25

26
namespace
27
{
28

29
class FunctionH3HexAreaKm2 final : public IFunction
30
{
31
public:
32
    static constexpr auto name = "h3HexAreaKm2";
33

34
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionH3HexAreaKm2>(); }
35

36
    std::string getName() const override { return name; }
37

38
    size_t getNumberOfArguments() const override { return 1; }
39
    bool useDefaultImplementationForConstants() const override { return true; }
40
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
41

42
    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
43
    {
44
        const auto * arg = arguments[0].get();
45
        if (!WhichDataType(arg).isUInt8())
46
            throw Exception(
47
                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
48
                "Illegal type {} of argument {} of function {}. Must be UInt8",
49
                arg->getName(), 1, getName());
50

51
        return std::make_shared<DataTypeFloat64>();
52
    }
53

54
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
55
    {
56
        auto non_const_arguments = arguments;
57
        for (auto & argument : non_const_arguments)
58
            argument.column = argument.column->convertToFullColumnIfConst();
59

60
        const auto * column = checkAndGetColumn<ColumnUInt8>(non_const_arguments[0].column.get());
61
        if (!column)
62
            throw Exception(
63
                ErrorCodes::ILLEGAL_COLUMN,
64
                "Illegal type {} of argument {} of function {}. Must be UInt8",
65
                arguments[0].column->getName(),
66
                1,
67
                getName());
68

69
        const auto & data = column->getData();
70

71
        auto dst = ColumnVector<Float64>::create();
72
        auto & dst_data = dst->getData();
73
        dst_data.resize(input_rows_count);
74

75
        for (size_t row = 0; row < input_rows_count; ++row)
76
        {
77
            const UInt8 resolution = data[row];
78
            if (resolution > MAX_H3_RES)
79
                throw Exception(
80
                    ErrorCodes::ARGUMENT_OUT_OF_BOUND,
81
                    "The argument 'resolution' ({}) of function {} is out of bounds because the maximum resolution in H3 library is {}",
82
                    toString(resolution),
83
                    getName(),
84
                    MAX_H3_RES);
85

86
            Float64 res = getHexagonAreaAvgKm2(resolution);
87
            dst_data[row] = res;
88
        }
89

90
        return dst;
91
    }
92
};
93

94
}
95

96
REGISTER_FUNCTION(H3HexAreaKm2)
97
{
98
    factory.registerFunction<FunctionH3HexAreaKm2>();
99
}
100

101
}
102

103
#endif
104