feat(vortex-geo): per-row bounding-box scalar function (vortex.geo.envelope)

vortex-geo/Cargo.toml

@@ -31,9 +31,14 @@ vortex-session = { workspace = true }
 wkb = { workspace = true }
 
 [dev-dependencies]
+divan = { workspace = true }
 rstest = { workspace = true }
 vortex-array = { workspace = true, features = ["_test-harness"] }
 vortex-layout = { workspace = true }
 
+[[bench]]
+name = "envelope_write"
+harness = false
+
 [lints]
 workspace = true

vortex-geo/benches/envelope_write.rs

@@ -0,0 +1,94 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Benchmarks for the envelope row loop's output-writing shape.
+//!
+//! `push` is the original loop: `BufferMut::with_capacity` plus four `push` calls per row —
+//! each an out-of-line call, since `push`'s grow path defeats inlining. `indexed` is the
+//! shipped loop: `BufferMut::zeroed` up front, plain indexed stores, and adjacent offset
+//! pairs instead of a bounds-checked `row_offsets[r + 1]`. The reduction is identical in
+//! both, so any difference is the write path.
+
+use divan::Bencher;
+use vortex_buffer::Buffer;
+use vortex_buffer::BufferMut;
+
+fn main() {
+    divan::main();
+}
+
+/// `(rows, vertices per row)`: per-row write overhead matters most when rows are short.
+const SHAPES: &[(usize, usize)] = &[(65536, 1), (65536, 4), (65536, 16), (16384, 64)];
+
+fn setup(rows: usize, verts: usize) -> (Vec<usize>, Vec<f64>, Vec<f64>) {
+    let n = rows * verts;
+    let val = |i: usize, salt: usize| ((i * 2654435761 + salt) % 1_000_003) as f64 * 0.001;
+    (
+        (0..=rows).map(|r| r * verts).collect(),
+        (0..n).map(|i| val(i, 17)).collect(),
+        (0..n).map(|i| val(i, 89)).collect(),
+    )
+}
+
+fn box_corners(xs: &[f64], ys: &[f64]) -> [f64; 4] {
+    let (mut xmin, mut ymin) = (f64::INFINITY, f64::INFINITY);
+    let (mut xmax, mut ymax) = (f64::NEG_INFINITY, f64::NEG_INFINITY);
+    for (&x, &y) in xs.iter().zip(ys) {
+        xmin = xmin.min(x);
+        ymin = ymin.min(y);
+        xmax = xmax.max(x);
+        ymax = ymax.max(y);
+    }
+    [xmin, ymin, xmax, ymax]
+}
+
+#[divan::bench(args = SHAPES)]
+fn push(bencher: Bencher, (rows, verts): (usize, usize)) {
+    let (row_offsets, xs, ys) = setup(rows, verts);
+    bencher.bench(|| -> [Buffer<f64>; 4] {
+        let len = rows;
+        let mut xmins = BufferMut::with_capacity(len);
+        let mut ymins = BufferMut::with_capacity(len);
+        let mut xmaxs = BufferMut::with_capacity(len);
+        let mut ymaxs = BufferMut::with_capacity(len);
+        for r in 0..len {
+            let (start, end) = (row_offsets[r], row_offsets[r + 1]);
+            let [xmin, ymin, xmax, ymax] = box_corners(&xs[start..end], &ys[start..end]);
+            xmins.push(xmin);
+            ymins.push(ymin);
+            xmaxs.push(xmax);
+            ymaxs.push(ymax);
+        }
+        [
+            xmins.freeze(),
+            ymins.freeze(),
+            xmaxs.freeze(),
+            ymaxs.freeze(),
+        ]
+    });
+}
+
+#[divan::bench(args = SHAPES)]
+fn indexed(bencher: Bencher, (rows, verts): (usize, usize)) {
+    let (row_offsets, xs, ys) = setup(rows, verts);
+    bencher.bench(|| -> [Buffer<f64>; 4] {
+        let len = rows;
+        let mut xmins = BufferMut::zeroed(len);
+        let mut ymins = BufferMut::zeroed(len);
+        let mut xmaxs = BufferMut::zeroed(len);
+        let mut ymaxs = BufferMut::zeroed(len);
+        for (r, (&start, &end)) in row_offsets.iter().zip(&row_offsets[1..]).enumerate() {
+            let [xmin, ymin, xmax, ymax] = box_corners(&xs[start..end], &ys[start..end]);
+            xmins[r] = xmin;
+            ymins[r] = ymin;
+            xmaxs[r] = xmax;
+            ymaxs[r] = ymax;
+        }
+        [
+            xmins.freeze(),
+            ymins.freeze(),
+            xmaxs.freeze(),
+            ymaxs.freeze(),
+        ]
+    });
+}

vortex-geo/src/aggregate_fn/aabb.rs

@@ -13,8 +13,6 @@ use vortex_array::aggregate_fn::AggregateFnRef;
 use vortex_array::aggregate_fn::AggregateFnVTable;
 use vortex_array::aggregate_fn::AggregateFnVTableExt;
 use vortex_array::aggregate_fn::EmptyOptions;
-use vortex_array::arrays::PrimitiveArray;
-use vortex_array::arrays::struct_::StructArrayExt;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::extension::ExtDType;
@@ -29,6 +27,8 @@ use crate::extension::GeoMetadata;
 use crate::extension::Rect;
 use crate::extension::box_storage_dtype;
 use crate::extension::coordinate::Dimension;
+use crate::extension::coordinate::box_corners;
+use crate::extension::coordinate::ordinates;
 use crate::extension::flatten_coordinates;
 use crate::extension::is_native_geometry;
 
@@ -78,18 +78,10 @@ fn aabb_storage_dtype() -> DType {
 
 /// The AABB of the raw `x`/`y` slices, or `None` when empty.
 fn aabb_of(xs: &[f64], ys: &[f64]) -> Option<GeoRect<f64>> {
-    if xs.is_empty() {
-        return None;
-    }
-    let min_max = |vals: &[f64]| {
-        vals.iter()
-            .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), &v| {
-                (lo.min(v), hi.max(v))
-            })
-    };
-    let (xmin, xmax) = min_max(xs);
-    let (ymin, ymax) = min_max(ys);
-    Some(GeoRect::new((xmin, ymin), (xmax, ymax)))
+    (!xs.is_empty()).then(|| {
+        let [xmin, ymin, xmax, ymax] = box_corners(xs, ys);
+        GeoRect::new((xmin, ymin), (xmax, ymax))
+    })
 }
 
 /// Read an AABB stat scalar (a nullable native `geoarrow.box`) into a [`GeoRect`], or `None` when
@@ -217,15 +209,9 @@ impl AggregateFnVTable for GeometryAabb {
         // `unmasked_field_by_name` reads are therefore safe. Min/max the raw x/y buffers directly:
         // cheap, and avoids `to_geometry`'s panic on empty points (which decoding would hit).
         let coords = flatten_coordinates(&array, ctx)?;
-        let xs = coords
-            .unmasked_field_by_name("x")?
-            .clone()
-            .execute::<PrimitiveArray>(ctx)?;
-        let ys = coords
-            .unmasked_field_by_name("y")?
-            .clone()
-            .execute::<PrimitiveArray>(ctx)?;
-        if let Some(rect) = aabb_of(xs.as_slice::<f64>(), ys.as_slice::<f64>()) {
+        let xs = ordinates(&coords, "x", ctx)?;
+        let ys = ordinates(&coords, "y", ctx)?;
+        if let Some(rect) = aabb_of(&xs, &ys) {
             partial.merge(rect);
         }
         Ok(())

vortex-geo/src/extension/coordinate.rs

@@ -15,12 +15,16 @@ use std::fmt::Display;
 use std::fmt::Formatter;
 
 use geoarrow::datatypes::Dimension as GeoArrowDimension;
+use vortex_array::ExecutionCtx;
+use vortex_array::arrays::StructArray;
+use vortex_array::arrays::struct_::StructArrayExt;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldNames;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
 use vortex_array::dtype::StructFields;
 use vortex_array::scalar::Scalar;
+use vortex_buffer::Buffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
@@ -189,6 +193,33 @@ pub(crate) fn coordinate_from_struct(scalar: &Scalar) -> VortexResult<Coordinate
     })
 }
 
+/// Materialize the named ordinate (`x`, `y`, ...) of a coordinate `Struct` column as a flat
+/// [`Buffer`] for bulk reads.
+pub(crate) fn ordinates(
+    coords: &StructArray,
+    name: &str,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Buffer<f64>> {
+    coords
+        .unmasked_field_by_name(name)?
+        .clone()
+        .execute::<Buffer<f64>>(ctx)
+}
+
+/// The corners of the box containing the `(xs, ys)` coordinates, in
+/// `[xmin, ymin, xmax, ymax]` order.
+pub(crate) fn box_corners(xs: &[f64], ys: &[f64]) -> [f64; 4] {
+    let (mut xmin, mut ymin) = (f64::INFINITY, f64::INFINITY);
+    let (mut xmax, mut ymax) = (f64::NEG_INFINITY, f64::NEG_INFINITY);
+    for (&x, &y) in xs.iter().zip(ys) {
+        xmin = xmin.min(x);
+        ymin = ymin.min(y);
+        xmax = xmax.max(x);
+        ymax = ymax.max(y);
+    }
+    [xmin, ymin, xmax, ymax]
+}
+
 #[cfg(test)]
 mod tests {
     use rstest::rstest;

vortex-geo/src/extension/mod.rs

@@ -47,12 +47,18 @@ use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::extension::ExtensionArrayExt;
+use vortex_array::arrays::list::ListArrayExt;
 use vortex_array::arrays::listview::ListViewArrayExt;
+use vortex_array::arrays::listview::list_from_list_view;
+use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::dtype::DType;
+use vortex_array::dtype::Nullability;
+use vortex_array::dtype::PType;
 use vortex_array::dtype::extension::ExtDType;
 use vortex_array::dtype::extension::ExtVTable;
 use vortex_array::scalar::Scalar;
 use vortex_arrow::FromArrowArray;
+use vortex_buffer::Buffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
@@ -114,6 +120,49 @@ pub(crate) fn flatten_coordinates(
     node.execute::<StructArray>(ctx)
 }
 
+/// Flatten a native geometry `storage` array to its leaf coordinates, keeping track of which
+/// row owns which coordinates.
+///
+/// Returns the flat coordinate `Struct` plus `row_offsets` (one entry per row plus a final
+/// cap): row `r`'s coordinates are `coordinates[row_offsets[r]..row_offsets[r + 1]]`, an empty
+/// range if the row has none.
+///
+/// Row boundaries are pushed down one `List` level at a time. A boundary at list `e` moves to
+/// `offsets[e]`, where that list's children start. For example, a 2-row `MultiPolygon` column —
+/// row 0 = one polygon of two rings (3 + 4 vertices), row 1 = one polygon of one ring (5
+/// vertices):
+///
+/// ```text
+/// level                offsets       row_offsets after the level
+/// rows     → polygons  [0,1,2]       [0,1,2]    row 1 starts at polygon 1
+/// polygons → rings     [0,2,3]       [0,2,3]    row 1 starts at ring 2
+/// rings    → vertices  [0,3,7,12]    [0,7,12]   row 1 starts at vertex 7
+/// ```
+pub(crate) fn flatten_row_offsets(
+    storage: ArrayRef,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<(Vec<usize>, StructArray)> {
+    let len = storage.len();
+
+    // At the outermost level, row `r` starts at element `r`; the extra entry caps the last row.
+    let mut row_offsets: Vec<usize> = (0..=len).collect();
+    let mut level = storage;
+    while matches!(level.dtype(), DType::List(..)) {
+        let list = list_from_list_view(level.execute::<Canonical>(ctx)?.into_listview(), ctx)?;
+        let offsets = list
+            .offsets()
+            .clone()
+            .cast(DType::Primitive(PType::U64, Nullability::NonNullable))?
+            .execute::<Buffer<u64>>(ctx)?;
+        for row_offset in &mut row_offsets {
+            *row_offset = usize::try_from(offsets[*row_offset])
+                .map_err(|_| vortex_err!("geo: list offset exceeds usize"))?;
+        }
+        level = list.elements().clone();
+    }
+    Ok((row_offsets, level.execute::<StructArray>(ctx)?))
+}
+
 /// Decode a native geometry column to `geo_types`. A non-geometry operand is an error.
 pub(crate) fn geometries(
     array: &ArrayRef,

vortex-geo/src/extension/rect.rs

@@ -93,7 +93,7 @@ impl ExtVTable for Rect {
 
 /// The `geoarrow.box` storage field names for `dim`: the lower corner's ordinates followed by the
 /// upper corner's.
-fn box_field_names(dim: Dimension) -> &'static [&'static str] {
+pub(crate) fn box_field_names(dim: Dimension) -> &'static [&'static str] {
     match dim {
         Dimension::Xy => &["xmin", "ymin", "xmax", "ymax"],
         Dimension::Xyz => &["xmin", "ymin", "zmin", "xmax", "ymax", "zmax"],

vortex-geo/src/lib.rs

@@ -23,6 +23,7 @@ use crate::prune::GeoDistancePrune;
 use crate::prune::GeoIntersectsPrune;
 use crate::scalar_fn::contains::GeoContains;
 use crate::scalar_fn::distance::GeoDistance;
+use crate::scalar_fn::envelope::GeoEnvelope;
 use crate::scalar_fn::intersects::GeoIntersects;
 
 pub mod aggregate_fn;
@@ -63,6 +64,7 @@ pub fn initialize(session: &VortexSession) {
     session.arrow().register_importer(Arc::new(Rect));
 
     // Register the geometry scalar functions.
+    session.scalar_fns().register(GeoEnvelope);
     session.scalar_fns().register(GeoContains);
     session.scalar_fns().register(GeoDistance);
     session.scalar_fns().register(GeoIntersects);