ClickHouse

1
#include "config.h"
2

3
#if USE_H3
4

5
#include <Columns/ColumnArray.h>
6
#include <Columns/ColumnsNumber.h>
7
#include <Columns/ColumnTuple.h>
8
#include <DataTypes/DataTypesNumber.h>
9
#include <DataTypes/DataTypeArray.h>
10
#include <DataTypes/DataTypeTuple.h>
11
#include <Functions/FunctionFactory.h>
12
#include <Functions/IFunction.h>
13
#include <Common/typeid_cast.h>
14
#include <IO/WriteHelpers.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

29
class FunctionH3GetUnidirectionalEdgeBoundary : public IFunction
30
{
31
public:
32
    static constexpr auto name = "h3GetUnidirectionalEdgeBoundary";
33

34
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionH3GetUnidirectionalEdgeBoundary>(); }
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 true; }
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<DataTypeArray>(
52
            std::make_shared<DataTypeTuple>(
53
                DataTypes{std::make_shared<DataTypeFloat64>(), std::make_shared<DataTypeFloat64>()}));
54
    }
55

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

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

71
        const auto & data_hindex_edge = col_hindex_edge->getData();
72

73
        auto latitude = ColumnFloat64::create();
74
        auto longitude = ColumnFloat64::create();
75
        auto offsets = DataTypeNumber<IColumn::Offset>().createColumn();
76
        offsets->reserve(input_rows_count);
77
        IColumn::Offset current_offset = 0;
78

79
        for (size_t row = 0; row < input_rows_count; ++row)
80
        {
81
            H3Index edge = data_hindex_edge[row];
82
            CellBoundary boundary{};
83

84
            directedEdgeToBoundary(edge, &boundary);
85

86
            for (int vert = 0; vert < boundary.numVerts; ++vert)
87
            {
88
                latitude->getData().push_back(radsToDegs(boundary.verts[vert].lat));
89
                longitude->getData().push_back(radsToDegs(boundary.verts[vert].lng));
90
            }
91

92
            current_offset += boundary.numVerts;
93
            offsets->insert(current_offset);
94
        }
95

96
        return ColumnArray::create(
97
            ColumnTuple::create(Columns{std::move(latitude), std::move(longitude)}),
98
            std::move(offsets));
99
    }
100
};
101

102
}
103

104
REGISTER_FUNCTION(H3GetUnidirectionalEdgeBoundary)
105
{
106
    factory.registerFunction<FunctionH3GetUnidirectionalEdgeBoundary>();
107
}
108

109
}
110

111
#endif
112