Add ST_ClipByRect

src/sgl/sgl.cpp

@@ -1,12 +1,15 @@
 #include "sgl.hpp"
 
 #include <algorithm>
+#include <cassert>
 #include <cstdlib>
 #include <cstring>
 #include <cmath>
+#include <iostream>
 #include <new>
 #include <queue>
 #include <string> // TODO: Remove this, it is only used for error messages in WKT reader
+#include <vector>
 
 //======================================================================================================================
 // Helpers
@@ -745,6 +748,105 @@ void ops::visit_polygon_geometries(const geometry &geom, void *state,
 	});
 }
 
+void ops::clip_point_to_box(const geometry &geom, const geometry &bbox, geometry &out_geom) {
+	SGL_ASSERT(geom.get_type() == geometry_type::POINT);
+	SGL_ASSERT(bbox.get_type() == geometry_type::POLYGON);
+	SGL_ASSERT(!geom.is_empty());
+
+	out_geom.set_type(geometry_type::POINT);
+	// TODO: determine if we want to keep Z.
+	out_geom.set_z(false);
+	out_geom.set_m(false);
+
+	auto rect = sgl::extent_xy::smallest();
+	if (sgl::ops::get_total_extent_xy(bbox, rect) == 0) {
+		out_geom.set_vertex_array(nullptr, 0);
+		return;
+	}
+
+	const auto vertex = geom.get_vertex_xy(0);
+	if (rect.contains(vertex)) {
+		out_geom.set_vertex_array(geom.get_vertex_array(), 1);
+	} else {
+		out_geom.set_vertex_array(nullptr, 0);
+	}
+}
+
+void ops::clip_multipoint_to_box(allocator &alloc, const geometry &geom, const geometry &bbox, geometry &out_geom) {
+	SGL_ASSERT(geom.get_type() == geometry_type::MULTI_POINT);
+	SGL_ASSERT(bbox.get_type() == geometry_type::POLYGON);
+	SGL_ASSERT(!geom.is_empty());
+
+	auto rect = sgl::extent_xy::smallest();
+	if (sgl::ops::get_total_extent_xy(bbox, rect) == 0) {
+		out_geom.set_vertex_array(nullptr, 0);
+		return;
+	}
+
+	// first pass
+	uint32_t num_points_to_return = 0;
+	visit_points(geom, [&rect, &num_points_to_return](const geometry &part) {
+		if (part.is_empty()) {
+			return;
+		}
+		const auto v_array = part.get_vertex_array();
+		const auto v_width = part.get_vertex_width();
+
+		vertex_xyzm vertex = {0, 0, 0, 0};
+		memcpy(&vertex, v_array, v_width);
+		if ((vertex.x < rect.min.x) || (vertex.x > rect.max.x) || (vertex.y < rect.min.y) || (vertex.y > rect.max.y)) {
+			return;
+		}
+		num_points_to_return += 1;
+	});
+
+	if (num_points_to_return == 0) {
+		out_geom.set_type(sgl::geometry_type::GEOMETRY_COLLECTION);
+	} else if (num_points_to_return == 1) {
+		out_geom.set_type(sgl::geometry_type::POINT);
+		// second pass
+		visit_points(geom, [&rect, &out_geom, &alloc](const geometry &part) {
+			if (part.is_empty()) {
+				return;
+			}
+			const auto v_array = part.get_vertex_array();
+			const auto v_width = part.get_vertex_width();
+
+			vertex_xyzm vertex = {0, 0, 0, 0};
+			memcpy(&vertex, v_array, v_width);
+			if ((vertex.x < rect.min.x) || (vertex.x > rect.max.x) || (vertex.y < rect.min.y) || (vertex.y > rect.max.y)) {
+				return;
+			}
+			const auto data_mem = static_cast<char *>(alloc.alloc(v_width));
+			memcpy(data_mem, &vertex, v_width);
+			out_geom.set_vertex_array(data_mem, 1);
+		});
+	} else {
+		out_geom.set_type(sgl::geometry_type::MULTI_POINT);
+		// second pass
+		visit_points(geom, [&rect, &out_geom, &alloc](const geometry &part) {
+			if (part.is_empty()) {
+				return;
+			}
+			const auto v_array = part.get_vertex_array();
+			const auto v_width = part.get_vertex_width();
+
+			vertex_xyzm vertex = {0, 0, 0, 0};
+			memcpy(&vertex, v_array, v_width);
+			if ((vertex.x < rect.min.x) || (vertex.x > rect.max.x) || (vertex.y < rect.min.y) || (vertex.y > rect.max.y)) {
+				return;
+			}
+			const auto data_mem = static_cast<char *>(alloc.alloc(v_width));
+			memcpy(data_mem, &vertex, v_width);
+			auto point_mem = static_cast<char *>(alloc.alloc(sizeof(sgl::geometry)));
+			auto point_ptr = new (point_mem) sgl::geometry(geometry_type::POINT, part.has_z(), part.has_m());
+			point_ptr->set_vertex_array(data_mem, 1);
+			out_geom.append_part(point_ptr);
+		});
+	}
+}
+
+
 } // namespace sgl
 
 //======================================================================================================================

src/sgl/sgl.hpp

@@ -703,6 +703,9 @@ void locate_along(allocator &alloc, const geometry &linear_geom, double measure,
 void locate_between(allocator &alloc, const geometry &linear_geom, double measure_beg, double measure_end,
                     double offset, geometry &out_geom);
 
+void clip_point_to_box(const geometry &geom, const geometry &bbox, geometry &out_geom);
+void clip_multipoint_to_box(allocator &alloc, const geometry &geom, const geometry &bbox, geometry &out_geom);
+
 } // namespace ops
 
 // TODO: Move these

src/spatial/modules/geos/geos_geometry.hpp

@@ -91,6 +91,7 @@ class GeosGeometry {
 	GeosGeometry get_coverage_invalid_edges(double tolerance) const;
 	GeosGeometry get_coverage_simplified(double tolerance, bool preserve_boundary) const;
 	GeosGeometry get_coverage_union() const;
+	GeosGeometry clip_by_rect(double x_min, double y_min, double x_max, double y_max) const;
 
 	PreparedGeosGeometry get_prepared() const;
 
@@ -480,6 +481,10 @@ inline PreparedGeosGeometry GeosGeometry::get_prepared() const {
 	return PreparedGeosGeometry(handle, *this);
 }
 
+inline GeosGeometry GeosGeometry::clip_by_rect(double x_min, double y_min, double x_max, double y_max) const {
+	return GeosGeometry(handle, GEOSClipByRect_r(handle, geom, x_min, y_min, x_max, y_max));
+}
+
 //-- PreparedGeosGeometry --//
 
 inline bool PreparedGeosGeometry::contains(const GeosGeometry &other) const {

src/spatial/modules/geos/geos_module.cpp

@@ -2773,4 +2773,26 @@ void RegisterGEOSModule(ExtensionLoader &loader) {
 	ST_CoverageSimplify_Agg::Register(loader);
 }
 
+string_t GeosOperations::clip_to_rect(Vector &result, const string_t &blob, double x_min, double y_min, double x_max, double y_max) {
+	auto ctx = GEOS_init_r();
+
+	GEOSContext_setErrorMessageHandler_r(ctx, [](const char *message, void *) { throw InvalidInputException(message); }, nullptr);
+
+	const auto geom = GeosSerde::Deserialize(ctx, blob.GetData(), blob.GetSize());
+	if (geom == nullptr) {
+		GEOS_finish_r(ctx);
+		throw InvalidInputException("Could not deserialize geometry");
+	}
+	auto geos = GeosGeometry(ctx, geom);
+	auto clipped = geos.clip_by_rect(x_min, y_min, x_max, y_max);
+	const auto raw = clipped.get_raw();
+	const auto size = GeosSerde::GetRequiredSize(ctx, raw);
+	string_t serialized = StringVector::EmptyString(result, size);
+	const auto ptr = serialized.GetDataWriteable();
+	GeosSerde::Serialize(ctx, raw, ptr, size);
+	serialized.Finalize();
+	GEOS_finish_r(ctx);
+	return serialized;
+}
+
 } // namespace duckdb

src/spatial/modules/geos/geos_module.hpp

@@ -1,9 +1,16 @@
 #pragma once
 
+#include "duckdb/common/types/string_type.hpp"
+
 namespace duckdb {
 
 class ExtensionLoader;
 
 void RegisterGEOSModule(ExtensionLoader &loader);
 
+class GeosOperations {
+public:
+    static string_t clip_to_rect(Vector &result, const string_t &blob, double x_min, double y_min, double x_max, double y_max);
+};
+
 } // namespace duckdb

src/spatial/modules/main/spatial_functions_scalar.cpp

@@ -22,6 +22,10 @@
 // Extra
 #include "yyjson.h"
 
+#if SPATIAL_USE_GEOS
+#include "spatial/modules/geos/geos_module.hpp"
+#endif
+
 namespace duckdb {
 
 namespace {
@@ -1641,6 +1645,84 @@ struct ST_Centroid {
 	}
 };
 
+//======================================================================================================================
+// ST_ClipByBox2D
+//======================================================================================================================
+
+struct ST_ClipByBox2D {
+
+	//------------------------------------------------------------------------------------------------------------------
+	// GEOMETRY
+	//------------------------------------------------------------------------------------------------------------------
+	static void ExecuteGeometry(DataChunk &args, ExpressionState &state, Vector &result) {
+		auto &lstate = LocalState::ResetAndGet(state);
+		auto &alloc = lstate.GetAllocator();
+
+		BinaryExecutor::Execute<string_t, string_t, string_t>(
+		    args.data[0], args.data[1], result, args.size(),
+		    [&](const string_t &geom_blob, const string_t &bbox_blob) {
+			    sgl::geometry geom;
+				sgl::geometry bbox;
+			    lstate.Deserialize(geom_blob, geom);
+				lstate.Deserialize(bbox_blob, bbox);
+
+				if (geom.get_type() == sgl::geometry_type::POINT) {
+					sgl::geometry clipped;
+					sgl::ops::clip_point_to_box(geom, bbox, clipped);
+					return lstate.Serialize(result, clipped);
+				} else if (geom.get_type() == sgl::geometry_type::MULTI_POINT) {
+					sgl::geometry clipped;
+					sgl::ops::clip_multipoint_to_box(alloc, geom, bbox, clipped);
+					return lstate.Serialize(result, clipped);
+				} else {
+#if SPATIAL_USE_GEOS
+					auto rect = sgl::extent_xy::smallest();
+					if (sgl::ops::get_total_extent_xy(bbox, rect) == 0) {
+						sgl::geometry clipped;
+						clipped.set_vertex_array(nullptr, 0);
+						return lstate.Serialize(result, clipped);
+					}
+					return GeosOperations::clip_to_rect(result, geom_blob, rect.min.x, rect.min.y, rect.max.x, rect.max.y);
+#else
+				    throw InvalidInputException("ST_ClipByBox2D: input is not a POINT or MULTIPOINT, and GEOS fallback is not available");
+#endif
+			    }
+		    });
+	}
+
+	//------------------------------------------------------------------------------------------------------------------
+	// Documentation
+	//------------------------------------------------------------------------------------------------------------------
+	static constexpr auto DESCRIPTION = R"(
+		Clips a geometry by a rectangular bounding box (BBOX).
+
+		This may produce a different kind of geometry. For example, a linestring can be converted into a multilinestring.
+	)";
+	static constexpr auto EXAMPLE = "";
+
+	//------------------------------------------------------------------------------------------------------------------
+	// Register
+	//------------------------------------------------------------------------------------------------------------------
+	static void Register(ExtensionLoader &loader) {
+		FunctionBuilder::RegisterScalar(loader, "ST_ClipByBox2D", [](ScalarFunctionBuilder &func) {
+			func.AddVariant([](ScalarFunctionVariantBuilder &variant) {
+				variant.AddParameter("geom", GeoTypes::GEOMETRY());
+				variant.AddParameter("box", GeoTypes::GEOMETRY());
+				variant.SetReturnType(GeoTypes::GEOMETRY());
+
+				variant.SetFunction(ExecuteGeometry);
+				variant.SetInit(LocalState::Init);
+			});
+
+			func.SetDescription(DESCRIPTION);
+			func.SetExample(EXAMPLE);
+
+			func.SetTag("ext", "spatial");
+			func.SetTag("category", "construction");
+		});
+	}
+};
+
 //======================================================================================================================
 // ST_Collect
 //======================================================================================================================
@@ -9271,6 +9353,7 @@ void RegisterSpatialScalarFunctions(ExtensionLoader &loader) {
 	ST_AsSVG::Register(loader);
 	ST_Azimuth::Register(loader);
 	ST_Centroid::Register(loader);
+	ST_ClipByBox2D::Register(loader);
 	ST_Collect::Register(loader);
 	ST_CollectionExtract::Register(loader);
 	ST_Contains::Register(loader);