ClickHouse

1
#include <Functions/FunctionFactory.h>
2
#include <Functions/GeoHash.h>
3
#include <Functions/FunctionHelpers.h>
4

5
#include <Columns/ColumnString.h>
6
#include <DataTypes/DataTypeString.h>
7
#include <DataTypes/DataTypesNumber.h>
8

9
#include <string>
10

11
#define GEOHASH_MAX_TEXT_LENGTH 16
12

13

14
namespace DB
15
{
16

17
namespace ErrorCodes
18
{
19
    extern const int LOGICAL_ERROR;
20
    extern const int TOO_MANY_ARGUMENTS_FOR_FUNCTION;
21
}
22

23
namespace
24
{
25

26
// geohashEncode(lon float32/64, lat float32/64, length UInt8) => string
27
class FunctionGeohashEncode : public IFunction
28
{
29
public:
30
    static constexpr auto name = "geohashEncode";
31
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionGeohashEncode>(); }
32

33
    String getName() const override
34
    {
35
        return name;
36
    }
37

38
    bool isVariadic() const override { return true; }
39
    size_t getNumberOfArguments() const override { return 0; }
40
    bool useDefaultImplementationForConstants() const override { return true; }
41
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
42

43
    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
44
    {
45
        validateArgumentType(*this, arguments, 0, isFloat, "float");
46
        validateArgumentType(*this, arguments, 1, isFloat, "float");
47
        if (arguments.size() == 3)
48
        {
49
            validateArgumentType(*this, arguments, 2, isInteger, "integer");
50
        }
51
        if (arguments.size() > 3)
52
        {
53
            throw Exception(ErrorCodes::TOO_MANY_ARGUMENTS_FOR_FUNCTION, "Too many arguments for function {} expected at most 3",
54
                            getName());
55
        }
56

57
        return std::make_shared<DataTypeString>();
58
    }
59

60
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t /*input_rows_count*/) const override
61
    {
62
        const IColumn * longitude = arguments[0].column.get();
63
        const IColumn * latitude = arguments[1].column.get();
64

65
        ColumnPtr precision;
66
        if (arguments.size() < 3)
67
            precision = DataTypeUInt8().createColumnConst(longitude->size(), GEOHASH_MAX_TEXT_LENGTH);
68
        else
69
            precision = arguments[2].column;
70

71
        ColumnPtr res_column;
72
        vector(longitude, latitude, precision.get(), res_column);
73
        return res_column;
74
    }
75

76
private:
77
    void vector(const IColumn * lon_column, const IColumn * lat_column, const IColumn * precision_column, ColumnPtr & result) const
78
    {
79
        auto col_str = ColumnString::create();
80
        ColumnString::Chars & out_vec = col_str->getChars();
81
        ColumnString::Offsets & out_offsets = col_str->getOffsets();
82

83
        const size_t size = lat_column->size();
84

85
        out_offsets.resize(size);
86
        out_vec.resize(size * (GEOHASH_MAX_TEXT_LENGTH + 1));
87

88
        char * begin = reinterpret_cast<char *>(out_vec.data());
89
        char * pos = begin;
90

91
        for (size_t i = 0; i < size; ++i)
92
        {
93
            const Float64 longitude_value = lon_column->getFloat64(i);
94
            const Float64 latitude_value = lat_column->getFloat64(i);
95
            const UInt64 precision_value = std::min<UInt64>(precision_column->get64(i), GEOHASH_MAX_TEXT_LENGTH);
96

97
            const size_t encoded_size = geohashEncode(longitude_value, latitude_value, precision_value, pos);
98

99
            pos += encoded_size;
100
            *pos = '\0';
101
            out_offsets[i] = ++pos - begin;
102
        }
103
        out_vec.resize(pos - begin);
104

105
        if (!out_offsets.empty() && out_offsets.back() != out_vec.size())
106
            throw Exception(ErrorCodes::LOGICAL_ERROR, "Column size mismatch (internal logical error)");
107

108
        result = std::move(col_str);
109
    }
110
};
111

112
}
113

114
REGISTER_FUNCTION(GeohashEncode)
115
{
116
    factory.registerFunction<FunctionGeohashEncode>();
117
}
118

119
}
120