mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-11-21 23:21:59 +00:00
add some cap and rect functions
This commit is contained in:
parent
4d75450e54
commit
fe075f407e
121
src/Functions/S2CapContains.cpp
Normal file
121
src/Functions/S2CapContains.cpp
Normal file
@ -0,0 +1,121 @@
|
||||
#include "config_functions.h"
|
||||
|
||||
#include <Columns/ColumnsNumber.h>
|
||||
#include <Columns/ColumnTuple.h>
|
||||
#include <DataTypes/DataTypesNumber.h>
|
||||
#include <DataTypes/DataTypeTuple.h>
|
||||
#include <Functions/FunctionFactory.h>
|
||||
#include <Functions/IFunctionImpl.h>
|
||||
#include <Common/typeid_cast.h>
|
||||
#include <ext/range.h>
|
||||
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2cell_id.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2cap.h"
|
||||
#include "../contrib/s2geometry/src/s2/s1angle.h"
|
||||
|
||||
namespace DB
|
||||
{
|
||||
|
||||
namespace ErrorCodes
|
||||
{
|
||||
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
|
||||
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
|
||||
}
|
||||
|
||||
namespace
|
||||
{
|
||||
|
||||
/// TODO: Comment this
|
||||
class FunctionS2CapContains : public IFunction
|
||||
{
|
||||
public:
|
||||
static constexpr auto name = "S2CapContains";
|
||||
|
||||
static FunctionPtr create(ContextPtr)
|
||||
{
|
||||
return std::make_shared<FunctionS2CapContains>();
|
||||
}
|
||||
|
||||
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
|
||||
{
|
||||
size_t number_of_arguments = arguments.size();
|
||||
|
||||
if (number_of_arguments != 3) {
|
||||
throw Exception("Number of arguments for function " + getName() + " doesn't match: passed "
|
||||
+ toString(number_of_arguments) + ", should be 3",
|
||||
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
|
||||
}
|
||||
|
||||
const 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 (!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).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(3) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
return std::make_shared<DataTypeUInt8>();
|
||||
}
|
||||
|
||||
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
|
||||
{
|
||||
const auto * col_center = arguments[0].column.get();
|
||||
const auto * col_degrees = arguments[1].column.get();
|
||||
const auto * col_point = arguments[2].column.get();
|
||||
|
||||
auto dst = ColumnVector<UInt8>::create();
|
||||
auto & dst_data = dst->getData();
|
||||
dst_data.resize(input_rows_count);
|
||||
|
||||
for (const auto row : ext::range(0, input_rows_count))
|
||||
{
|
||||
const UInt64 center = col_center->getUInt(row);
|
||||
const Float64 degrees = col_degrees->getFloat64(row);
|
||||
const UInt64 point = col_point->getInt(row);
|
||||
|
||||
S1Angle angle = S1Angle::Degrees(degrees);
|
||||
S2Cap cap(S2CellId(center).ToPoint(), angle);
|
||||
|
||||
dst_data[row] = cap.Contains(S2CellId(point).ToPoint());
|
||||
}
|
||||
|
||||
return dst;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
void registerFunctionS2CapContains(FunctionFactory & factory)
|
||||
{
|
||||
factory.registerFunction<FunctionS2CapContains>();
|
||||
}
|
||||
|
||||
|
||||
}
|
144
src/Functions/S2CapUnion.cpp
Normal file
144
src/Functions/S2CapUnion.cpp
Normal file
@ -0,0 +1,144 @@
|
||||
#include "config_functions.h"
|
||||
|
||||
#include <Columns/ColumnsNumber.h>
|
||||
#include <Columns/ColumnTuple.h>
|
||||
#include <DataTypes/DataTypesNumber.h>
|
||||
#include <DataTypes/DataTypeTuple.h>
|
||||
#include <Functions/FunctionFactory.h>
|
||||
#include <Functions/IFunctionImpl.h>
|
||||
#include <Common/typeid_cast.h>
|
||||
#include <ext/range.h>
|
||||
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2cell_id.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2cap.h"
|
||||
#include "../contrib/s2geometry/src/s2/s1angle.h"
|
||||
|
||||
class S2CellId;
|
||||
|
||||
namespace DB
|
||||
{
|
||||
|
||||
namespace ErrorCodes
|
||||
{
|
||||
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
|
||||
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
|
||||
}
|
||||
|
||||
namespace
|
||||
{
|
||||
|
||||
/// TODO: Comment this
|
||||
class FunctionS2CapUnion : public IFunction
|
||||
{
|
||||
public:
|
||||
static constexpr auto name = "S2CapUnion";
|
||||
|
||||
static FunctionPtr create(ContextPtr)
|
||||
{
|
||||
return std::make_shared<FunctionS2CapUnion>();
|
||||
}
|
||||
|
||||
std::string getName() const override
|
||||
{
|
||||
return name;
|
||||
}
|
||||
|
||||
size_t getNumberOfArguments() const override { return 4; }
|
||||
|
||||
bool useDefaultImplementationForConstants() const override { return true; }
|
||||
|
||||
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
|
||||
{
|
||||
size_t number_of_arguments = arguments.size();
|
||||
|
||||
if (number_of_arguments != 4) {
|
||||
throw Exception("Number of arguments for function " + getName() + " doesn't match: passed "
|
||||
+ toString(number_of_arguments) + ", should be 4",
|
||||
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
|
||||
}
|
||||
|
||||
const 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 (!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).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(3) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
arg = arguments[3].get();
|
||||
|
||||
if (!WhichDataType(arg).isFloat64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(4) + " of function " + getName() + ". Must be Float64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
DataTypePtr center = std::make_shared<DataTypeUInt64>();
|
||||
DataTypePtr radius = std::make_shared<DataTypeFloat64>();
|
||||
|
||||
return std::make_shared<DataTypeTuple>(DataTypes{center, radius});
|
||||
}
|
||||
|
||||
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
|
||||
{
|
||||
const auto * col_center1 = arguments[0].column.get();
|
||||
const auto * col_radius1 = arguments[1].column.get();
|
||||
const auto * col_center2 = arguments[2].column.get();
|
||||
const auto * col_radius2 = arguments[3].column.get();
|
||||
|
||||
auto col_res_first = ColumnUInt64::create();
|
||||
auto col_res_second = ColumnFloat64::create();
|
||||
|
||||
auto & vec_res_first = col_res_first->getData();
|
||||
vec_res_first.resize(input_rows_count);
|
||||
|
||||
auto & vec_res_second = col_res_second->getData();
|
||||
vec_res_second.resize(input_rows_count);
|
||||
|
||||
for (const auto row : ext::range(0, input_rows_count))
|
||||
{
|
||||
const UInt64 center1 = col_center1->getUInt(row);
|
||||
const Float64 radius1 = col_radius1->getFloat64(row);
|
||||
const UInt64 center2 = col_center2->getUInt(row);
|
||||
const Float64 radius2 = col_radius2->getFloat64(row);
|
||||
|
||||
S2Cap cap1(S2CellId(center1).ToPoint(), S1Angle::Degrees(radius1));
|
||||
S2Cap cap2(S2CellId(center2).ToPoint(), S1Angle::Degrees(radius2));
|
||||
|
||||
S2Cap cap_union = cap1.Union(cap2);
|
||||
|
||||
vec_res_first[row] = S2CellId(cap_union.center()).id();
|
||||
vec_res_second[row] = cap_union.GetRadius().degrees();
|
||||
}
|
||||
|
||||
return ColumnTuple::create(Columns{std::move(col_res_first), std::move(col_res_second)});
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
void registerFunctionS2CapUnion(FunctionFactory & factory)
|
||||
{
|
||||
factory.registerFunction<FunctionS2CapUnion>();
|
||||
}
|
||||
|
||||
|
||||
}
|
@ -60,7 +60,7 @@ public:
|
||||
|
||||
if (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(1) + " of function " + getName() + ". Must be Float64",
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(1) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
@ -68,7 +68,7 @@ public:
|
||||
|
||||
if (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(1) + " of function " + getName() + ". Must be Float64",
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(2) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
|
136
src/Functions/S2RectAdd.cpp
Normal file
136
src/Functions/S2RectAdd.cpp
Normal file
@ -0,0 +1,136 @@
|
||||
#include "config_functions.h"
|
||||
|
||||
#include <Columns/ColumnsNumber.h>
|
||||
#include <Columns/ColumnTuple.h>
|
||||
#include <DataTypes/DataTypesNumber.h>
|
||||
#include <DataTypes/DataTypeTuple.h>
|
||||
#include <Functions/FunctionFactory.h>
|
||||
#include <Functions/IFunctionImpl.h>
|
||||
#include <Common/typeid_cast.h>
|
||||
#include <ext/range.h>
|
||||
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2cell_id.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2point.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng_rect.h"
|
||||
|
||||
class S2CellId;
|
||||
|
||||
namespace DB
|
||||
{
|
||||
|
||||
namespace ErrorCodes
|
||||
{
|
||||
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
|
||||
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
|
||||
}
|
||||
|
||||
namespace
|
||||
{
|
||||
|
||||
/// TODO: Comment this
|
||||
class FunctionS2RectAdd : public IFunction
|
||||
{
|
||||
public:
|
||||
static constexpr auto name = "S2RectAdd";
|
||||
|
||||
static FunctionPtr create(ContextPtr)
|
||||
{
|
||||
return std::make_shared<FunctionS2RectAdd>();
|
||||
}
|
||||
|
||||
std::string getName() const override
|
||||
{
|
||||
return name;
|
||||
}
|
||||
|
||||
size_t getNumberOfArguments() const override { return 4; }
|
||||
|
||||
bool useDefaultImplementationForConstants() const override { return true; }
|
||||
|
||||
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
|
||||
{
|
||||
size_t number_of_arguments = arguments.size();
|
||||
|
||||
if (number_of_arguments != 3) {
|
||||
throw Exception("Number of arguments for function " + getName() + " doesn't match: passed "
|
||||
+ toString(number_of_arguments) + ", should be 3",
|
||||
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
|
||||
}
|
||||
|
||||
const 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 (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(2) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
arg = arguments[2].get();
|
||||
|
||||
if (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(3) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
DataTypePtr element = std::make_shared<DataTypeUInt64>();
|
||||
|
||||
return std::make_shared<DataTypeTuple>(DataTypes{element, element});
|
||||
}
|
||||
|
||||
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
|
||||
{
|
||||
const auto * col_lo = arguments[0].column.get();
|
||||
const auto * col_hi = arguments[1].column.get();
|
||||
const auto * col_point = arguments[2].column.get();
|
||||
|
||||
auto col_res_first = ColumnUInt64::create();
|
||||
auto col_res_second = ColumnUInt64::create();
|
||||
|
||||
auto & vec_res_first = col_res_first->getData();
|
||||
vec_res_first.resize(input_rows_count);
|
||||
|
||||
auto & vec_res_second = col_res_second->getData();
|
||||
vec_res_second.resize(input_rows_count);
|
||||
|
||||
for (const auto row : ext::range(0, input_rows_count))
|
||||
{
|
||||
const UInt64 lo = col_lo->getUInt(row);
|
||||
const UInt64 hi = col_hi->getUInt(row);
|
||||
const UInt64 point = col_point->getUInt(row);
|
||||
|
||||
S2CellId id_lo(lo);
|
||||
S2CellId id_hi(hi);
|
||||
S2CellId id_point(point);
|
||||
|
||||
S2LatLngRect rect(id_lo.ToLatLng(), id_hi.ToLatLng());
|
||||
|
||||
rect.AddPoint(id_point.ToPoint());
|
||||
|
||||
vec_res_first[row] = S2CellId(rect.lo()).id();
|
||||
vec_res_second[row] = S2CellId(rect.hi()).id();
|
||||
}
|
||||
|
||||
return ColumnTuple::create(Columns{std::move(col_res_first), std::move(col_res_second)});
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
void registerFunctionS2RectAdd(FunctionFactory & factory)
|
||||
{
|
||||
factory.registerFunction<FunctionS2RectAdd>();
|
||||
}
|
||||
|
||||
|
||||
}
|
126
src/Functions/S2RectContains.cpp
Normal file
126
src/Functions/S2RectContains.cpp
Normal file
@ -0,0 +1,126 @@
|
||||
#include "config_functions.h"
|
||||
|
||||
#include <Columns/ColumnsNumber.h>
|
||||
#include <Columns/ColumnTuple.h>
|
||||
#include <DataTypes/DataTypesNumber.h>
|
||||
#include <DataTypes/DataTypeTuple.h>
|
||||
#include <Functions/FunctionFactory.h>
|
||||
#include <Functions/IFunctionImpl.h>
|
||||
#include <Common/typeid_cast.h>
|
||||
#include <ext/range.h>
|
||||
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2cell_id.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2point.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng_rect.h"
|
||||
|
||||
class S2CellId;
|
||||
|
||||
namespace DB
|
||||
{
|
||||
|
||||
namespace ErrorCodes
|
||||
{
|
||||
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
|
||||
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
|
||||
}
|
||||
|
||||
namespace
|
||||
{
|
||||
|
||||
/// TODO: Comment this
|
||||
class FunctionS2RectContains : public IFunction
|
||||
{
|
||||
public:
|
||||
static constexpr auto name = "S2RectContains";
|
||||
|
||||
static FunctionPtr create(ContextPtr)
|
||||
{
|
||||
return std::make_shared<FunctionS2RectContains>();
|
||||
}
|
||||
|
||||
std::string getName() const override
|
||||
{
|
||||
return name;
|
||||
}
|
||||
|
||||
size_t getNumberOfArguments() const override { return 4; }
|
||||
|
||||
bool useDefaultImplementationForConstants() const override { return true; }
|
||||
|
||||
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
|
||||
{
|
||||
size_t number_of_arguments = arguments.size();
|
||||
|
||||
if (number_of_arguments != 3) {
|
||||
throw Exception("Number of arguments for function " + getName() + " doesn't match: passed "
|
||||
+ toString(number_of_arguments) + ", should be 3",
|
||||
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
|
||||
}
|
||||
|
||||
const 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 (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(2) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
arg = arguments[2].get();
|
||||
|
||||
if (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(3) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
return std::make_shared<DataTypeUInt8>();
|
||||
}
|
||||
|
||||
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
|
||||
{
|
||||
const auto * col_lo = arguments[0].column.get();
|
||||
const auto * col_hi = arguments[1].column.get();
|
||||
const auto * col_point = arguments[2].column.get();
|
||||
|
||||
auto dst = ColumnVector<UInt8>::create();
|
||||
auto & dst_data = dst->getData();
|
||||
dst_data.resize(input_rows_count);
|
||||
|
||||
for (const auto row : ext::range(0, input_rows_count))
|
||||
{
|
||||
const UInt64 lo = col_lo->getUInt(row);
|
||||
const UInt64 hi = col_hi->getUInt(row);
|
||||
const UInt64 point = col_point->getUInt(row);
|
||||
|
||||
S2CellId id_lo(lo);
|
||||
S2CellId id_hi(hi);
|
||||
S2CellId id_point(point);
|
||||
|
||||
S2LatLngRect rect(id_lo.ToLatLng(), id_hi.ToLatLng());
|
||||
|
||||
dst_data[row] = rect.Contains(id_point.ToLatLng());
|
||||
}
|
||||
|
||||
return dst;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
void registerFunctionS2RectContains(FunctionFactory & factory)
|
||||
{
|
||||
factory.registerFunction<FunctionS2RectContains>();
|
||||
}
|
||||
|
||||
|
||||
}
|
148
src/Functions/S2RectIntersection.cpp
Normal file
148
src/Functions/S2RectIntersection.cpp
Normal file
@ -0,0 +1,148 @@
|
||||
#include "config_functions.h"
|
||||
|
||||
#include <Columns/ColumnsNumber.h>
|
||||
#include <Columns/ColumnTuple.h>
|
||||
#include <DataTypes/DataTypesNumber.h>
|
||||
#include <DataTypes/DataTypeTuple.h>
|
||||
#include <Functions/FunctionFactory.h>
|
||||
#include <Functions/IFunctionImpl.h>
|
||||
#include <Common/typeid_cast.h>
|
||||
#include <ext/range.h>
|
||||
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2cell_id.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2point.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng_rect.h"
|
||||
|
||||
class S2CellId;
|
||||
|
||||
namespace DB
|
||||
{
|
||||
|
||||
namespace ErrorCodes
|
||||
{
|
||||
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
|
||||
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
|
||||
}
|
||||
|
||||
namespace
|
||||
{
|
||||
|
||||
/// TODO: Comment this
|
||||
class FunctionS2RectIntersecion : public IFunction
|
||||
{
|
||||
public:
|
||||
static constexpr auto name = "S2RectIntersecion";
|
||||
|
||||
static FunctionPtr create(ContextPtr)
|
||||
{
|
||||
return std::make_shared<FunctionS2RectIntersecion>();
|
||||
}
|
||||
|
||||
std::string getName() const override
|
||||
{
|
||||
return name;
|
||||
}
|
||||
|
||||
size_t getNumberOfArguments() const override { return 4; }
|
||||
|
||||
bool useDefaultImplementationForConstants() const override { return true; }
|
||||
|
||||
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
|
||||
{
|
||||
size_t number_of_arguments = arguments.size();
|
||||
|
||||
if (number_of_arguments != 4) {
|
||||
throw Exception("Number of arguments for function " + getName() + " doesn't match: passed "
|
||||
+ toString(number_of_arguments) + ", should be 4",
|
||||
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
|
||||
}
|
||||
|
||||
const 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 (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(2) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
arg = arguments[2].get();
|
||||
|
||||
if (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(3) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
arg = arguments[3].get();
|
||||
|
||||
if (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(4) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
DataTypePtr element = std::make_shared<DataTypeUInt64>();
|
||||
|
||||
return std::make_shared<DataTypeTuple>(DataTypes{element, element});
|
||||
}
|
||||
|
||||
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
|
||||
{
|
||||
const auto * col_lo1 = arguments[0].column.get();
|
||||
const auto * col_hi1 = arguments[1].column.get();
|
||||
const auto * col_lo2 = arguments[2].column.get();
|
||||
const auto * col_hi2 = arguments[3].column.get();
|
||||
|
||||
auto col_res_first = ColumnUInt64::create();
|
||||
auto col_res_second = ColumnUInt64::create();
|
||||
|
||||
auto & vec_res_first = col_res_first->getData();
|
||||
vec_res_first.resize(input_rows_count);
|
||||
|
||||
auto & vec_res_second = col_res_second->getData();
|
||||
vec_res_second.resize(input_rows_count);
|
||||
|
||||
for (const auto row : ext::range(0, input_rows_count))
|
||||
{
|
||||
const UInt64 lo1 = col_lo1->getUInt(row);
|
||||
const UInt64 hi1 = col_hi1->getUInt(row);
|
||||
const UInt64 lo2 = col_lo2->getUInt(row);
|
||||
const UInt64 hi2 = col_hi2->getUInt(row);
|
||||
|
||||
S2CellId id_lo1(lo1);
|
||||
S2CellId id_hi1(hi1);
|
||||
S2CellId id_lo2(lo2);
|
||||
S2CellId id_hi2(hi2);
|
||||
|
||||
S2LatLngRect rect1(id_lo1.ToLatLng(), id_hi1.ToLatLng());
|
||||
S2LatLngRect rect2(id_lo2.ToLatLng(), id_hi2.ToLatLng());
|
||||
|
||||
S2LatLngRect rect_intersection = rect1.Intersection(rect2);
|
||||
|
||||
vec_res_first[row] = S2CellId(rect_intersection.lo()).id();
|
||||
vec_res_second[row] = S2CellId(rect_intersection.hi()).id();
|
||||
}
|
||||
|
||||
return ColumnTuple::create(Columns{std::move(col_res_first), std::move(col_res_second)});
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
void registerFunctionS2RectIntersecion(FunctionFactory & factory)
|
||||
{
|
||||
factory.registerFunction<FunctionS2RectIntersecion>();
|
||||
}
|
||||
|
||||
|
||||
}
|
148
src/Functions/S2RectUnion.cpp
Normal file
148
src/Functions/S2RectUnion.cpp
Normal file
@ -0,0 +1,148 @@
|
||||
#include "config_functions.h"
|
||||
|
||||
#include <Columns/ColumnsNumber.h>
|
||||
#include <Columns/ColumnTuple.h>
|
||||
#include <DataTypes/DataTypesNumber.h>
|
||||
#include <DataTypes/DataTypeTuple.h>
|
||||
#include <Functions/FunctionFactory.h>
|
||||
#include <Functions/IFunctionImpl.h>
|
||||
#include <Common/typeid_cast.h>
|
||||
#include <ext/range.h>
|
||||
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2cell_id.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2point.h"
|
||||
#include "../contrib/s2geometry/src/s2/s2latlng_rect.h"
|
||||
|
||||
class S2CellId;
|
||||
|
||||
namespace DB
|
||||
{
|
||||
|
||||
namespace ErrorCodes
|
||||
{
|
||||
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
|
||||
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
|
||||
}
|
||||
|
||||
namespace
|
||||
{
|
||||
|
||||
/// TODO: Comment this
|
||||
class FunctionS2RectUnion : public IFunction
|
||||
{
|
||||
public:
|
||||
static constexpr auto name = "S2RectUnion";
|
||||
|
||||
static FunctionPtr create(ContextPtr)
|
||||
{
|
||||
return std::make_shared<FunctionS2RectUnion>();
|
||||
}
|
||||
|
||||
std::string getName() const override
|
||||
{
|
||||
return name;
|
||||
}
|
||||
|
||||
size_t getNumberOfArguments() const override { return 4; }
|
||||
|
||||
bool useDefaultImplementationForConstants() const override { return true; }
|
||||
|
||||
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
|
||||
{
|
||||
size_t number_of_arguments = arguments.size();
|
||||
|
||||
if (number_of_arguments != 4) {
|
||||
throw Exception("Number of arguments for function " + getName() + " doesn't match: passed "
|
||||
+ toString(number_of_arguments) + ", should be 4",
|
||||
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
|
||||
}
|
||||
|
||||
const 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 (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(2) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
arg = arguments[2].get();
|
||||
|
||||
if (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(3) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
arg = arguments[3].get();
|
||||
|
||||
if (!WhichDataType(arg).isUInt64()) {
|
||||
throw Exception(
|
||||
"Illegal type " + arg->getName() + " of argument " + std::to_string(4) + " of function " + getName() + ". Must be UInt64",
|
||||
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
|
||||
}
|
||||
|
||||
DataTypePtr element = std::make_shared<DataTypeUInt64>();
|
||||
|
||||
return std::make_shared<DataTypeTuple>(DataTypes{element, element});
|
||||
}
|
||||
|
||||
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
|
||||
{
|
||||
const auto * col_lo1 = arguments[0].column.get();
|
||||
const auto * col_hi1 = arguments[1].column.get();
|
||||
const auto * col_lo2 = arguments[2].column.get();
|
||||
const auto * col_hi2 = arguments[3].column.get();
|
||||
|
||||
auto col_res_first = ColumnUInt64::create();
|
||||
auto col_res_second = ColumnUInt64::create();
|
||||
|
||||
auto & vec_res_first = col_res_first->getData();
|
||||
vec_res_first.resize(input_rows_count);
|
||||
|
||||
auto & vec_res_second = col_res_second->getData();
|
||||
vec_res_second.resize(input_rows_count);
|
||||
|
||||
for (const auto row : ext::range(0, input_rows_count))
|
||||
{
|
||||
const UInt64 lo1 = col_lo1->getUInt(row);
|
||||
const UInt64 hi1 = col_hi1->getUInt(row);
|
||||
const UInt64 lo2 = col_lo2->getUInt(row);
|
||||
const UInt64 hi2 = col_hi2->getUInt(row);
|
||||
|
||||
S2CellId id_lo1(lo1);
|
||||
S2CellId id_hi1(hi1);
|
||||
S2CellId id_lo2(lo2);
|
||||
S2CellId id_hi2(hi2);
|
||||
|
||||
S2LatLngRect rect1(id_lo1.ToLatLng(), id_hi1.ToLatLng());
|
||||
S2LatLngRect rect2(id_lo2.ToLatLng(), id_hi2.ToLatLng());
|
||||
|
||||
S2LatLngRect rect_union = rect1.Union(rect2);
|
||||
|
||||
vec_res_first[row] = S2CellId(rect_union.lo()).id();
|
||||
vec_res_second[row] = S2CellId(rect_union.hi()).id();
|
||||
}
|
||||
|
||||
return ColumnTuple::create(Columns{std::move(col_res_first), std::move(col_res_second)});
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
void registerFunctionS2RectUnion(FunctionFactory & factory)
|
||||
{
|
||||
factory.registerFunction<FunctionS2RectUnion>();
|
||||
}
|
||||
|
||||
|
||||
}
|
@ -47,6 +47,12 @@ void registerFunctionRadiansToS2(FunctionFactory &);
|
||||
void registerFunctionS2GetNeighbors(FunctionFactory &);
|
||||
void registerFunctionS2ToGeo(FunctionFactory &);
|
||||
void registerFunctionS2CellsIntersect(FunctionFactory &);
|
||||
void registerFunctionS2CapContains(FunctionFactory &);
|
||||
void registerFunctionS2CapUnion(FunctionFactory &);
|
||||
void registerFunctionS2RectAdd(FunctionFactory &);
|
||||
void registerFunctionS2RectContains(FunctionFactory &);
|
||||
void registerFunctionS2RectUnion(FunctionFactory &);
|
||||
void registerFunctionS2RectIntersection(FunctionFactory &);
|
||||
|
||||
|
||||
void registerFunctionsGeo(FunctionFactory & factory)
|
||||
@ -91,6 +97,12 @@ void registerFunctionsGeo(FunctionFactory & factory)
|
||||
registerFunctionS2GetNeighbors(factory);
|
||||
registerFunctionS2ToGeo(factory);
|
||||
registerFunctionS2CellsIntersect(factory);
|
||||
registerFunctionS2CapContains(factory);
|
||||
registerFunctionS2CapUnion(factory);
|
||||
registerFunctionS2RectAdd(factory);
|
||||
registerFunctionS2RectContains(factory);
|
||||
registerFunctionS2RectUnion(factory);
|
||||
registerFunctionS2RectIntersection(factory);
|
||||
}
|
||||
|
||||
}
|
||||
|
@ -0,0 +1,2 @@
|
||||
1
|
||||
0
|
2
tests/queries/0_stateless/01854_s2_cap_contains.sql
Normal file
2
tests/queries/0_stateless/01854_s2_cap_contains.sql
Normal file
@ -0,0 +1,2 @@
|
||||
select S2CapContains(1157339245694594829, 1.0, 1157347770437378819)
|
||||
select S2CapContains(1157339245694594829, 1.0, 1152921504606846977)
|
Loading…
Reference in New Issue
Block a user