ClickHouse

1
#include "config.h"
2

3
#if USE_S2_GEOMETRY
4

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

14
#include "s2_fwd.h"
15

16
namespace DB
17
{
18

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

26
namespace
27
{
28

29
/**
30
 * The cap represents a portion of the sphere that has been cut off by a plane.
31
 * See comment for s2CapContains function.
32
 * This function returns the smallest cap that contains both of input caps.
33
 * It is represented by identifier of the center and a radius.
34
 */
35
class FunctionS2CapUnion : public IFunction
36
{
37
public:
38
    static constexpr auto name = "s2CapUnion";
39

40
    static FunctionPtr create(ContextPtr)
41
    {
42
        return std::make_shared<FunctionS2CapUnion>();
43
    }
44

45
    std::string getName() const override
46
    {
47
        return name;
48
    }
49

50
    size_t getNumberOfArguments() const override { return 4; }
51

52
    bool useDefaultImplementationForConstants() const override { return true; }
53

54
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
55

56
    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
57
    {
58
        for (size_t index = 0; index < getNumberOfArguments(); ++index)
59
        {
60
            const auto * arg = arguments[index].get();
61
            if (index == 1 || index == 3)
62
            {
63
                if (!WhichDataType(arg).isFloat64())
64
                    throw Exception(
65
                        ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
66
                        "Illegal type {} of argument {} of function {}. Must be Float64",
67
                        arg->getName(), index + 1, getName());
68
            }
69
            else if (!WhichDataType(arg).isUInt64())
70
                throw Exception(
71
                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
72
                    "Illegal type {} of argument {} of function {}. Must be UInt64",
73
                    arg->getName(), index + 1, getName()
74
                    );
75
        }
76

77
        DataTypePtr center = std::make_shared<DataTypeUInt64>();
78
        DataTypePtr radius = std::make_shared<DataTypeFloat64>();
79

80
        return std::make_shared<DataTypeTuple>(DataTypes{center, radius});
81
    }
82

83
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
84
    {
85
        auto non_const_arguments = arguments;
86
        for (auto & argument : non_const_arguments)
87
            argument.column = argument.column->convertToFullColumnIfConst();
88

89
        const auto * col_center1 = checkAndGetColumn<ColumnUInt64>(non_const_arguments[0].column.get());
90
        if (!col_center1)
91
            throw Exception(
92
                ErrorCodes::ILLEGAL_COLUMN,
93
                "Illegal type {} of argument {} of function {}. Must be UInt64",
94
                arguments[0].type->getName(),
95
                1,
96
                getName());
97
        const auto & data_center1 = col_center1->getData();
98

99
        const auto * col_radius1 = checkAndGetColumn<ColumnFloat64>(non_const_arguments[1].column.get());
100
        if (!col_radius1)
101
            throw Exception(
102
                ErrorCodes::ILLEGAL_COLUMN,
103
                "Illegal type {} of argument {} of function {}. Must be Float64",
104
                arguments[1].type->getName(),
105
                2,
106
                getName());
107
        const auto & data_radius1 = col_radius1->getData();
108

109
        const auto * col_center2 = checkAndGetColumn<ColumnUInt64>(non_const_arguments[2].column.get());
110
        if (!col_center2)
111
            throw Exception(
112
                ErrorCodes::ILLEGAL_COLUMN,
113
                "Illegal type {} of argument {} of function {}. Must be UInt64",
114
                arguments[2].type->getName(),
115
                3,
116
                getName());
117
        const auto & data_center2 = col_center2->getData();
118

119
        const auto * col_radius2 = checkAndGetColumn<ColumnFloat64>(non_const_arguments[3].column.get());
120
        if (!col_radius2)
121
            throw Exception(
122
                ErrorCodes::ILLEGAL_COLUMN,
123
                "Illegal type {} of argument {} of function {}. Must be Float64",
124
                arguments[3].type->getName(),
125
                4,
126
                getName());
127
        const auto & data_radius2 = col_radius2->getData();
128

129
        auto col_res_center = ColumnUInt64::create();
130
        auto col_res_radius = ColumnFloat64::create();
131

132
        auto & vec_res_center = col_res_center->getData();
133
        vec_res_center.reserve(input_rows_count);
134

135
        auto & vec_res_radius = col_res_radius->getData();
136
        vec_res_radius.reserve(input_rows_count);
137

138
        for (size_t row = 0; row < input_rows_count; ++row)
139
        {
140
            const UInt64 first_center = data_center1[row];
141
            const Float64 first_radius = data_radius1[row];
142
            const UInt64 second_center = data_center2[row];
143
            const Float64 second_radius = data_radius2[row];
144

145
            if (isNaN(first_radius) || isNaN(second_radius))
146
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Radius of the cap must not be nan");
147

148
            if (std::isinf(first_radius) || std::isinf(second_radius))
149
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Radius of the cap must not be infinite");
150

151
            auto first_center_cell = S2CellId(first_center);
152
            auto second_center_cell = S2CellId(second_center);
153

154
            if (!first_center_cell.is_valid() || !second_center_cell.is_valid())
155
                throw Exception(ErrorCodes::BAD_ARGUMENTS, "Center of the cap is not valid");
156

157
            S2Cap cap1(first_center_cell.ToPoint(), S1Angle::Degrees(first_radius));
158
            S2Cap cap2(second_center_cell.ToPoint(), S1Angle::Degrees(second_radius));
159

160
            S2Cap cap_union = cap1.Union(cap2);
161

162
            vec_res_center.emplace_back(S2CellId(cap_union.center()).id());
163
            vec_res_radius.emplace_back(cap_union.GetRadius().degrees());
164
        }
165

166
        return ColumnTuple::create(Columns{std::move(col_res_center), std::move(col_res_radius)});
167
    }
168
};
169

170
}
171

172
REGISTER_FUNCTION(S2CapUnion)
173
{
174
    factory.registerFunction<FunctionS2CapUnion>();
175
}
176

177
}
178

179
#endif
180