ClickHouse
124 строки · 3.2 Кб
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 <base/range.h>
12
13#include "s2_fwd.h"
14
15class S2CellId;
16
17namespace DB
18{
19
20namespace ErrorCodes
21{
22extern const int ILLEGAL_TYPE_OF_ARGUMENT;
23extern const int BAD_ARGUMENTS;
24extern const int ILLEGAL_COLUMN;
25}
26
27namespace
28{
29
30/**
31* Returns a point (longitude, latitude) in degrees
32*/
33class FunctionS2ToGeo : public IFunction
34{
35public:
36static constexpr auto name = "s2ToGeo";
37
38static FunctionPtr create(ContextPtr)
39{
40return std::make_shared<FunctionS2ToGeo>();
41}
42
43std::string getName() const override
44{
45return name;
46}
47
48size_t getNumberOfArguments() const override { return 1; }
49
50bool useDefaultImplementationForConstants() const override { return true; }
51
52bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
53
54DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
55{
56const auto * arg = arguments[0].get();
57
58if (!WhichDataType(arg).isUInt64())
59throw Exception(
60ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
61"Illegal type {} of argument {} of function {}. Must be UInt64",
62arg->getName(), 1, getName());
63
64DataTypePtr element = std::make_shared<DataTypeFloat64>();
65
66return std::make_shared<DataTypeTuple>(DataTypes{element, element});
67}
68
69ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
70{
71auto non_const_arguments = arguments;
72for (auto & argument : non_const_arguments)
73argument.column = argument.column->convertToFullColumnIfConst();
74
75const auto * col_id = checkAndGetColumn<ColumnUInt64>(non_const_arguments[0].column.get());
76if (!col_id)
77throw Exception(
78ErrorCodes::ILLEGAL_COLUMN,
79"Illegal type {} of argument {} of function {}. Must be UInt64",
80arguments[0].type->getName(),
811,
82getName());
83
84const auto & data_id = col_id->getData();
85
86auto col_longitude = ColumnFloat64::create();
87auto col_latitude = ColumnFloat64::create();
88
89auto & longitude = col_longitude->getData();
90longitude.reserve(input_rows_count);
91
92auto & latitude = col_latitude->getData();
93latitude.reserve(input_rows_count);
94
95for (size_t row = 0; row < input_rows_count; ++row)
96{
97const auto id = S2CellId(data_id[row]);
98
99if (!id.is_valid())
100throw Exception(ErrorCodes::BAD_ARGUMENTS, "CellId is invalid.");
101
102S2Point point = id.ToPoint();
103S2LatLng ll(point);
104
105longitude.emplace_back(ll.lng().degrees());
106latitude.emplace_back(ll.lat().degrees());
107}
108
109return ColumnTuple::create(Columns{std::move(col_longitude), std::move(col_latitude)});
110}
111
112};
113
114}
115
116REGISTER_FUNCTION(S2ToGeo)
117{
118factory.registerFunction<FunctionS2ToGeo>();
119}
120
121
122}
123
124#endif
125