feat(vortex-geo): per-row 2-D AABB compute primitive

Adds vortex_geo::bounds::aabbs_2d, decoding a native geometry column to
per-row (min, max) corner pairs for row-oriented consumers, e.g.
bulk-loading an in-memory R-tree in a spatial-join operator. Deliberately
an internal compute primitive rather than a scalar function: no
registration, no SQL name, and no columnar Rect output, since the
consumer reads the boxes row by row. Empty geometries yield None; 2-D
only. Columnar consumers keep the Rect extension type, and the
GeometryAabb aggregate remains the whole-column box.

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

Author: HarukiMoriarty

vortex-geo/src/bounds.rs

@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Row-shaped bounding-box computation over native geometry columns.
+
+use geo::BoundingRect;
+use vortex_array::ArrayRef;
+use vortex_array::ExecutionCtx;
+use vortex_error::VortexResult;
+
+use crate::extension::geometries;
+
+/// Per-row 2-D axis-aligned bounding boxes (AABBs) of a native geometry column, as `(min, max)`
+/// pairs of `[x, y]` corners.
+///
+/// The per-row counterpart of the [`GeometryAabb`](crate::aggregate_fn::GeometryAabb) aggregate,
+/// which unions the whole column into one box. A point yields a degenerate box with `min == max`;
+/// a geometry with no extent (e.g. an empty polygon) yields `None`.
+///
+/// A non-geometry column is an error. Any `z`/`m` ordinates carry no bounds, this decodes to 2-D `geo_types`.
+pub fn aabbs_2d(
+    array: &ArrayRef,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<impl Iterator<Item = Option<([f64; 2], [f64; 2])>> + use<>> {
+    Ok(geometries(array, ctx)?.into_iter().map(|geometry| {
+        geometry
+            .bounding_rect()
+            .map(|rect| ([rect.min().x, rect.min().y], [rect.max().x, rect.max().y]))
+    }))
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::IntoArray;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_error::VortexResult;
+
+    use super::aabbs_2d;
+    use crate::test_harness::linestring_column;
+    use crate::test_harness::multilinestring_column;
+    use crate::test_harness::multipoint_column;
+    use crate::test_harness::multipolygon_column;
+    use crate::test_harness::point_column;
+    use crate::test_harness::polygon_column;
+    use crate::test_harness::rect_column;
+
+    /// A shorthand box from its four corners.
+    fn bounds(xmin: f64, ymin: f64, xmax: f64, ymax: f64) -> Option<([f64; 2], [f64; 2])> {
+        Some(([xmin, ymin], [xmax, ymax]))
+    }
+
+    /// A point's box is degenerate: both corners are the point itself.
+    #[test]
+    fn point_box_is_degenerate() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let points = point_column(vec![1.0, 3.0], vec![2.0, 4.0])?;
+        assert_eq!(
+            aabbs_2d(&points, &mut ctx)?.collect::<Vec<_>>(),
+            vec![bounds(1.0, 2.0, 1.0, 2.0), bounds(3.0, 4.0, 3.0, 4.0)]
+        );
+        Ok(())
+    }
+
+    /// A polygon's box is its extent: the min/max over every ring vertex, one box per row.
+    #[test]
+    fn polygon_box_is_extent() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let polygons = polygon_column(vec![
+            vec![vec![(0.0, 0.0), (4.0, 0.0), (2.0, 5.0)]],
+            vec![vec![(-3.0, 7.0), (-2.0, 7.0), (-2.0, 9.0), (-3.0, 9.0)]],
+        ])?;
+        assert_eq!(
+            aabbs_2d(&polygons, &mut ctx)?.collect::<Vec<_>>(),
+            vec![bounds(0.0, 0.0, 4.0, 5.0), bounds(-3.0, 7.0, -2.0, 9.0)]
+        );
+        Ok(())
+    }
+
+    /// A `Rect` row is its own bounding box.
+    #[test]
+    fn rect_box_is_itself() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let rects = rect_column(vec![(0.0, 0.0, 2.0, 3.0), (-1.0, -1.0, 1.0, 1.0)])?;
+        assert_eq!(
+            aabbs_2d(&rects, &mut ctx)?.collect::<Vec<_>>(),
+            vec![bounds(0.0, 0.0, 2.0, 3.0), bounds(-1.0, -1.0, 1.0, 1.0)]
+        );
+        Ok(())
+    }
+
+    /// Every multi-vertex native type over the same vertex set yields that set's box, so the
+    /// whole type family is covered (`Point` has its own degenerate-box test above).
+    #[test]
+    fn covers_every_native_geometry_type() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let vertices = vec![(1.0, 2.0), (-1.0, 5.0), (3.0, 4.0)];
+        let columns = vec![
+            linestring_column(vec![vertices.clone()])?,
+            multipoint_column(vec![vertices.clone()])?,
+            polygon_column(vec![vec![vertices.clone()]])?,
+            multilinestring_column(vec![vec![vertices.clone()]])?,
+            multipolygon_column(vec![vec![vec![vertices]]])?,
+        ];
+        for column in columns {
+            assert_eq!(
+                aabbs_2d(&column, &mut ctx)?.collect::<Vec<_>>(),
+                vec![bounds(-1.0, 2.0, 3.0, 5.0)],
+                "box mismatch for {}",
+                column.dtype()
+            );
+        }
+        Ok(())
+    }
+
+    /// An empty geometry (here a zero-part multipolygon) has no extent and yields `None`; other
+    /// rows keep their boxes.
+    #[test]
+    fn empty_geometry_has_no_box() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let multipolygons = multipolygon_column(vec![
+            vec![vec![vec![(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0)]]],
+            vec![],
+        ])?;
+        assert_eq!(
+            aabbs_2d(&multipolygons, &mut ctx)?.collect::<Vec<_>>(),
+            vec![bounds(0.0, 0.0, 1.0, 1.0), None]
+        );
+        Ok(())
+    }
+
+    /// A zero-row column yields a zero-row result.
+    #[test]
+    fn empty_column_yields_nothing() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let points = point_column(vec![], vec![])?;
+        assert_eq!(aabbs_2d(&points, &mut ctx)?.count(), 0);
+        Ok(())
+    }
+
+    /// A non-geometry column is an error.
+    #[test]
+    fn non_geometry_column_is_rejected() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let numbers = PrimitiveArray::from_iter([1.0f64, 2.0]).into_array();
+        assert!(aabbs_2d(&numbers, &mut ctx).is_err());
+        Ok(())
+    }
+}

vortex-geo/src/lib.rs

@@ -26,6 +26,7 @@ use crate::scalar_fn::distance::GeoDistance;
 use crate::scalar_fn::intersects::GeoIntersects;
 
 pub mod aggregate_fn;
+pub mod bounds;
 pub mod extension;
 pub mod prune;
 pub mod scalar_fn;