ClickHouse

1
#include "config.h"
2

3
#if USE_H3
4

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

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

17

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

26
namespace
27
{
28
template <class Impl>
29
class FunctionH3PointDist final : public IFunction
30
{
31
public:
32
    static constexpr auto name = Impl::name;
33
    static constexpr auto function = Impl::function;
34

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

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

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

43
    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
44
    {
45
        for (size_t i = 0; i < getNumberOfArguments(); ++i)
46
        {
47
            const auto * arg = arguments[i].get();
48
            if (!WhichDataType(arg).isFloat64())
49
                throw Exception(
50
                        ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
51
                        "Illegal type {} of argument {} of function {}. Must be Float64",
52
                        arg->getName(), i, getName());
53
        }
54
        return std::make_shared<DataTypeFloat64>();
55
    }
56

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

63
        const auto * col_lat1 = checkAndGetColumn<ColumnFloat64>(non_const_arguments[0].column.get());
64
        if (!col_lat1)
65
            throw Exception(
66
                    ErrorCodes::ILLEGAL_COLUMN,
67
                    "Illegal type {} of argument {} of function {}. Must be Float64",
68
                    arguments[0].type->getName(),
69
                    1,
70
                    getName());
71
        const auto & data_lat1 = col_lat1->getData();
72

73
        const auto * col_lon1 = checkAndGetColumn<ColumnFloat64>(non_const_arguments[1].column.get());
74
        if (!col_lon1)
75
            throw Exception(
76
                    ErrorCodes::ILLEGAL_COLUMN,
77
                    "Illegal type {} of argument {} of function {}. Must be Float64",
78
                    arguments[1].type->getName(),
79
                    2,
80
                    getName());
81
        const auto & data_lon1 = col_lon1->getData();
82

83
        const auto * col_lat2 = checkAndGetColumn<ColumnFloat64>(non_const_arguments[2].column.get());
84
        if (!col_lat2)
85
            throw Exception(
86
                    ErrorCodes::ILLEGAL_COLUMN,
87
                    "Illegal type {} of argument {} of function {}. Must be Float64",
88
                    arguments[2].type->getName(),
89
                    3,
90
                    getName());
91
        const auto & data_lat2 = col_lat2->getData();
92

93
        const auto * col_lon2 = checkAndGetColumn<ColumnFloat64>(non_const_arguments[3].column.get());
94
        if (!col_lon2)
95
            throw Exception(
96
                    ErrorCodes::ILLEGAL_COLUMN,
97
                    "Illegal type {} of argument {} of function {}. Must be Float64",
98
                    arguments[3].type->getName(),
99
                    4,
100
                    getName());
101
        const auto & data_lon2 = col_lon2->getData();
102

103
        auto dst = ColumnVector<Float64>::create();
104
        auto & dst_data = dst->getData();
105
        dst_data.resize(input_rows_count);
106

107
        for (size_t row = 0; row < input_rows_count; ++row)
108
        {
109
            const double lat1 = data_lat1[row];
110
            const double lon1 = data_lon1[row];
111
            const auto lat2 = data_lat2[row];
112
            const auto lon2 = data_lon2[row];
113

114
            LatLng point1 = {degsToRads(lat1), degsToRads(lon1)};
115
            LatLng point2 = {degsToRads(lat2), degsToRads(lon2)};
116

117
            // function will be equivalent to distanceM or distanceKm or distanceRads
118
            Float64 res = function(&point1, &point2);
119
            dst_data[row] = res;
120
        }
121

122
        return dst;
123
    }
124
};
125

126
}
127

128
struct H3PointDistM
129
{
130
    static constexpr auto name = "h3PointDistM";
131
    static constexpr auto function = distanceM;
132
};
133

134
struct H3PointDistKm
135
{
136
    static constexpr auto name = "h3PointDistKm";
137
    static constexpr auto function = distanceKm;
138
};
139

140
struct H3PointDistRads
141
{
142
    static constexpr auto name = "h3PointDistRads";
143
    static constexpr auto function = distanceRads;
144
};
145

146

147
REGISTER_FUNCTION(H3PointDistM) { factory.registerFunction<FunctionH3PointDist<H3PointDistM>>(); }
148
REGISTER_FUNCTION(H3PointDistKm) { factory.registerFunction<FunctionH3PointDist<H3PointDistKm>>(); }
149
REGISTER_FUNCTION(H3PointDistRads) { factory.registerFunction<FunctionH3PointDist<H3PointDistRads>>(); }
150

151
}
152

153
#endif
154