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

vortex-geo/Cargo.toml

@@ -30,10 +30,20 @@ vortex-mask = { workspace = true }
 vortex-session = { workspace = true }
 wkb = { workspace = true }
 
+[features]
+# Exposes the `test_harness` module to benches; not part of the public API.
+_test-harness = []
+
 [dev-dependencies]
+divan = { workspace = true }
 rstest = { workspace = true }
 vortex-array = { workspace = true, features = ["_test-harness"] }
+vortex-geo = { path = ".", features = ["_test-harness"] }
 vortex-layout = { workspace = true }
 
+[[bench]]
+name = "envelope"
+harness = false
+
 [lints]
 workspace = true

vortex-geo/benches/envelope.rs

@@ -0,0 +1,167 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Microbenchmark for the `vortex.geo.envelope` scalar function: per-row bounding boxes over
+//! native geometry storage.
+//!
+//! Cases vary the two axes that dominate the kernel's cost profile:
+//! - nesting depth: `Point` (no `List` level, pure per-row reduction), `MultiPoint` (one level),
+//!   `MultiPolygon` (three levels);
+//! - validity: non-nullable operands, predictable sparse nulls (~10%, periodic), and
+//!   unpredictable dense nulls (~50%, pseudo-random) — the worst case for branching on validity.
+//!
+//! All cases share the same row count, so numbers are comparable across shapes.
+//!
+//! Run with `cargo bench -p vortex-geo --bench envelope`.
+
+#![expect(clippy::unwrap_used)]
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use divan::counter::ItemsCount;
+use vortex_array::ArrayRef;
+use vortex_array::Canonical;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_geo::scalar_fn::envelope::GeoEnvelope;
+use vortex_geo::test_harness::MultiPolygonRings;
+use vortex_geo::test_harness::geo_session;
+use vortex_geo::test_harness::multipoint_column;
+use vortex_geo::test_harness::multipolygon_column;
+use vortex_geo::test_harness::nullable_multipolygon_column;
+use vortex_geo::test_harness::nullable_point_column;
+use vortex_geo::test_harness::point_column;
+use vortex_session::VortexSession;
+
+fn main() {
+    divan::main();
+}
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(geo_session);
+
+/// Every case has the same row count so results are comparable across shapes: differences then
+/// reflect per-row cost (nesting depth, validity handling) rather than input size.
+const ROWS: usize = 1 << 17;
+
+/// Deterministic pseudo-random ordinate in `[0, 1000)`.
+fn ordinate(i: usize) -> f64 {
+    (i.wrapping_mul(2654435761) % 1000) as f64
+}
+
+/// A deterministic but unpredictable ~50% null pattern — the worst case for branching on
+/// validity, since the branch predictor cannot learn it (unlike a periodic `i % k` pattern).
+fn coin(i: usize) -> bool {
+    (i.wrapping_mul(2654435761) >> 13) & 1 == 0
+}
+
+/// Execute the envelope of `column` to completion.
+fn envelope(column: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef {
+    GeoEnvelope::try_new_array(column.clone())
+        .unwrap()
+        .into_array()
+        .execute::<Canonical>(ctx)
+        .unwrap()
+        .into_array()
+}
+
+#[divan::bench]
+fn point_non_nullable(bencher: Bencher) {
+    let xs = (0..ROWS).map(ordinate).collect();
+    let ys = (0..ROWS).map(|i| ordinate(i + 1)).collect();
+    let column = point_column(xs, ys).unwrap();
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| envelope(&column, &mut ctx));
+}
+
+#[divan::bench]
+fn point_mixed_validity(bencher: Bencher) {
+    let points = (0..ROWS)
+        .map(|i| (!i.is_multiple_of(10)).then(|| (ordinate(i), ordinate(i + 1))))
+        .collect();
+    let column = nullable_point_column(points).unwrap();
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| envelope(&column, &mut ctx));
+}
+
+#[divan::bench]
+fn point_random_nulls(bencher: Bencher) {
+    let points = (0..ROWS)
+        .map(|i| (!coin(i)).then(|| (ordinate(i), ordinate(i + 1))))
+        .collect();
+    let column = nullable_point_column(points).unwrap();
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| envelope(&column, &mut ctx));
+}
+
+const POINTS_PER_ROW: usize = 32;
+
+#[divan::bench]
+fn multipoint_non_nullable(bencher: Bencher) {
+    let rows = (0..ROWS)
+        .map(|r| {
+            (0..POINTS_PER_ROW)
+                .map(|i| (ordinate(r + i), ordinate(r + i + 1)))
+                .collect()
+        })
+        .collect();
+    let column = multipoint_column(rows).unwrap();
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| envelope(&column, &mut ctx));
+}
+
+const VERTICES_PER_RING: usize = 8;
+
+/// A multipolygon of two polygons with two rings each (8 vertices per ring, 32 per row), so the
+/// intermediate list levels have non-identity offsets.
+fn multipolygon_row(r: usize) -> MultiPolygonRings {
+    let ring = |p: usize| {
+        (0..VERTICES_PER_RING)
+            .map(|i| (ordinate(r + p + i), ordinate(r + p + i + 1)))
+            .collect()
+    };
+    vec![vec![ring(0), ring(1)], vec![ring(2), ring(3)]]
+}
+
+#[divan::bench]
+fn multipolygon_non_nullable(bencher: Bencher) {
+    let rows = (0..ROWS).map(multipolygon_row).collect();
+    let column = multipolygon_column(rows).unwrap();
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| envelope(&column, &mut ctx));
+}
+
+#[divan::bench]
+fn multipolygon_mixed_validity(bencher: Bencher) {
+    let rows = (0..ROWS)
+        .map(|r| (!r.is_multiple_of(10)).then(|| multipolygon_row(r)))
+        .collect();
+    let column = nullable_multipolygon_column(rows).unwrap();
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| envelope(&column, &mut ctx));
+}
+
+#[divan::bench]
+fn multipolygon_random_nulls(bencher: Bencher) {
+    let rows = (0..ROWS)
+        .map(|r| (!coin(r)).then(|| multipolygon_row(r)))
+        .collect();
+    let column = nullable_multipolygon_column(rows).unwrap();
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| envelope(&column, &mut ctx));
+}

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;