ClickHouse
127 строк · 4.2 Кб
1#include "config.h"
2
3#if USE_H3
4
5#include <Columns/ColumnArray.h>
6#include <Columns/ColumnsNumber.h>
7#include <DataTypes/DataTypeArray.h>
8#include <DataTypes/DataTypesNumber.h>
9#include <Functions/FunctionFactory.h>
10#include <Functions/IFunction.h>
11#include <Common/typeid_cast.h>
12#include <IO/WriteHelpers.h>
13#include <h3api.h>
14
15
16namespace DB
17{
18namespace ErrorCodes
19{
20extern const int ILLEGAL_TYPE_OF_ARGUMENT;
21extern const int ILLEGAL_COLUMN;
22}
23
24namespace
25{
26
27class FunctionH3GetUnidirectionalEdge : public IFunction
28{
29public:
30static constexpr auto name = "h3GetUnidirectionalEdge";
31
32static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionH3GetUnidirectionalEdge>(); }
33
34std::string getName() const override { return name; }
35
36size_t getNumberOfArguments() const override { return 2; }
37bool useDefaultImplementationForConstants() const override { return true; }
38bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
39
40DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
41{
42const auto * arg = arguments[0].get();
43if (!WhichDataType(arg).isUInt64())
44throw Exception(
45ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
46"Illegal type {} of argument {} of function {}. Must be UInt64",
47arg->getName(), 1, getName());
48
49arg = arguments[1].get();
50if (!WhichDataType(arg).isUInt64())
51throw Exception(
52ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
53"Illegal type {} of argument {} of function {}. Must be UInt64",
54arg->getName(), 2, getName());
55
56return std::make_shared<DataTypeUInt64>();
57}
58
59ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
60{
61auto non_const_arguments = arguments;
62for (auto & argument : non_const_arguments)
63argument.column = argument.column->convertToFullColumnIfConst();
64
65const auto * col_hindex_origin = checkAndGetColumn<ColumnUInt64>(non_const_arguments[0].column.get());
66if (!col_hindex_origin)
67throw Exception(
68ErrorCodes::ILLEGAL_COLUMN,
69"Illegal type {} of argument {} of function {}. Must be UInt64.",
70arguments[0].type->getName(),
711,
72getName());
73
74const auto & data_hindex_origin = col_hindex_origin->getData();
75
76const auto * col_hindex_dest = checkAndGetColumn<ColumnUInt64>(non_const_arguments[1].column.get());
77if (!col_hindex_dest)
78throw Exception(
79ErrorCodes::ILLEGAL_COLUMN,
80"Illegal type {} of argument {} of function {}. Must be UInt64.",
81arguments[1].type->getName(),
822,
83getName());
84
85const auto & data_hindex_dest = col_hindex_dest->getData();
86
87auto dst = ColumnVector<UInt64>::create();
88auto & dst_data = dst->getData();
89dst_data.resize(input_rows_count);
90
91for (size_t row = 0; row < input_rows_count; ++row)
92{
93const UInt64 origin = data_hindex_origin[row];
94const UInt64 dest = data_hindex_dest[row];
95
96if (!isValidCell(origin))
97throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Invalid origin H3 index: {}", origin);
98if (!isValidCell(dest))
99throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Invalid dest H3 index: {}", dest);
100
101UInt64 res = getUnidirectionalEdge(origin, dest);
102dst_data[row] = res;
103}
104
105return dst;
106}
107
108/// suppress asan errors generated by the following:
109/// 'NEW_ADJUSTMENT_III' defined in '../contrib/h3/src/h3lib/lib/algos.c:142:24
110/// 'NEW_DIGIT_III' defined in '../contrib/h3/src/h3lib/lib/algos.c:121:24
111__attribute__((no_sanitize_address)) static inline UInt64 getUnidirectionalEdge(const UInt64 origin, const UInt64 dest)
112{
113const UInt64 res = cellsToDirectedEdge(origin, dest);
114return res;
115}
116};
117
118}
119
120REGISTER_FUNCTION(H3GetUnidirectionalEdge)
121{
122factory.registerFunction<FunctionH3GetUnidirectionalEdge>();
123}
124
125}
126
127#endif
128