1 | #include <Functions/FunctionFactory.h> |
2 | #include <Functions/GeoUtils.h> |
3 | #include <Functions/FunctionHelpers.h> |
4 | |
5 | #include <Columns/ColumnString.h> |
6 | #include <DataTypes/DataTypeString.h> |
7 | |
8 | #include <string> |
9 | |
10 | #define GEOHASH_MAX_TEXT_LENGTH 16 |
11 | |
12 | |
13 | namespace DB |
14 | { |
15 | |
16 | namespace ErrorCodes |
17 | { |
18 | extern const int LOGICAL_ERROR; |
19 | extern const int ILLEGAL_COLUMN; |
20 | extern const int TOO_MANY_ARGUMENTS_FOR_FUNCTION; |
21 | } |
22 | |
23 | // geohashEncode(lon float32/64, lat float32/64, length UInt8) => string |
24 | class FunctionGeohashEncode : public IFunction |
25 | { |
26 | public: |
27 | static constexpr auto name = "geohashEncode" ; |
28 | static FunctionPtr create(const Context &) { return std::make_shared<FunctionGeohashEncode>(); } |
29 | |
30 | String getName() const override |
31 | { |
32 | return name; |
33 | } |
34 | |
35 | bool isVariadic() const override { return true; } |
36 | size_t getNumberOfArguments() const override { return 0; } |
37 | ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {2}; } |
38 | bool useDefaultImplementationForConstants() const override { return true; } |
39 | |
40 | DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override |
41 | { |
42 | validateArgumentType(*this, arguments, 0, isFloat, "float" ); |
43 | validateArgumentType(*this, arguments, 1, isFloat, "float" ); |
44 | if (arguments.size() == 3) |
45 | { |
46 | validateArgumentType(*this, arguments, 2, isInteger, "integer" ); |
47 | } |
48 | if (arguments.size() > 3) |
49 | { |
50 | throw Exception("Too many arguments for function " + getName() + |
51 | " expected at most 3" , |
52 | ErrorCodes::TOO_MANY_ARGUMENTS_FOR_FUNCTION); |
53 | } |
54 | |
55 | return std::make_shared<DataTypeString>(); |
56 | } |
57 | |
58 | template <typename LonType, typename LatType> |
59 | bool tryExecute(const IColumn * lon_column, const IColumn * lat_column, UInt64 precision_value, ColumnPtr & result) |
60 | { |
61 | const ColumnVector<LonType> * longitude = checkAndGetColumn<ColumnVector<LonType>>(lon_column); |
62 | const ColumnVector<LatType> * latitude = checkAndGetColumn<ColumnVector<LatType>>(lat_column); |
63 | if (!latitude || !longitude) |
64 | return false; |
65 | |
66 | auto col_str = ColumnString::create(); |
67 | ColumnString::Chars & out_vec = col_str->getChars(); |
68 | ColumnString::Offsets & out_offsets = col_str->getOffsets(); |
69 | |
70 | const size_t size = lat_column->size(); |
71 | |
72 | out_offsets.resize(size); |
73 | out_vec.resize(size * (GEOHASH_MAX_TEXT_LENGTH + 1)); |
74 | |
75 | char * begin = reinterpret_cast<char *>(out_vec.data()); |
76 | char * pos = begin; |
77 | |
78 | for (size_t i = 0; i < size; ++i) |
79 | { |
80 | const Float64 longitude_value = longitude->getElement(i); |
81 | const Float64 latitude_value = latitude->getElement(i); |
82 | |
83 | const size_t encoded_size = GeoUtils::geohashEncode(longitude_value, latitude_value, precision_value, pos); |
84 | |
85 | pos += encoded_size; |
86 | *pos = '\0'; |
87 | out_offsets[i] = ++pos - begin; |
88 | } |
89 | out_vec.resize(pos - begin); |
90 | |
91 | if (!out_offsets.empty() && out_offsets.back() != out_vec.size()) |
92 | throw Exception("Column size mismatch (internal logical error)" , ErrorCodes::LOGICAL_ERROR); |
93 | |
94 | result = std::move(col_str); |
95 | |
96 | return true; |
97 | |
98 | } |
99 | |
100 | void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t /*input_rows_count*/) override |
101 | { |
102 | const IColumn * longitude = block.getByPosition(arguments[0]).column.get(); |
103 | const IColumn * latitude = block.getByPosition(arguments[1]).column.get(); |
104 | |
105 | const UInt64 precision_value = std::min<UInt64>(GEOHASH_MAX_TEXT_LENGTH, |
106 | arguments.size() == 3 ? block.getByPosition(arguments[2]).column->get64(0) : GEOHASH_MAX_TEXT_LENGTH); |
107 | |
108 | ColumnPtr & res_column = block.getByPosition(result).column; |
109 | |
110 | if (tryExecute<Float32, Float32>(longitude, latitude, precision_value, res_column) || |
111 | tryExecute<Float64, Float32>(longitude, latitude, precision_value, res_column) || |
112 | tryExecute<Float32, Float64>(longitude, latitude, precision_value, res_column) || |
113 | tryExecute<Float64, Float64>(longitude, latitude, precision_value, res_column)) |
114 | return; |
115 | |
116 | std::string arguments_description; |
117 | for (size_t i = 0; i < arguments.size(); ++i) |
118 | { |
119 | if (i != 0) |
120 | arguments_description += ", " ; |
121 | arguments_description += block.getByPosition(arguments[i]).column->getName(); |
122 | } |
123 | |
124 | throw Exception("Unsupported argument types: " + arguments_description + |
125 | + " for function " + getName(), |
126 | ErrorCodes::ILLEGAL_COLUMN); |
127 | } |
128 | }; |
129 | |
130 | |
131 | void registerFunctionGeohashEncode(FunctionFactory & factory) |
132 | { |
133 | factory.registerFunction<FunctionGeohashEncode>(); |
134 | } |
135 | |
136 | } |
137 | |