H3 library integration

.gitmodules

@@ -76,3 +76,6 @@
 [submodule "contrib/brotli"]
 	path = contrib/brotli
 	url = https://github.com/google/brotli.git
+[submodule "contrib/h3"]
+	path = contrib/h3
+	url = https://github.com/uber/h3

CMakeLists.txt

@@ -247,6 +247,7 @@ include (cmake/find_re2.cmake)
 include (cmake/find_rdkafka.cmake)
 include (cmake/find_capnp.cmake)
 include (cmake/find_llvm.cmake)
+include (cmake/find_h3.cmake)
 include (cmake/find_cpuid.cmake) # Freebsd, bundled
 if (NOT USE_CPUID)
     include (cmake/find_cpuinfo.cmake) # Debian

cmake/find_h3.cmake

@@ -0,0 +1,17 @@
+option (USE_INTERNAL_H3_LIBRARY "Set to FALSE to use system h3 library instead of bundled" ${NOT_UNBUNDLED})
+
+if (USE_INTERNAL_H3_LIBRARY)
+    set (H3_LIBRARY h3)
+    set (H3_INCLUDE_DIR ${ClickHouse_SOURCE_DIR}/contrib/h3/src/h3lib/include)
+else ()
+    find_library (H3_LIBRARY h3)
+    find_path (H3_INCLUDE_DIR NAMES geoCoord.h PATHS ${H3_INCLUDE_PATHS})
+endif ()
+
+if (H3_LIBRARY AND H3_INCLUDE_DIR)
+    set (USE_H3 1)
+else ()
+    set (USE_H3 0)
+endif ()
+
+message (STATUS "Using h3=${USE_H3}: ${H3_INCLUDE_DIR} : ${H3_LIBRARY}")

contrib/CMakeLists.txt

@@ -106,6 +106,11 @@ if (USE_INTERNAL_CPUID_LIBRARY)
     add_subdirectory (libcpuid)
 endif ()
 
+if (USE_INTERNAL_H3_LIBRARY)
+    add_subdirectory(h3)
+endif ()
+
+
 if (USE_INTERNAL_SSL_LIBRARY)
     if (NOT MAKE_STATIC_LIBRARIES)
         set (BUILD_SHARED 1)

dbms/src/Functions/CMakeLists.txt

@@ -19,6 +19,7 @@ target_link_libraries(clickhouse_functions
         ${FARMHASH_LIBRARIES}
         ${METROHASH_LIBRARIES}
         murmurhash
+        m
         ${BASE64_LIBRARY}
         ${OPENSSL_CRYPTO_LIBRARY})
 
@@ -60,3 +61,8 @@ if (USE_XXHASH)
     target_link_libraries(clickhouse_functions PRIVATE ${XXHASH_LIBRARY})
     target_include_directories(clickhouse_functions SYSTEM PRIVATE ${XXHASH_INCLUDE_DIR})
 endif()
+
+if (USE_H3)
+    target_link_libraries(clickhouse_functions PRIVATE ${H3_LIBRARY})
+    target_include_directories(clickhouse_functions SYSTEM PRIVATE ${H3_INCLUDE_DIR})
+endif()

dbms/src/Functions/geoToH3.cpp

@@ -0,0 +1,171 @@
+#include <array>
+#include <math.h>
+#include <Functions/FunctionFactory.h>
+#include <Columns/ColumnConst.h>
+#include <Columns/ColumnsNumber.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <Functions/IFunction.h>
+#include <Common/typeid_cast.h>
+#include <ext/range.h>
+
+
+extern "C" {
+#include <h3Index.h>
+}
+
+namespace DB
+{
+namespace ErrorCodes
+{
+    extern const int ILLEGAL_COLUMN;
+}
+
+/// Implements the function geoToH3 which takes 3 arguments (latitude, longitude and h3 resolution)
+/// and returns h3 index of this point
+class FunctionGeoToH3 : public IFunction
+{
+public:
+    static constexpr auto name = "geoToH3";
+
+    FunctionGeoToH3(const Context & context) : context(context) {}
+
+    static FunctionPtr create(const Context & context) { return std::make_shared<FunctionGeoToH3>(context); }
+
+    std::string getName() const override { return name; }
+
+    size_t getNumberOfArguments() const override { return 3; }
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
+    {
+        auto arg = arguments[0].get();
+        if (!WhichDataType(arg).isFloat64())
+            throw Exception(
+                "Illegal type " + arg->getName() + " of argument " + std::to_string(1) + " of function " + getName() + ". Must be Float64",
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+        arg = arguments[1].get();
+        if (!WhichDataType(arg).isFloat64())
+            throw Exception(
+                "Illegal type " + arg->getName() + " of argument " + std::to_string(2) + " of function " + getName() + ". Must be Float64",
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+        arg = arguments[2].get();
+        if (!WhichDataType(arg).isUInt8())
+            throw Exception(
+                "Illegal type " + arg->getName() + " of argument " + std::to_string(3) + " of function " + getName() + ". Must be UInt8",
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+        return std::make_shared<DataTypeUInt64>();
+    }
+
+    void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override
+    {
+        int const_cnt = 0;
+        const auto size = input_rows_count;
+
+        for (const auto idx : ext::range(0, 2))
+        {
+            const auto column = block.getByPosition(arguments[idx]).column.get();
+            if (typeid_cast<const ColumnConst *>(column))
+            {
+                ++const_cnt;
+            }
+            else if (!typeid_cast<const ColumnVector<Float64> *>(column))
+            {
+                throw Exception(
+                    "Illegal column " + column->getName() + " of argument of function " + getName(), ErrorCodes::ILLEGAL_COLUMN);
+            }
+        }
+
+        double resolution = 0;
+        bool is_const_resulution = false;
+        {
+            const auto column = block.getByPosition(arguments[2]).column.get();
+            if (typeid_cast<const ColumnConst *>(column))
+            {
+                is_const_resulution = true;
+                const auto col_const_res = static_cast<const ColumnConst *>(column);
+                resolution = col_const_res->getValue<UInt8>();
+            }
+            else if (!typeid_cast<const ColumnVector<UInt8> *>(column))
+            {
+                throw Exception(
+                    "Illegal column " + column->getName() + " of argument of function " + getName(), ErrorCodes::ILLEGAL_COLUMN);
+            }
+            else if (const_cnt == 2)
+            {
+                throw Exception(
+                    "Illegal type " + column->getName() + " of arguments 3 of function " + getName()
+                        + ". It must be const if arguments 1 and 2 are consts.",
+                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+            }
+        }
+
+
+        const auto col_lat = block.getByPosition(arguments[0]).column.get();
+        const auto col_lon = block.getByPosition(arguments[1]).column.get();
+        const auto col_res = block.getByPosition(arguments[2]).column.get();
+        if (const_cnt == 0)
+        {
+            const auto col_vec_lat = static_cast<const ColumnVector<Float64> *>(col_lat);
+            const auto col_vec_lon = static_cast<const ColumnVector<Float64> *>(col_lon);
+            const auto col_vec_res = static_cast<const ColumnVector<UInt8> *>(col_res);
+
+            auto dst = ColumnVector<UInt64>::create();
+            auto & dst_data = dst->getData();
+            dst_data.resize(size);
+
+            for (const auto row : ext::range(0, size))
+            {
+                const double lat = col_vec_lat->getData()[row];
+                const double lon = col_vec_lon->getData()[row];
+                if (!is_const_resulution)
+                {
+                    resolution = col_vec_res->getData()[row];
+                }
+
+                GeoCoord coord;
+                setGeoDegs(&coord, lat, lon);
+
+                H3Index hindex = H3_EXPORT(geoToH3)(&coord, resolution);
+
+                dst_data[row] = hindex;
+            }
+
+            block.getByPosition(result).column = std::move(dst);
+        }
+        else if (const_cnt == 2)
+        {
+            const auto col_const_lat = static_cast<const ColumnConst *>(col_lat);
+            const auto col_const_lon = static_cast<const ColumnConst *>(col_lon);
+
+            const double lat = col_const_lat->getValue<Float64>();
+            const double lon = col_const_lon->getValue<Float64>();
+
+            GeoCoord coord;
+            setGeoDegs(&coord, lat, lon);
+            H3Index hindex = H3_EXPORT(geoToH3)(&coord, resolution);
+
+            block.getByPosition(result).column = DataTypeUInt64().createColumnConst(size, hindex);
+        }
+        else
+        {
+            throw Exception(
+                "Illegal types " + col_lat->getName() + ", " + col_lon->getName() + " of arguments 1, 2 of function " + getName()
+                    + ". All must be either const or vector",
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+        }
+    }
+
+private:
+    const Context & context;
+};
+
+
+void registerFunctionGeoToH3(FunctionFactory & factory)
+{
+    factory.registerFunction<FunctionGeoToH3>(FunctionFactory::CaseInsensitive);
+}
+
+}

dbms/src/Functions/registerFunctions.cpp

@@ -42,6 +42,7 @@ void registerFunctionsGeo(FunctionFactory &);
 void registerFunctionsNull(FunctionFactory &);
 void registerFunctionsFindCluster(FunctionFactory &);
 void registerFunctionTransform(FunctionFactory &);
+void registerFunctionGeoToH3(FunctionFactory &);
 
 #if USE_ICU
 void registerFunctionConvertCharset(FunctionFactory &);
@@ -83,6 +84,7 @@ void registerFunctions()
     registerFunctionsNull(factory);
     registerFunctionsFindCluster(factory);
     registerFunctionTransform(factory);
+    registerFunctionGeoToH3(factory);
 
 #if USE_ICU
     registerFunctionConvertCharset(factory);

dbms/tests/queries/0_stateless/00926_geo_to_h3.reference

@@ -0,0 +1,20 @@
+644325529094369568
+639821928864584823
+644325528491955313
+644325528491955313
+644325528627451570
+644325529094369568
+644325528491955313
+644325528491955313
+644325528491955313
+644325528627451570
+644325529094369568
+55.720762	37.598135	644325528491955313
+55.720762	37.598135	644325528491955313
+55.72076201	37.598135	644325528491955313
+55.763241	37.660183	644325528627451570
+55.77922738	37.63098076	644325529094369568
+639821928864584823	1
+644325528491955313	2
+644325528627451570	1
+644325529094369568	1

dbms/tests/queries/0_stateless/00926_geo_to_h3.sql

@@ -0,0 +1,19 @@
+USE test;
+
+DROP TABLE IF EXISTS table1;
+
+CREATE TABLE table1 (lat Float64, lon Float64, resolution UInt8) ENGINE = Memory;
+
+INSERT INTO table1 VALUES(55.77922738, 37.63098076, 15);
+INSERT INTO table1 VALUES(55.76324100, 37.66018300, 15);
+INSERT INTO table1 VALUES(55.72076200, 37.59813500, 15);
+INSERT INTO table1 VALUES(55.72076201, 37.59813500, 15);
+INSERT INTO table1 VALUES(55.72076200, 37.59813500, 14);
+
+select geoToH3(55.77922738, 37.63098076, 15);
+select geoToH3(lat, lon, resolution) from table1 order by lat, lon, resolution;
+select geoToH3(lat, lon, 15) from table1 order by lat, lon, geoToH3(lat, lon, 15);
+select lat, lon, geoToH3(lat, lon, 15) from table1 order by lat, lon, geoToH3(lat, lon, 15);
+select geoToH3(lat, lon, resolution), count(*) from table1 group by geoToH3(lat, lon, resolution) order by geoToH3(lat, lon, resolution);
+
+DROP TABLE table1

docs/ru/query_language/functions/geo.md

@@ -99,4 +99,37 @@ SELECT pointInPolygon((3., 3.), [(6, 0), (8, 4), (5, 8), (0, 2)]) AS res
 └─────┘
 ```
 
+## geoToH3
+
+Получает H3 индекс точки (lat, lon) с заданным разрешением
+
+```
+pointInPolygon(lat, lon, resolution)
+```
+
+**Входные значения**
+
+- `lat` - географическая широта. Тип данных — [Float64](../../data_types/float.md).
+- `lon` - географическая долгота. Тип данных — [Float64](../../data_types/float.md).
+- `resolution` - требуемое разрешение индекса. Тип данных — [UInt8](../../data_types/int_uint.md). Диапазон возможных значение — `[0, 15]`.
+
+Параметры `lat` и `lon` должны быть одновременно или константными, или нет. Если параметры `lat` и `lon` не являются константными, то параметр `resolution` не может быть константным.
+
+**Возвращаемые значения**
+
+Возвращает значение с типом [UInt64] (../../data_types/int_uint.md).
+`0` в случае ошибки.
+Иначе возвращается индексный номер шестиугольника.
+
+**Пример**
+
+``` sql
+SELECT geoToH3(55.71290588, 37.79506683, 15) as h3Index
+```
+```
+┌────────────h3Index─┐
+│ 644325524701193974 │
+└────────────────────┘
+```
+
 [Оригинальная статья](https://clickhouse.yandex/docs/ru/query_language/functions/geo/) <!--hide-->