diff --git a/dbms/src/Functions/h3kRing.cpp b/dbms/src/Functions/h3kRing.cpp new file mode 100644 index 00000000000..5fcbad883cf --- /dev/null +++ b/dbms/src/Functions/h3kRing.cpp @@ -0,0 +1,103 @@ +#include "config_functions.h" +#if USE_H3 +# include +# include +# include +# include +# include +# include +# include +# include +# include +# include + +extern "C" { +# ifdef __clang__ +# pragma clang diagnostic push +# pragma clang diagnostic ignored "-Wdocumentation" +# endif + +# include + +# ifdef __clang__ +# pragma clang diagnostic pop +# endif +} + +namespace DB +{ +class FunctionH3KRing : public IFunction +{ +public: + static constexpr auto name = "h3kRing"; + + static FunctionPtr create(const Context &) { return std::make_shared(); } + + std::string getName() const override { return name; } + + size_t getNumberOfArguments() const override { return 2; } + bool useDefaultImplementationForConstants() const override { return true; } + + DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override + { + auto arg = arguments[0].get(); + if (!WhichDataType(arg).isUInt64()) + throw Exception( + "Illegal type " + arg->getName() + " of argument " + std::to_string(1) + " of function " + getName() + ". Must be UInt64", + ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); + + arg = arguments[1].get(); + if (!isInteger(arg)) + throw Exception( + "Illegal type " + arg->getName() + " of argument " + std::to_string(2) + " of function " + getName() + ". Must be integer", + ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); + + return std::make_shared(std::make_shared()); + } + + void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result, size_t input_rows_count) override + { + const auto col_hindex = block.getByPosition(arguments[0]).column.get(); + const auto col_k = block.getByPosition(arguments[1]).column.get(); + + auto dst = ColumnArray::create(ColumnUInt64::create()); + auto & dst_data = dst->getData(); + auto & dst_offsets = dst->getOffsets(); + dst_offsets.resize(input_rows_count); + auto current_offset = 0; + + std::vector hindex_vec; + + for (const auto row : ext::range(0, input_rows_count)) + { + const H3Index origin_hindex = col_hindex->getUInt(row); + const int k = col_k->getInt(row); + + const auto vec_size = H3_EXPORT(maxKringSize)(k); + hindex_vec.resize(vec_size); + H3_EXPORT(kRing)(origin_hindex, k, hindex_vec.data()); + + dst_data.reserve(dst_data.size() + vec_size); + for (auto hindex : hindex_vec) + { + if (hindex != 0) + { + current_offset++; + dst_data.insert(hindex); + } + } + dst_offsets[row] = current_offset; + } + + block.getByPosition(result).column = std::move(dst); + } +}; + + +void registerFunctionH3KRing(FunctionFactory & factory) +{ + factory.registerFunction(); +} + +} +#endif diff --git a/dbms/src/Functions/registerFunctionsGeo.cpp b/dbms/src/Functions/registerFunctionsGeo.cpp index a6af3471233..8390d910818 100644 --- a/dbms/src/Functions/registerFunctionsGeo.cpp +++ b/dbms/src/Functions/registerFunctionsGeo.cpp @@ -16,6 +16,7 @@ void registerFunctionGeohashesInBox(FunctionFactory & factory); void registerFunctionGeoToH3(FunctionFactory &); void registerFunctionH3GetResolution(FunctionFactory &); void registerFunctionH3IsValid(FunctionFactory &); +void registerFunctionH3KRing(FunctionFactory &); #endif void registerFunctionsGeo(FunctionFactory & factory) @@ -31,6 +32,7 @@ void registerFunctionsGeo(FunctionFactory & factory) registerFunctionGeoToH3(factory); registerFunctionH3GetResolution(factory); registerFunctionH3IsValid(factory); + registerFunctionH3KRing(factory); #endif } diff --git a/dbms/tests/queries/0_stateless/01042_h3_k_ring.reference b/dbms/tests/queries/0_stateless/01042_h3_k_ring.reference new file mode 100644 index 00000000000..860ddac7547 --- /dev/null +++ b/dbms/tests/queries/0_stateless/01042_h3_k_ring.reference @@ -0,0 +1,3 @@ +[581250224954015743,581259021047037951,581267817140060159,581276613233082367,581500913605148671,581518505791193087,581764796395814911] +[581276613233082367] +[581276613233082367] diff --git a/dbms/tests/queries/0_stateless/01042_h3_k_ring.sql b/dbms/tests/queries/0_stateless/01042_h3_k_ring.sql new file mode 100644 index 00000000000..30bd2f58801 --- /dev/null +++ b/dbms/tests/queries/0_stateless/01042_h3_k_ring.sql @@ -0,0 +1,3 @@ +SELECT arraySort(h3kRing(581276613233082367, 1)); +SELECT h3kRing(581276613233082367, 0); +SELECT h3kRing(581276613233082367, -1);