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

vortex-geo/src/bounds.rs

@@ -0,0 +1,236 @@
+// 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 null row, or a geometry with no extent (e.g. an empty polygon), yields `None`: either way
+/// the row has no box, and matches nothing in a spatial join.
+///
+/// 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<>> {
+    let len = array.len();
+    // Drop the null rows before decoding, since a null row has no geometry to decode (`filter`
+    // collapses an all-true mask, so an all-valid column passes through unchanged), then weave
+    // the boxes back into row order with `None` in every null slot.
+    let valid = array.validity()?.execute_mask(len, ctx)?;
+    let decoded = geometries(&array.filter(valid.clone())?, ctx)?;
+    let mut boxes = decoded.into_iter().map(|geometry| {
+        geometry
+            .bounding_rect()
+            .map(|rect| ([rect.min().x, rect.min().y], [rect.max().x, rect.max().y]))
+    });
+    Ok((0..len).map(move |row| {
+        if valid.value(row) {
+            boxes.next().flatten()
+        } else {
+            None
+        }
+    }))
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::IntoArray;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::arrays::ConstantArray;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::scalar::Scalar;
+    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::nullable_multipolygon_column;
+    use crate::test_harness::nullable_point_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 null geometry row yields `None`, just like an empty geometry: no box, so it matches
+    /// nothing in a spatial join. Valid rows keep their boxes.
+    #[test]
+    fn null_row_yields_none() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let points = nullable_point_column(vec![Some((1.0, 2.0)), None, Some((3.0, 4.0))])?;
+        assert_eq!(
+            aabbs_2d(&points, &mut ctx)?.collect::<Vec<_>>(),
+            vec![bounds(1.0, 2.0, 1.0, 2.0), None, bounds(3.0, 4.0, 3.0, 4.0)]
+        );
+        Ok(())
+    }
+
+    /// Valid, empty, and null rows stay positionally aligned: the valid row keeps its box, and
+    /// the empty and null rows are both `None`.
+    #[test]
+    fn mixed_valid_empty_null_rows_align() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let multipolygons = nullable_multipolygon_column(vec![
+            Some(vec![vec![vec![
+                (0.0, 0.0),
+                (1.0, 0.0),
+                (1.0, 1.0),
+                (0.0, 1.0),
+            ]]]),
+            Some(vec![]),
+            None,
+        ])?;
+        assert_eq!(
+            aabbs_2d(&multipolygons, &mut ctx)?.collect::<Vec<_>>(),
+            vec![bounds(0.0, 0.0, 1.0, 1.0), None, None]
+        );
+        Ok(())
+    }
+
+    /// A constant-null column carries its nulls in the array's logical validity rather than in
+    /// materialized storage; every row is `None`.
+    #[test]
+    fn constant_null_column_is_all_none() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable();
+        let nulls = ConstantArray::new(Scalar::null(dtype), 2).into_array();
+        assert_eq!(
+            aabbs_2d(&nulls, &mut ctx)?.collect::<Vec<_>>(),
+            vec![None, 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/extension/mod.rs

@@ -89,6 +89,10 @@ pub(crate) fn validate_geometry_operands(dtypes: &[DType]) -> VortexResult<()> {
 
 /// Flatten a native geometry column into a single coordinate `Struct<x, y, ...>` containing
 /// every vertex of every geometry.
+///
+/// Reads the coordinate buffers directly and ignores row validity, so a null row's placeholder
+/// coordinates are included. Callers that must exclude nulls (e.g. the AABB stat) filter the
+/// column first.
 pub(crate) fn flatten_coordinates(
     array: &ArrayRef,
     ctx: &mut ExecutionCtx,
@@ -115,6 +119,9 @@ pub(crate) fn flatten_coordinates(
 }
 
 /// Decode a native geometry column to `geo_types`. A non-geometry operand is an error.
+///
+/// Every row must be present: a null row has no decodable geometry (the per-type decoders reject
+/// it), so callers filter out null rows first.
 pub(crate) fn geometries(
     array: &ArrayRef,
     ctx: &mut ExecutionCtx,

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;

vortex-geo/src/test_harness.rs

@@ -180,6 +180,33 @@ pub(crate) fn multipolygon_column(multipolygons: Vec<MultiPolygonRings>) -> Vort
     )
 }
 
+/// A nullable `MultiPolygon` column: `None` rows are null (an empty-list placeholder in storage).
+pub(crate) fn nullable_multipolygon_column(
+    multipolygons: Vec<Option<MultiPolygonRings>>,
+) -> VortexResult<ArrayRef> {
+    let rows: Vec<MultiPolygonRings> = multipolygons
+        .iter()
+        .map(|row| row.clone().unwrap_or_default())
+        .collect();
+    let polygons: Vec<Vec<Vec<(f64, f64)>>> = rows.iter().flatten().cloned().collect();
+    let mut offsets = vec![0i32];
+    let mut len = 0usize;
+    for row in &rows {
+        len += row.len();
+        offsets.push(offset(len)?);
+    }
+    let storage = ListArray::try_new(
+        vertex_list_lists(&polygons)?,
+        PrimitiveArray::from_iter(offsets).into_array(),
+        Validity::from_iter(multipolygons.iter().map(Option::is_some)),
+    )?
+    .into_array();
+    geo_column::<MultiPolygon>(
+        storage,
+        multipolygon_storage_dtype(Dimension::Xy, Nullability::Nullable),
+    )
+}
+
 /// A 2D `Rect` (`geoarrow.box`) column over `(xmin, ymin, xmax, ymax)` boxes, stored as
 /// `Struct<xmin, ymin, xmax, ymax>`.
 pub(crate) fn rect_column(boxes: Vec<(f64, f64, f64, f64)>) -> VortexResult<ArrayRef> {