diff --git a/spatial/src/spatial/core/functions/scalar/st_hilbert.cpp b/spatial/src/spatial/core/functions/scalar/st_hilbert.cpp index b65f5048..27c80a00 100644 --- a/spatial/src/spatial/core/functions/scalar/st_hilbert.cpp +++ b/spatial/src/spatial/core/functions/scalar/st_hilbert.cpp @@ -1,10 +1,11 @@ -#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" -#include "duckdb/common/vector_operations/generic_executor.hpp" #include "duckdb/common/constants.hpp" +#include "duckdb/common/vector_operations/generic_executor.hpp" +#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" #include "spatial/common.hpp" -#include "spatial/core/functions/scalar.hpp" #include "spatial/core/functions/common.hpp" +#include "spatial/core/functions/scalar.hpp" #include "spatial/core/geometry/geometry.hpp" +#include "spatial/core/util/math.hpp" #include "spatial/core/types.hpp" #include @@ -76,6 +77,21 @@ inline uint32_t HilbertEncode(uint32_t n, uint32_t x, uint32_t y) { return ((Interleave(i1) << 1) | Interleave(i0)) >> (32 - 2 * n); } +static uint32_t FloatToUint32(float f) +{ + if (std::isnan(f)) { + return 0xFFFFFFFF; + } + uint32_t res; + memcpy(&res, &f, sizeof(res)); + if((res & 0x80000000) != 0) { + res ^= 0xFFFFFFFF; + } else { + res |= 0x80000000; + } + return res; +} + //------------------------------------------------------------------------------ // Coordinates //------------------------------------------------------------------------------ @@ -120,35 +136,61 @@ static void HilbertEncodeBoundsFunction(DataChunk &args, ExpressionState &state, using GEOM_TYPE = PrimitiveType; using UINT32_TYPE = PrimitiveType; - ArenaAllocator dummy_allocator(Allocator::DefaultAllocator()); - GenericExecutor::ExecuteBinary( input_vec, bounds_vec, result, count, [&](const GEOM_TYPE &geom_type, const BOX_TYPE &bounds) { const auto geom = geom_type.val; - if (geom.GetType() != GeometryType::POINT) { - throw InvalidInputException("ST_Hilbert only supports points"); - } - const auto point = Geometry::Deserialize(dummy_allocator, geom); - if (Point::IsEmpty(point)) { - throw InvalidInputException("ST_Hilbert does not support empty points"); - } + Box2D geom_bounds; + if(!geom.TryGetCachedBounds(geom_bounds)) { + throw InvalidInputException("ST_Hilbert(geom, bounds) requires that all geometries have a bounding box"); + } - const auto v = Point::GetVertex(point); - const auto x = v.x; - const auto y = v.y; + const auto dx = geom_bounds.min.x + (geom_bounds.max.x - geom_bounds.min.x) / 2; + const auto dy = geom_bounds.min.y + (geom_bounds.max.y - geom_bounds.min.y) / 2; const auto hilbert_width = max_hilbert / (bounds.c_val - bounds.a_val); const auto hilbert_height = max_hilbert / (bounds.d_val - bounds.b_val); // TODO: Check for overflow - const auto hilbert_x = static_cast((x - bounds.a_val) * hilbert_width); - const auto hilbert_y = static_cast((y - bounds.b_val) * hilbert_height); + const auto hilbert_x = static_cast((dx - bounds.a_val) * hilbert_width); + const auto hilbert_y = static_cast((dy - bounds.b_val) * hilbert_height); const auto h = HilbertEncode(16, hilbert_x, hilbert_y); return UINT32_TYPE {h}; }); } +//------------------------------------------------------------------------------ +// GEOMETRY +//------------------------------------------------------------------------------ +static void HilbertEncodeGeometryFunction(DataChunk &args, ExpressionState &state, Vector &result) { + const auto count = args.size(); + auto &input_vec = args.data[0]; + + UnaryExecutor::ExecuteWithNulls( + input_vec, result, count, [&](const geometry_t &geom, ValidityMask &mask, idx_t out_idx) -> uint32_t { + Box2D bounds; + if(!geom.TryGetCachedBounds(bounds)) { + mask.SetInvalid(out_idx); + return 0; + } + + Box2D bounds_f; + bounds_f.min.x = MathUtil::DoubleToFloatDown(bounds.min.x); + bounds_f.min.y = MathUtil::DoubleToFloatDown(bounds.min.y); + bounds_f.max.x = MathUtil::DoubleToFloatUp(bounds.max.x); + bounds_f.max.y = MathUtil::DoubleToFloatUp(bounds.max.y); + + const auto dx = bounds_f.min.x + (bounds_f.max.x - bounds_f.min.x) / 2; + const auto dy = bounds_f.min.y + (bounds_f.max.y - bounds_f.min.y) / 2; + + const auto hx = FloatToUint32(dx); + const auto hy = FloatToUint32(dy); + + return HilbertEncode(16, hx, hy); + }); +} + + //------------------------------------------------------------------------------ // BOX_2D/BOX_2DF //------------------------------------------------------------------------------ @@ -205,6 +247,7 @@ void CoreScalarFunctions::RegisterStHilbert(DatabaseInstance &db) { HilbertEncodeBoxFunction)); set.AddFunction(ScalarFunction({GeoTypes::BOX_2DF(), GeoTypes::BOX_2DF()}, LogicalType::UINTEGER, HilbertEncodeBoxFunction)); + set.AddFunction(ScalarFunction({GeoTypes::GEOMETRY()}, LogicalType::UINTEGER, HilbertEncodeGeometryFunction)); ExtensionUtil::RegisterFunction(db, set); DocUtil::AddDocumentation(db, "ST_Hilbert", DOC_DESCRIPTION, DOC_EXAMPLE, DOC_TAGS);