diff options
author | vitalyisaev <vitalyisaev@ydb.tech> | 2023-11-14 09:58:56 +0300 |
---|---|---|
committer | vitalyisaev <vitalyisaev@ydb.tech> | 2023-11-14 10:20:20 +0300 |
commit | c2b2dfd9827a400a8495e172a56343462e3ceb82 (patch) | |
tree | cd4e4f597d01bede4c82dffeb2d780d0a9046bd0 /contrib/clickhouse/src/Functions/h3GetUnidirectionalEdge.cpp | |
parent | d4ae8f119e67808cb0cf776ba6e0cf95296f2df7 (diff) | |
download | ydb-c2b2dfd9827a400a8495e172a56343462e3ceb82.tar.gz |
YQ Connector: move tests from yql to ydb (OSS)
Перенос папки с тестами на Коннектор из папки yql в папку ydb (синхронизируется с github).
Diffstat (limited to 'contrib/clickhouse/src/Functions/h3GetUnidirectionalEdge.cpp')
-rw-r--r-- | contrib/clickhouse/src/Functions/h3GetUnidirectionalEdge.cpp | 127 |
1 files changed, 127 insertions, 0 deletions
diff --git a/contrib/clickhouse/src/Functions/h3GetUnidirectionalEdge.cpp b/contrib/clickhouse/src/Functions/h3GetUnidirectionalEdge.cpp new file mode 100644 index 0000000000..052b1f86a8 --- /dev/null +++ b/contrib/clickhouse/src/Functions/h3GetUnidirectionalEdge.cpp @@ -0,0 +1,127 @@ +#include "clickhouse_config.h" + +#if USE_H3 + +#include <Columns/ColumnArray.h> +#include <Columns/ColumnsNumber.h> +#include <DataTypes/DataTypeArray.h> +#include <DataTypes/DataTypesNumber.h> +#include <Functions/FunctionFactory.h> +#include <Functions/IFunction.h> +#include <Common/typeid_cast.h> +#include <IO/WriteHelpers.h> +#include <h3api.h> + + +namespace DB +{ +namespace ErrorCodes +{ + extern const int ILLEGAL_TYPE_OF_ARGUMENT; + extern const int ILLEGAL_COLUMN; +} + +namespace +{ + +class FunctionH3GetUnidirectionalEdge : public IFunction +{ +public: + static constexpr auto name = "h3GetUnidirectionalEdge"; + + static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionH3GetUnidirectionalEdge>(); } + + std::string getName() const override { return name; } + + size_t getNumberOfArguments() const override { return 2; } + bool useDefaultImplementationForConstants() const override { return true; } + bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; } + + DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override + { + const auto * arg = arguments[0].get(); + if (!WhichDataType(arg).isUInt64()) + throw Exception( + ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, + "Illegal type {} of argument {} of function {}. Must be UInt64", + arg->getName(), 1, getName()); + + arg = arguments[1].get(); + if (!WhichDataType(arg).isUInt64()) + throw Exception( + ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, + "Illegal type {} of argument {} of function {}. Must be UInt64", + arg->getName(), 2, getName()); + + return std::make_shared<DataTypeUInt64>(); + } + + ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override + { + auto non_const_arguments = arguments; + for (auto & argument : non_const_arguments) + argument.column = argument.column->convertToFullColumnIfConst(); + + const auto * col_hindex_origin = checkAndGetColumn<ColumnUInt64>(non_const_arguments[0].column.get()); + if (!col_hindex_origin) + throw Exception( + ErrorCodes::ILLEGAL_COLUMN, + "Illegal type {} of argument {} of function {}. Must be UInt64.", + arguments[0].type->getName(), + 1, + getName()); + + const auto & data_hindex_origin = col_hindex_origin->getData(); + + const auto * col_hindex_dest = checkAndGetColumn<ColumnUInt64>(non_const_arguments[1].column.get()); + if (!col_hindex_dest) + throw Exception( + ErrorCodes::ILLEGAL_COLUMN, + "Illegal type {} of argument {} of function {}. Must be UInt64.", + arguments[1].type->getName(), + 2, + getName()); + + const auto & data_hindex_dest = col_hindex_dest->getData(); + + auto dst = ColumnVector<UInt64>::create(); + auto & dst_data = dst->getData(); + dst_data.resize(input_rows_count); + + for (size_t row = 0; row < input_rows_count; ++row) + { + const UInt64 origin = data_hindex_origin[row]; + const UInt64 dest = data_hindex_dest[row]; + + if (!isValidCell(origin)) + throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Invalid origin H3 index: {}", origin); + if (!isValidCell(dest)) + throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Invalid dest H3 index: {}", dest); + + UInt64 res = getUnidirectionalEdge(origin, dest); + dst_data[row] = res; + } + + return dst; + } + + /// suppress asan errors generated by the following: + /// 'NEW_ADJUSTMENT_III' defined in '../contrib/h3/src/h3lib/lib/algos.c:142:24 + /// 'NEW_DIGIT_III' defined in '../contrib/h3/src/h3lib/lib/algos.c:121:24 + __attribute__((no_sanitize_address)) static inline UInt64 getUnidirectionalEdge(const UInt64 origin, const UInt64 dest) + { + const UInt64 res = cellsToDirectedEdge(origin, dest); + return res; + } +}; + +} + +REGISTER_FUNCTION(H3GetUnidirectionalEdge) +{ + factory.registerFunction<FunctionH3GetUnidirectionalEdge>(); +} + +} + +#endif |