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 ILLEGAL_COLUMN;
23
extern const int INCORRECT_DATA;
24
}
25

26
namespace
27
{
28

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

34
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionH3CellAreaM2>(); }
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).isUInt64())
46
            throw Exception(
47
                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
48
                "Illegal type {} of argument {} of function {}. Must be UInt64",
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<ColumnUInt64>(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 UInt64.",
65
                arguments[0].type->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 UInt64 index = data[row];
78

79
            CellBoundary boundary{};
80
            auto err = cellToBoundary(index, &boundary);
81
            if (err)
82
                throw Exception(ErrorCodes::INCORRECT_DATA, "Incorrect H3 index: {}, error: {}", index, err);
83

84
            Float64 res = cellAreaM2(index);
85
            dst_data[row] = res;
86
        }
87

88
        return dst;
89
    }
90
};
91

92
}
93

94
REGISTER_FUNCTION(H3CellAreaM2)
95
{
96
    factory.registerFunction<FunctionH3CellAreaM2>();
97
}
98

99
}
100

101
#endif
102