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

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,7 @@ use crate::extension::GeoMetadata;
 use crate::extension::Rect;
 use crate::extension::box_storage_dtype;
 use crate::extension::coordinate::Dimension;
+use crate::extension::coordinate::f64_field;
 use crate::extension::flatten_coordinates;
 use crate::extension::is_native_geometry;
 
@@ -217,14 +216,8 @@ 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)?;
+        let xs = f64_field(&coords, "x", ctx)?;
+        let ys = f64_field(&coords, "y", ctx)?;
         if let Some(rect) = aabb_of(xs.as_slice::<f64>(), ys.as_slice::<f64>()) {
             partial.merge(rect);
         }

vortex-geo/src/extension/coordinate.rs

@@ -15,6 +15,10 @@ use std::fmt::Display;
 use std::fmt::Formatter;
 
 use geoarrow::datatypes::Dimension as GeoArrowDimension;
+use vortex_array::ExecutionCtx;
+use vortex_array::arrays::PrimitiveArray;
+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;
@@ -189,6 +193,19 @@ pub(crate) fn coordinate_from_struct(scalar: &Scalar) -> VortexResult<Coordinate
     })
 }
 
+/// Materialize a named non-nullable `f64` field of a coordinate `Struct` column as a
+/// [`PrimitiveArray`], for bulk per-ordinate reads.
+pub(crate) fn f64_field(
+    coords: &StructArray,
+    name: &str,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<PrimitiveArray> {
+    coords
+        .unmasked_field_by_name(name)?
+        .clone()
+        .execute::<PrimitiveArray>(ctx)
+}
+
 #[cfg(test)]
 mod tests {
     use rstest::rstest;

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);