ClickHouse/src/Functions/geohashesInBox.cpp

183 lines
7.8 KiB
C++
Raw Normal View History

2021-05-17 07:30:42 +00:00
#include <Functions/IFunction.h>
#include <Functions/FunctionFactory.h>
#include <Functions/FunctionHelpers.h>
2020-03-25 18:51:48 +00:00
#include <Functions/GeoHash.h>
#include <Columns/ColumnArray.h>
#include <Columns/ColumnString.h>
#include <DataTypes/DataTypeArray.h>
#include <DataTypes/DataTypeString.h>
#include <memory>
#include <string>
2020-06-28 14:31:57 +00:00
namespace DB
{
namespace ErrorCodes
{
extern const int LOGICAL_ERROR;
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
extern const int TOO_LARGE_ARRAY_SIZE;
}
2020-09-07 18:00:37 +00:00
namespace
{
class FunctionGeohashesInBox : public IFunction
{
public:
static constexpr auto name = "geohashesInBox";
2021-06-01 12:20:52 +00:00
static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionGeohashesInBox>(); }
String getName() const override { return name; }
size_t getNumberOfArguments() const override { return 5; }
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
{
validateArgumentType(*this, arguments, 0, isFloat, "float");
validateArgumentType(*this, arguments, 1, isFloat, "float");
validateArgumentType(*this, arguments, 2, isFloat, "float");
validateArgumentType(*this, arguments, 3, isFloat, "float");
validateArgumentType(*this, arguments, 4, isUInt8, "integer");
if (!(arguments[0]->equals(*arguments[1]) &&
arguments[0]->equals(*arguments[2]) &&
arguments[0]->equals(*arguments[3])))
{
throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
"Illegal type of argument of {} all coordinate arguments must have the same type, "
"instead they are:{}, {}, {}, {}.", getName(), arguments[0]->getName(),
arguments[1]->getName(), arguments[2]->getName(), arguments[3]->getName());
}
return std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>());
}
bool useDefaultImplementationForConstants() const override { return true; }
2021-06-22 16:21:23 +00:00
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
template <typename LonAndLatType, typename PrecisionType>
2020-08-02 02:59:09 +00:00
void execute(
const IColumn * lon_min_column,
const IColumn * lat_min_column,
const IColumn * lon_max_column,
const IColumn * lat_max_column,
const IColumn * precision_column,
ColumnPtr & result,
size_t input_rows_count) const
{
static constexpr size_t max_array_size = 10'000'000;
2020-08-02 02:59:09 +00:00
const auto * lon_min_const = typeid_cast<const ColumnConst *>(lon_min_column);
const auto * lat_min_const = typeid_cast<const ColumnConst *>(lat_min_column);
const auto * lon_max_const = typeid_cast<const ColumnConst *>(lon_max_column);
const auto * lat_max_const = typeid_cast<const ColumnConst *>(lat_max_column);
const auto * precision_const = typeid_cast<const ColumnConst *>(precision_column);
if (lon_min_const)
lon_min_column = &lon_min_const->getDataColumn();
if (lat_min_const)
lat_min_column = &lat_min_const->getDataColumn();
if (lon_max_const)
lon_max_column = &lon_max_const->getDataColumn();
if (lat_max_const)
lat_max_column = &lat_max_const->getDataColumn();
if (precision_const)
precision_column = &precision_const->getDataColumn();
const auto * lon_min = checkAndGetColumn<ColumnVector<LonAndLatType>>(lon_min_column);
const auto * lat_min = checkAndGetColumn<ColumnVector<LonAndLatType>>(lat_min_column);
const auto * lon_max = checkAndGetColumn<ColumnVector<LonAndLatType>>(lon_max_column);
const auto * lat_max = checkAndGetColumn<ColumnVector<LonAndLatType>>(lat_max_column);
2020-08-02 02:59:09 +00:00
const auto * precision = checkAndGetColumn<ColumnVector<PrecisionType>>(precision_column);
if (!lon_min || !lat_min || !lon_max || !lat_max || !precision)
{
throw Exception(ErrorCodes::LOGICAL_ERROR, "Unsupported argument types for function {} : {}, {}, {}, {}.",
getName(), lon_min_column->getName(),
lat_min_column->getName(), lon_max_column->getName(), lat_max_column->getName());
}
auto col_res = ColumnArray::create(ColumnString::create());
ColumnString & res_strings = typeid_cast<ColumnString &>(col_res->getData());
ColumnArray::Offsets & res_offsets = col_res->getOffsets();
ColumnString::Chars & res_strings_chars = res_strings.getChars();
ColumnString::Offsets & res_strings_offsets = res_strings.getOffsets();
2020-08-02 02:59:09 +00:00
for (size_t row = 0; row < input_rows_count; ++row)
{
2020-08-02 02:59:09 +00:00
const Float64 lon_min_value = lon_min->getElement(lon_min_const ? 0 : row);
const Float64 lat_min_value = lat_min->getElement(lat_min_const ? 0 : row);
const Float64 lon_max_value = lon_max->getElement(lon_max_const ? 0 : row);
const Float64 lat_max_value = lat_max->getElement(lat_max_const ? 0 : row);
const PrecisionType precision_value = precision->getElement(precision_const ? 0 : row);
2020-03-25 18:51:48 +00:00
const auto prepared_args = geohashesInBoxPrepare(
2020-08-02 02:59:09 +00:00
lon_min_value, lat_min_value, lon_max_value, lat_max_value,
precision_value);
if (prepared_args.items_count > max_array_size)
{
throw Exception(ErrorCodes::TOO_LARGE_ARRAY_SIZE, "{} would produce {} array elements, "
"which is bigger than the allowed maximum of {}",
getName(), prepared_args.items_count, max_array_size);
}
res_strings_offsets.reserve(res_strings_offsets.size() + prepared_args.items_count);
res_strings_chars.resize(res_strings_chars.size() + prepared_args.items_count * (prepared_args.precision + 1));
const auto starting_offset = res_strings_offsets.empty() ? 0 : res_strings_offsets.back();
char * out = reinterpret_cast<char *>(res_strings_chars.data() + starting_offset);
// Actually write geohashes into preallocated buffer.
2020-03-25 18:51:48 +00:00
geohashesInBox(prepared_args, out);
2020-06-03 10:27:15 +00:00
for (UInt64 i = 1; i <= prepared_args.items_count ; ++i)
res_strings_offsets.push_back(starting_offset + (prepared_args.precision + 1) * i);
2020-08-02 02:59:09 +00:00
res_offsets.push_back(res_offsets.back() + prepared_args.items_count);
}
2020-08-02 02:59:09 +00:00
if (!res_strings_offsets.empty() && res_strings_offsets.back() != res_strings_chars.size())
{
throw Exception(ErrorCodes::LOGICAL_ERROR, "String column size mismatch (internal logical error)");
}
if (!res_offsets.empty() && res_offsets.back() != res_strings.size())
{
throw Exception(ErrorCodes::LOGICAL_ERROR, "Array column size mismatch (internal logical error){} != {}",
res_offsets.back(), std::to_string(res_strings.size()));
}
result = std::move(col_res);
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
{
2020-10-19 13:42:14 +00:00
const IColumn * lon_min = arguments[0].column.get();
const IColumn * lat_min = arguments[1].column.get();
const IColumn * lon_max = arguments[2].column.get();
const IColumn * lat_max = arguments[3].column.get();
const IColumn * precision = arguments[4].column.get();
ColumnPtr res;
if (checkColumn<ColumnVector<Float32>>(lon_min))
2020-08-02 02:59:09 +00:00
execute<Float32, UInt8>(lon_min, lat_min, lon_max, lat_max, precision, res, input_rows_count);
else
2020-08-02 02:59:09 +00:00
execute<Float64, UInt8>(lon_min, lat_min, lon_max, lat_max, precision, res, input_rows_count);
2020-10-19 13:42:14 +00:00
return res;
}
};
2020-09-07 18:00:37 +00:00
}
REGISTER_FUNCTION(GeohashesInBox)
{
factory.registerFunction<FunctionGeohashesInBox>();
}
}