Merge pull request duckdb#389 from Maxxen/hilbert

Add new st_hilbert overload not requiring bounds parameter

Author: Maxxen

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 <cmath>
@@ -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<geometry_t>;
 	using UINT32_TYPE = PrimitiveType<uint32_t>;
 
-	ArenaAllocator dummy_allocator(Allocator::DefaultAllocator());
-
 	GenericExecutor::ExecuteBinary<GEOM_TYPE, BOX_TYPE, UINT32_TYPE>(
 	    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<double> 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<uint32_t>((x - bounds.a_val) * hilbert_width);
-		    const auto hilbert_y = static_cast<uint32_t>((y - bounds.b_val) * hilbert_height);
+		    const auto hilbert_x = static_cast<uint32_t>((dx - bounds.a_val) * hilbert_width);
+		    const auto hilbert_y = static_cast<uint32_t>((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<geometry_t, uint32_t>(
+		input_vec, result, count, [&](const geometry_t &geom, ValidityMask &mask, idx_t out_idx) -> uint32_t {
+			Box2D<double> bounds;
+			if(!geom.TryGetCachedBounds(bounds)) {
+				mask.SetInvalid(out_idx);
+				return 0;
+			}
+
+			Box2D<float> 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<double>));
 	set.AddFunction(ScalarFunction({GeoTypes::BOX_2DF(), GeoTypes::BOX_2DF()}, LogicalType::UINTEGER,
 	                               HilbertEncodeBoxFunction<float>));
+	set.AddFunction(ScalarFunction({GeoTypes::GEOMETRY()}, LogicalType::UINTEGER, HilbertEncodeGeometryFunction));
 
 	ExtensionUtil::RegisterFunction(db, set);
 	DocUtil::AddDocumentation(db, "ST_Hilbert", DOC_DESCRIPTION, DOC_EXAMPLE, DOC_TAGS);