From 72aa947b4164161d69a77c2ec6818a50f8ef540b Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 21 Jul 2026 10:13:40 -0400 Subject: [PATCH 1/8] feat(vortex-geo): per-row bounding-box scalar function (vortex.geo.envelope) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `GeoEnvelope` computes the per-row 2-D axis-aligned bounding box of a native geometry column, returned as a native geoarrow.box (`Rect`) column — the per-row counterpart of the `GeometryAabb` aggregate (whole-column). Intended for row-oriented consumers such as bulk-loading an in-memory R-tree in a spatial-join operator, which read the resulting box column back row by row. Computed directly over the native coordinate storage — no `geo_types` decode and no Arrow round-trip: walk the nested `ListView` storage down to the leaf x/y buffers, tracking which top-level row owns each coordinate, then min/max per row. A `Rect` input is the identity; a null row or an empty geometry yields a null box, so the output is always nullable. Hoist the shared `f64_field` coordinate-column accessor into `extension::coordinate`, used by both this function and the aggregate. Signed-off-by: Nemo Yu --- vortex-geo/src/aggregate_fn/aabb.rs | 13 +- vortex-geo/src/extension/coordinate.rs | 17 + vortex-geo/src/lib.rs | 2 + vortex-geo/src/scalar_fn/envelope.rs | 421 +++++++++++++++++++++++++ vortex-geo/src/scalar_fn/mod.rs | 1 + vortex-geo/src/test_harness.rs | 56 ++++ 6 files changed, 500 insertions(+), 10 deletions(-) create mode 100644 vortex-geo/src/scalar_fn/envelope.rs diff --git a/vortex-geo/src/aggregate_fn/aabb.rs b/vortex-geo/src/aggregate_fn/aabb.rs index 81f4450a674..80bd412a121 100644 --- a/vortex-geo/src/aggregate_fn/aabb.rs +++ b/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::(ctx)?; - let ys = coords - .unmasked_field_by_name("y")? - .clone() - .execute::(ctx)?; + let xs = f64_field(&coords, "x", ctx)?; + let ys = f64_field(&coords, "y", ctx)?; if let Some(rect) = aabb_of(xs.as_slice::(), ys.as_slice::()) { partial.merge(rect); } diff --git a/vortex-geo/src/extension/coordinate.rs b/vortex-geo/src/extension/coordinate.rs index b6e324cd457..037e4b52b8e 100644 --- a/vortex-geo/src/extension/coordinate.rs +++ b/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 VortexResult { + coords + .unmasked_field_by_name(name)? + .clone() + .execute::(ctx) +} + #[cfg(test)] mod tests { use rstest::rstest; diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index 3a0c52f1eb9..e823c6301a6 100644 --- a/vortex-geo/src/lib.rs +++ b/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); diff --git a/vortex-geo/src/scalar_fn/envelope.rs b/vortex-geo/src/scalar_fn/envelope.rs new file mode 100644 index 00000000000..c8c5902446e --- /dev/null +++ b/vortex-geo/src/scalar_fn/envelope.rs @@ -0,0 +1,421 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `envelope`: the per-row 2-D axis-aligned bounding box (AABB) of a native geometry column. +//! +//! A row-oriented consumer (e.g. bulk-loading an in-memory R-tree in a spatial-join operator) +//! reads the resulting box column back row by row. + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::expr::Expression; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +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::validate_geometry_operands; + +/// `envelope`: the axis-aligned bounding box of each geometry in a native geometry operand (a column +/// or a constant literal), as a native 2-D `geoarrow.box` ([`Rect`]) column. +/// +/// 2-D only: only the `x`/`y` leaf ordinates are read, so any `z`/`m` are ignored and each box is +/// the XY extent — matching the [`GeometryAabb`](crate::aggregate_fn::GeometryAabb) aggregate. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct GeoEnvelope; + +impl GeoEnvelope { + /// A lazy `ScalarFnArray` computing the per-row bounding box of geometry operand `a`, which may + /// be constant. The output length is taken from `a`. + pub fn try_new_array(a: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(GeoEnvelope, EmptyOptions).erased(), + vec![a], + ) + } +} + +/// The output extension type: a nullable native 2-D `geoarrow.box` ([`Rect`]). Nullable because a +/// null row has no box. Metadata is defaulted. +fn output_ext_dtype() -> VortexResult> { + ExtDType::::try_new( + GeoMetadata::default(), + box_storage_dtype(Dimension::Xy, Nullability::Nullable), + ) +} + +/// The per-row 2-D bounding box of `array` as a nullable `Rect` (`geoarrow.box`) column, computed +/// directly over the native coordinate storage — no decode to `geo_types`, no Arrow round-trip. +/// +/// Walks the nested `List` storage down to the leaf coordinate `Struct`, carrying which top-level +/// row owns each coordinate, then min/maxes x/y per row over the raw `f64` buffers. The ring/part +/// nesting is irrelevant to a bounding box, only which coordinates belong to a row matters. A null +/// row, or a valid row that owns no coordinate (an empty geometry), yields a null box. +fn envelopes(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let len = array.len(); + let valid = array.validity()?.execute_mask(len, ctx)?; + let ext = array + .dtype() + .as_extension_opt() + .ok_or_else(|| vortex_err!("geo: envelope operand is not a geometry extension type"))?; + let storage = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + + // Per-row min/max accumulators; `seen[r]` marks a row that owns at least one coordinate. + let mut lo = vec![(f64::INFINITY, f64::INFINITY); len]; + let mut hi = vec![(f64::NEG_INFINITY, f64::NEG_INFINITY); len]; + let mut seen = vec![false; len]; + + if ext.is::() { + // The bounding box of a box is itself: read its min/max fields straight from storage. + let coords = storage.execute::(ctx)?; + let xmin = f64_field(&coords, "xmin", ctx)?; + let ymin = f64_field(&coords, "ymin", ctx)?; + let xmax = f64_field(&coords, "xmax", ctx)?; + let ymax = f64_field(&coords, "ymax", ctx)?; + let (xmin, ymin) = (xmin.as_slice::(), ymin.as_slice::()); + let (xmax, ymax) = (xmax.as_slice::(), ymax.as_slice::()); + for r in 0..len { + lo[r] = (xmin[r], ymin[r]); + hi[r] = (xmax[r], ymax[r]); + seen[r] = true; + } + } else { + // Walk any `List` nesting down to the leaf coordinate `Struct`, mapping each element to the + // top-level row that owns it. `ListView` offsets/sizes are arbitrary integer types and need + // not be contiguous, so map every element explicitly rather than composing offset arrays. + let u64_dtype = DType::Primitive(PType::U64, Nullability::NonNullable); + let mut owner: Vec = (0..len).collect(); + let mut node = storage; + while matches!(node.dtype(), DType::List(..)) { + let list = node.execute::(ctx)?.into_listview(); + let offsets = list + .offsets() + .clone() + .cast(u64_dtype.clone())? + .execute::(ctx)?; + let sizes = list + .sizes() + .clone() + .cast(u64_dtype.clone())? + .execute::(ctx)?; + let (offsets, sizes) = (offsets.as_slice::(), sizes.as_slice::()); + let elements = list.elements().clone(); + let mut child = vec![0usize; elements.len()]; + for (i, &row) in owner.iter().enumerate() { + let start = usize::try_from(offsets[i]) + .map_err(|_| vortex_err!("geo: list offset exceeds usize"))?; + let size = usize::try_from(sizes[i]) + .map_err(|_| vortex_err!("geo: list size exceeds usize"))?; + child[start..start + size].fill(row); + } + owner = child; + node = elements; + } + let coords = node.execute::(ctx)?; + let xs = f64_field(&coords, "x", ctx)?; + let ys = f64_field(&coords, "y", ctx)?; + for ((&r, &x), &y) in owner + .iter() + .zip(xs.as_slice::()) + .zip(ys.as_slice::()) + { + lo[r] = (lo[r].0.min(x), lo[r].1.min(y)); + hi[r] = (hi[r].0.max(x), hi[r].1.max(y)); + seen[r] = true; + } + } + + // Build the nullable `Rect` column straight from the accumulators: a row is present iff it owns + // a coordinate and its operand row is valid; absent rows are null (physical placeholder `0.0`). + let present: Vec = (0..len).map(|r| seen[r] && valid.value(r)).collect(); + let ordinate = |src: &[(f64, f64)], axis: fn((f64, f64)) -> f64| { + PrimitiveArray::from_iter((0..len).map(|r| if present[r] { axis(src[r]) } else { 0.0 })) + .into_array() + }; + let storage = StructArray::try_new( + FieldNames::from(["xmin", "ymin", "xmax", "ymax"]), + vec![ + ordinate(&lo, |c| c.0), + ordinate(&lo, |c| c.1), + ordinate(&hi, |c| c.0), + ordinate(&hi, |c| c.1), + ], + len, + Validity::from_iter(present.iter().copied()), + )? + .into_array(); + Ok(ExtensionArray::try_new(output_ext_dtype()?.erased(), storage)?.into_array()) +} + +impl ScalarFnVTable for GeoEnvelope { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.envelope"); + *ID + } + + fn serialize(&self, _: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("geometry"), + _ => unreachable!("envelope has exactly one child"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_geometry_operands(dtypes)?; + // Always nullable: an empty geometry has no box, so nulls can appear even over a + // non-nullable operand. + Ok(DType::Extension(output_ext_dtype()?.erased())) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let array = args.get(0)?; + envelopes(&array, ctx) + } + + fn validity(&self, _: &Self::Options, _: &Expression) -> VortexResult> { + // The output null mask is not derivable from the operand's validity alone: an empty + // geometry yields a null box even where the operand is valid. Let the planner execute. + Ok(None) + } + + fn is_null_sensitive(&self, _: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_error::VortexResult; + + use super::GeoEnvelope; + 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::nullable_rect_column; + use crate::test_harness::point_column; + use crate::test_harness::polygon_column; + use crate::test_harness::rect_column; + + /// Execute a `GeoEnvelope` over `array`, returning the lazy box column. + fn boxes(array: ArrayRef) -> VortexResult { + Ok(GeoEnvelope::try_new_array(array)?.into_array()) + } + + /// A point's box is degenerate: both corners are the point itself. + #[test] + fn point_box_is_degenerate() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + + let points = point_column(vec![1.0, 3.0], vec![2.0, 4.0])?; + let expected = + nullable_rect_column(vec![Some((1.0, 2.0, 1.0, 2.0)), Some((3.0, 4.0, 3.0, 4.0))])?; + assert_arrays_eq!(boxes(points)?, expected, &mut ctx); + 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 = crate::test_harness::geo_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)]], + ])?; + let expected = nullable_rect_column(vec![ + Some((0.0, 0.0, 4.0, 5.0)), + Some((-3.0, 7.0, -2.0, 9.0)), + ])?; + assert_arrays_eq!(boxes(polygons)?, expected, &mut ctx); + Ok(()) + } + + /// A `Rect` row is its own bounding box. + #[test] + fn rect_box_is_itself() -> VortexResult<()> { + let session = crate::test_harness::geo_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)])?; + let expected = nullable_rect_column(vec![ + Some((0.0, 0.0, 2.0, 3.0)), + Some((-1.0, -1.0, 1.0, 1.0)), + ])?; + assert_arrays_eq!(boxes(rects)?, expected, &mut ctx); + 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 = crate::test_harness::geo_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]]])?, + ]; + let expected = nullable_rect_column(vec![Some((-1.0, 2.0, 3.0, 5.0))])?; + for column in columns { + assert_arrays_eq!(boxes(column)?, expected.clone(), &mut ctx); + } + Ok(()) + } + + /// An empty geometry (here a zero-part multipolygon) has no extent and yields a null box; other + /// rows keep their boxes. + #[test] + fn empty_geometry_has_no_box() -> VortexResult<()> { + let session = crate::test_harness::geo_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![], + ])?; + let expected = nullable_rect_column(vec![Some((0.0, 0.0, 1.0, 1.0)), None])?; + assert_arrays_eq!(boxes(multipolygons)?, expected, &mut ctx); + Ok(()) + } + + /// A null geometry row yields a null box, just like an empty geometry; valid rows keep theirs. + #[test] + fn null_row_yields_null_box() -> VortexResult<()> { + let session = crate::test_harness::geo_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))])?; + let expected = nullable_rect_column(vec![ + Some((1.0, 2.0, 1.0, 2.0)), + None, + Some((3.0, 4.0, 3.0, 4.0)), + ])?; + assert_arrays_eq!(boxes(points)?, expected, &mut ctx); + Ok(()) + } + + /// Valid, empty, and null rows stay positionally aligned: the valid row keeps its box, and the + /// empty and null rows are both null. + #[test] + fn mixed_valid_empty_null_rows_align() -> VortexResult<()> { + let session = crate::test_harness::geo_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, + ])?; + let expected = nullable_rect_column(vec![Some((0.0, 0.0, 1.0, 1.0)), None, None])?; + assert_arrays_eq!(boxes(multipolygons)?, expected, &mut ctx); + Ok(()) + } + + /// A constant-null operand yields an all-null box column. + #[test] + fn constant_null_is_all_null() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable(); + let null_const = ConstantArray::new(Scalar::null(point_dtype), 2).into_array(); + let expected = nullable_rect_column(vec![None, None])?; + assert_arrays_eq!(boxes(null_const)?, expected, &mut ctx); + Ok(()) + } + + /// Output is always nullable, even over a non-nullable operand, since an empty geometry has no + /// box. + #[test] + fn output_is_always_nullable() -> VortexResult<()> { + let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + assert!(!dtype.is_nullable()); + let out = GeoEnvelope.return_dtype(&EmptyOptions, &[dtype])?; + assert!(out.is_nullable()); + Ok(()) + } + + /// A non-geometry operand dtype is rejected up front, before execution. + #[test] + fn non_geometry_operand_is_rejected() -> VortexResult<()> { + let numeric = DType::Primitive(PType::I32, Nullability::NonNullable); + assert!(GeoEnvelope.return_dtype(&EmptyOptions, &[numeric]).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/mod.rs b/vortex-geo/src/scalar_fn/mod.rs index 7ea034896dc..e6770be4fff 100644 --- a/vortex-geo/src/scalar_fn/mod.rs +++ b/vortex-geo/src/scalar_fn/mod.rs @@ -5,5 +5,6 @@ pub mod contains; pub mod distance; +pub mod envelope; mod execute; pub mod intersects; diff --git a/vortex-geo/src/test_harness.rs b/vortex-geo/src/test_harness.rs index f023b727089..593462eb662 100644 --- a/vortex-geo/src/test_harness.rs +++ b/vortex-geo/src/test_harness.rs @@ -180,6 +180,33 @@ pub(crate) fn multipolygon_column(multipolygons: Vec) -> Vort ) } +/// A nullable `MultiPolygon` column: `None` rows are null (an empty-list placeholder in storage). +pub(crate) fn nullable_multipolygon_column( + multipolygons: Vec>, +) -> VortexResult { + let rows: Vec = multipolygons + .iter() + .map(|row| row.clone().unwrap_or_default()) + .collect(); + let polygons: Vec>> = 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::( + storage, + multipolygon_storage_dtype(Dimension::Xy, Nullability::Nullable), + ) +} + /// A 2D `Rect` (`geoarrow.box`) column over `(xmin, ymin, xmax, ymax)` boxes, stored as /// `Struct`. pub(crate) fn rect_column(boxes: Vec<(f64, f64, f64, f64)>) -> VortexResult { @@ -199,6 +226,35 @@ pub(crate) fn rect_column(boxes: Vec<(f64, f64, f64, f64)>) -> VortexResult>, +) -> VortexResult { + let len = boxes.len(); + let field = |select: fn(&(f64, f64, f64, f64)) -> f64| { + PrimitiveArray::from_iter(boxes.iter().map(|b| b.as_ref().map_or(0.0, select))).into_array() + }; + let storage = StructArray::try_new( + FieldNames::from(["xmin", "ymin", "xmax", "ymax"]), + vec![ + field(|b| b.0), + field(|b| b.1), + field(|b| b.2), + field(|b| b.3), + ], + len, + Validity::from_iter(boxes.iter().map(Option::is_some)), + )? + .into_array(); + let ext = ExtDType::::try_new( + GeoMetadata::default(), + box_storage_dtype(Dimension::Xy, Nullability::Nullable), + )?; + Ok(ExtensionArray::try_new(ext.erased(), storage)?.into_array()) +} + /// Decode a [`Coordinate`] from an extension-typed point scalar (unwrapped to its coordinate /// storage) or a bare coordinate `Struct` scalar — used to read back a single point in assertions. pub(crate) fn coordinate_from_scalar(scalar: &Scalar) -> VortexResult { From e4ee8e70f7b4e3620de22266e363ccc763a86eb4 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 21 Jul 2026 14:03:41 -0400 Subject: [PATCH 2/8] refactor(vortex-geo): compute envelope via kernel-based per-row flatten Address review feedback on the envelope scalar function: - Replace the per-element owner-scatter with flatten_row_offsets: each List level converts through list_from_list_view (exact ListArray layout), row boundaries compose through the offsets, and each row then min/maxes a contiguous slice of the raw f64 buffers. This also fixes sliced/scattered views, where unreferenced elements were previously attributed to row 0. - Short-circuit Rect operands: a box is its own envelope, so project the 2-D corner fields lazily and keep the operand validity lazy too. - ordinates (was f64_field) returns Buffer directly; share box_corners and box_field_names between the envelope scalar fn and the GeometryAabb aggregate instead of duplicating them. - New tests: sliced operands, uneven nesting across levels, inner-level empty geometries, and 3-D Rect operands. Signed-off-by: Nemo Yu --- vortex-geo/src/aggregate_fn/aabb.rs | 25 +-- vortex-geo/src/extension/coordinate.rs | 26 ++- vortex-geo/src/extension/mod.rs | 47 ++++ vortex-geo/src/extension/rect.rs | 2 +- vortex-geo/src/scalar_fn/envelope.rs | 300 +++++++++++++++---------- 5 files changed, 264 insertions(+), 136 deletions(-) diff --git a/vortex-geo/src/aggregate_fn/aabb.rs b/vortex-geo/src/aggregate_fn/aabb.rs index 80bd412a121..8ac5c9e2452 100644 --- a/vortex-geo/src/aggregate_fn/aabb.rs +++ b/vortex-geo/src/aggregate_fn/aabb.rs @@ -27,7 +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::f64_field; +use crate::extension::coordinate::box_corners; +use crate::extension::coordinate::ordinates; use crate::extension::flatten_coordinates; use crate::extension::is_native_geometry; @@ -77,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> { - 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 @@ -216,9 +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 = f64_field(&coords, "x", ctx)?; - let ys = f64_field(&coords, "y", ctx)?; - if let Some(rect) = aabb_of(xs.as_slice::(), ys.as_slice::()) { + let xs = ordinates(&coords, "x", ctx)?; + let ys = ordinates(&coords, "y", ctx)?; + if let Some(rect) = aabb_of(&xs, &ys) { partial.merge(rect); } Ok(()) diff --git a/vortex-geo/src/extension/coordinate.rs b/vortex-geo/src/extension/coordinate.rs index 037e4b52b8e..dc3537cfb30 100644 --- a/vortex-geo/src/extension/coordinate.rs +++ b/vortex-geo/src/extension/coordinate.rs @@ -16,7 +16,6 @@ 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; @@ -25,6 +24,7 @@ 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; @@ -193,17 +193,31 @@ pub(crate) fn coordinate_from_struct(scalar: &Scalar) -> VortexResult VortexResult { +) -> VortexResult> { coords .unmasked_field_by_name(name)? .clone() - .execute::(ctx) + .execute::>(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)] diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 1a3db5a7b3e..d7c3ee912cd 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/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,47 @@ pub(crate) fn flatten_coordinates( node.execute::(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, StructArray)> { + // At the outermost level, row `r` starts at element `r`; the extra entry caps the last row. + let mut row_offsets: Vec = (0..=storage.len()).collect(); + let mut level = storage; + while matches!(level.dtype(), DType::List(..)) { + let list = list_from_list_view(level.execute::(ctx)?.into_listview(), ctx)?; + let offsets = list + .offsets() + .clone() + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::>(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::(ctx)?)) +} + /// Decode a native geometry column to `geo_types`. A non-geometry operand is an error. pub(crate) fn geometries( array: &ArrayRef, diff --git a/vortex-geo/src/extension/rect.rs b/vortex-geo/src/extension/rect.rs index 3576966368c..4c373a1807f 100644 --- a/vortex-geo/src/extension/rect.rs +++ b/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"], diff --git a/vortex-geo/src/scalar_fn/envelope.rs b/vortex-geo/src/scalar_fn/envelope.rs index c8c5902446e..f83c1527745 100644 --- a/vortex-geo/src/scalar_fn/envelope.rs +++ b/vortex-geo/src/scalar_fn/envelope.rs @@ -7,20 +7,16 @@ //! reads the resulting box column back row by row. use vortex_array::ArrayRef; -use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrays::listview::ListViewArrayExt; -use vortex_array::builtins::ArrayBuiltins; +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::extension::ExtDType; use vortex_array::expr::Expression; use vortex_array::scalar_fn::Arity; @@ -31,16 +27,21 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::GeoMetadata; use crate::extension::Rect; +use crate::extension::box_field_names; use crate::extension::box_storage_dtype; use crate::extension::coordinate::Dimension; -use crate::extension::coordinate::f64_field; +use crate::extension::coordinate::box_corners; +use crate::extension::coordinate::ordinates; +use crate::extension::flatten_row_offsets; use crate::extension::validate_geometry_operands; /// `envelope`: the axis-aligned bounding box of each geometry in a native geometry operand (a column @@ -62,120 +63,62 @@ impl GeoEnvelope { } } -/// The output extension type: a nullable native 2-D `geoarrow.box` ([`Rect`]). Nullable because a -/// null row has no box. Metadata is defaulted. -fn output_ext_dtype() -> VortexResult> { +/// The output dtype: a nullable native 2-D box ([`Rect`], `geoarrow.box`) column. Nullable +/// because rows without a box — null or empty geometries — are null. Metadata is defaulted. +fn output_box_dtype() -> VortexResult> { ExtDType::::try_new( GeoMetadata::default(), box_storage_dtype(Dimension::Xy, Nullability::Nullable), ) } -/// The per-row 2-D bounding box of `array` as a nullable `Rect` (`geoarrow.box`) column, computed -/// directly over the native coordinate storage — no decode to `geo_types`, no Arrow round-trip. +/// Compute each row's 2-D bounding box: the smallest rectangle covering all of the row's +/// coordinates. /// -/// Walks the nested `List` storage down to the leaf coordinate `Struct`, carrying which top-level -/// row owns each coordinate, then min/maxes x/y per row over the raw `f64` buffers. The ring/part -/// nesting is irrelevant to a bounding box, only which coordinates belong to a row matters. A null -/// row, or a valid row that owns no coordinate (an empty geometry), yields a null box. -fn envelopes(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let len = array.len(); - let valid = array.validity()?.execute_mask(len, ctx)?; - let ext = array - .dtype() - .as_extension_opt() - .ok_or_else(|| vortex_err!("geo: envelope operand is not a geometry extension type"))?; - let storage = array - .clone() - .execute::(ctx)? - .storage_array() - .clone(); - - // Per-row min/max accumulators; `seen[r]` marks a row that owns at least one coordinate. - let mut lo = vec![(f64::INFINITY, f64::INFINITY); len]; - let mut hi = vec![(f64::NEG_INFINITY, f64::NEG_INFINITY); len]; - let mut seen = vec![false; len]; - - if ext.is::() { - // The bounding box of a box is itself: read its min/max fields straight from storage. - let coords = storage.execute::(ctx)?; - let xmin = f64_field(&coords, "xmin", ctx)?; - let ymin = f64_field(&coords, "ymin", ctx)?; - let xmax = f64_field(&coords, "xmax", ctx)?; - let ymax = f64_field(&coords, "ymax", ctx)?; - let (xmin, ymin) = (xmin.as_slice::(), ymin.as_slice::()); - let (xmax, ymax) = (xmax.as_slice::(), ymax.as_slice::()); - for r in 0..len { - lo[r] = (xmin[r], ymin[r]); - hi[r] = (xmax[r], ymax[r]); - seen[r] = true; - } - } else { - // Walk any `List` nesting down to the leaf coordinate `Struct`, mapping each element to the - // top-level row that owns it. `ListView` offsets/sizes are arbitrary integer types and need - // not be contiguous, so map every element explicitly rather than composing offset arrays. - let u64_dtype = DType::Primitive(PType::U64, Nullability::NonNullable); - let mut owner: Vec = (0..len).collect(); - let mut node = storage; - while matches!(node.dtype(), DType::List(..)) { - let list = node.execute::(ctx)?.into_listview(); - let offsets = list - .offsets() - .clone() - .cast(u64_dtype.clone())? - .execute::(ctx)?; - let sizes = list - .sizes() - .clone() - .cast(u64_dtype.clone())? - .execute::(ctx)?; - let (offsets, sizes) = (offsets.as_slice::(), sizes.as_slice::()); - let elements = list.elements().clone(); - let mut child = vec![0usize; elements.len()]; - for (i, &row) in owner.iter().enumerate() { - let start = usize::try_from(offsets[i]) - .map_err(|_| vortex_err!("geo: list offset exceeds usize"))?; - let size = usize::try_from(sizes[i]) - .map_err(|_| vortex_err!("geo: list size exceeds usize"))?; - child[start..start + size].fill(row); - } - owner = child; - node = elements; - } - let coords = node.execute::(ctx)?; - let xs = f64_field(&coords, "x", ctx)?; - let ys = f64_field(&coords, "y", ctx)?; - for ((&r, &x), &y) in owner - .iter() - .zip(xs.as_slice::()) - .zip(ys.as_slice::()) - { - lo[r] = (lo[r].0.min(x), lo[r].1.min(y)); - hi[r] = (hi[r].0.max(x), hi[r].1.max(y)); - seen[r] = true; +/// The output is columnar: four `f64` corner columns in `geoarrow.box` field order +/// (`xmin, ymin, xmax, ymax`), plus the output validity. +/// +/// The output validity narrows the input mask `valid`: a row's box is null when the input row +/// is null, but also when a valid row owns no coordinates: an empty geometry (e.g. +/// `MULTIPOLYGON EMPTY`) has no extent, and a box cannot represent "empty". Null boxes hold +/// placeholder `0.0` corners. +fn row_boxes( + storage: ArrayRef, + valid: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult<(Vec, Validity)> { + let len = valid.len(); + let (row_offsets, coords) = flatten_row_offsets(storage, ctx)?; + let xs = ordinates(&coords, "x", ctx)?; + let ys = ordinates(&coords, "y", ctx)?; + + // The output's four columns (xmin, ymin, xmax, ymax), built by transposing each row's + // corners into them (a box column is stored as one flat column per corner field). + let mut corner_columns: [BufferMut; 4] = + std::array::from_fn(|_| BufferMut::with_capacity(len)); + let mut has_boxes = Vec::with_capacity(len); + for r in 0..len { + let (start, end) = (row_offsets[r], row_offsets[r + 1]); + // A valid input row with at least one coordinate has a box; a null input row or an + // empty geometry (`start == end`) does not. + let has_box = valid.value(r) && start < end; + let corners = if has_box { + box_corners(&xs[start..end], &ys[start..end]) + } else { + [0.0; 4] + }; + has_boxes.push(has_box); + for (column, corner) in corner_columns.iter_mut().zip(corners) { + column.push(corner); } } - - // Build the nullable `Rect` column straight from the accumulators: a row is present iff it owns - // a coordinate and its operand row is valid; absent rows are null (physical placeholder `0.0`). - let present: Vec = (0..len).map(|r| seen[r] && valid.value(r)).collect(); - let ordinate = |src: &[(f64, f64)], axis: fn((f64, f64)) -> f64| { - PrimitiveArray::from_iter((0..len).map(|r| if present[r] { axis(src[r]) } else { 0.0 })) - .into_array() - }; - let storage = StructArray::try_new( - FieldNames::from(["xmin", "ymin", "xmax", "ymax"]), - vec![ - ordinate(&lo, |c| c.0), - ordinate(&lo, |c| c.1), - ordinate(&hi, |c| c.0), - ordinate(&hi, |c| c.1), - ], - len, - Validity::from_iter(present.iter().copied()), - )? - .into_array(); - Ok(ExtensionArray::try_new(output_ext_dtype()?.erased(), storage)?.into_array()) + Ok(( + corner_columns + .into_iter() + .map(|column| column.freeze().into_array()) + .collect(), + Validity::from_iter(has_boxes), + )) } impl ScalarFnVTable for GeoEnvelope { @@ -209,9 +152,12 @@ impl ScalarFnVTable for GeoEnvelope { validate_geometry_operands(dtypes)?; // Always nullable: an empty geometry has no box, so nulls can appear even over a // non-nullable operand. - Ok(DType::Extension(output_ext_dtype()?.erased())) + Ok(DType::Extension(output_box_dtype()?.erased())) } + /// Compute each row's box directly over the native coordinate storage — no decode to + /// `geo_types`, no Arrow round-trip. A null row, or a valid row that owns no coordinate (an + /// empty geometry), yields a null box. fn execute( &self, _: &Self::Options, @@ -219,7 +165,42 @@ impl ScalarFnVTable for GeoEnvelope { ctx: &mut ExecutionCtx, ) -> VortexResult { let array = args.get(0)?; - envelopes(&array, ctx) + let len = array.len(); + let ext = array + .dtype() + .as_extension_opt() + .ok_or_else(|| vortex_err!("geo: envelope operand is not a geometry extension type"))?; + let storage = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + + let (corners, output_validity) = if ext.is::() { + // A box is its own envelope: project the 2-D corner fields straight out of storage + // (dropping any z/m bounds). A stored box cannot be empty, so the output validity + // is exactly the operand's, kept lazy. + let coords = storage.execute::(ctx)?; + let corners = box_field_names(Dimension::Xy) + .iter() + .map(|name| coords.unmasked_field_by_name(name).cloned()) + .collect::>>()?; + (corners, array.validity()?.into_nullable()) + } else { + let valid = array.validity()?.execute_mask(len, ctx)?; + row_boxes(storage, &valid, ctx)? + }; + + // Nullness lives at the box (struct) level: the corner fields stay non-nullable `f64`, + // holding `0.0` placeholders under null rows. + let storage = StructArray::try_new( + FieldNames::from(box_field_names(Dimension::Xy)), + corners, + len, + output_validity, + )? + .into_array(); + Ok(ExtensionArray::try_new(output_box_dtype()?.erased(), storage)?.into_array()) } fn validity(&self, _: &Self::Options, _: &Expression) -> VortexResult> { @@ -239,16 +220,24 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexResult; use super::GeoEnvelope; + use crate::extension::GeoMetadata; + use crate::extension::Rect; + use crate::extension::box_storage_dtype; + use crate::extension::coordinate::Dimension; use crate::test_harness::linestring_column; use crate::test_harness::multilinestring_column; use crate::test_harness::multipoint_column; @@ -311,6 +300,34 @@ mod tests { Ok(()) } + /// The `Rect` fast path projects the 2-D corners by name, so a 3-D box — whose `zmin`/`zmax` + /// fields are interleaved between them in storage — yields its XY extent as a 2-D box. + #[test] + fn xyz_rect_drops_z_bounds() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + + let ordinate = |value: f64| PrimitiveArray::from_iter([value]).into_array(); + let storage = StructArray::from_fields(&[ + ("xmin", ordinate(0.0)), + ("ymin", ordinate(1.0)), + ("zmin", ordinate(-9.0)), + ("xmax", ordinate(2.0)), + ("ymax", ordinate(3.0)), + ("zmax", ordinate(9.0)), + ])? + .into_array(); + let ext = ExtDType::::try_new( + GeoMetadata::default(), + box_storage_dtype(Dimension::Xyz, Nullability::NonNullable), + )?; + let rects = ExtensionArray::try_new(ext.erased(), storage)?.into_array(); + + let expected = nullable_rect_column(vec![Some((0.0, 1.0, 2.0, 3.0))])?; + assert_arrays_eq!(boxes(rects)?, expected, &mut ctx); + 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] @@ -333,6 +350,29 @@ mod tests { Ok(()) } + /// Intermediate list levels with more than one part per row keep coordinates attributed to + /// the right rows: row 0 owns two polygons (the first with two rings) whose extremes live in + /// the second polygon, so composed bounds diverging from loop positions shows up here. + #[test] + fn uneven_nesting_keeps_rows_aligned() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + + let multipolygons = multipolygon_column(vec![ + vec![ + vec![vec![(0.0, 0.0), (1.0, 1.0)], vec![(0.2, 0.2), (0.8, 0.8)]], + vec![vec![(5.0, -3.0), (6.0, 7.0)]], + ], + vec![vec![vec![(10.0, 10.0), (11.0, 12.0)]]], + ])?; + let expected = nullable_rect_column(vec![ + Some((0.0, -3.0, 6.0, 7.0)), + Some((10.0, 10.0, 11.0, 12.0)), + ])?; + assert_arrays_eq!(boxes(multipolygons)?, expected, &mut ctx); + Ok(()) + } + /// An empty geometry (here a zero-part multipolygon) has no extent and yields a null box; other /// rows keep their boxes. #[test] @@ -349,6 +389,40 @@ mod tests { Ok(()) } + /// A geometry empty only at an inner level — here a polygon whose single ring has zero + /// vertices, in the first row — has no box, exactly like one empty at the outer level. + #[test] + fn inner_empty_ring_yields_null_box() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + + let polygons = polygon_column(vec![ + vec![vec![]], + vec![vec![(1.0, 2.0), (3.0, 4.0), (1.0, 4.0)]], + ])?; + let expected = nullable_rect_column(vec![None, Some((1.0, 2.0, 3.0, 4.0))])?; + assert_arrays_eq!(boxes(polygons)?, expected, &mut ctx); + Ok(()) + } + + /// A sliced operand keeps per-row boxes aligned: coordinates outside the slice window (still + /// present in the sliced list's element buffer) must not leak into any row's box. + #[test] + fn sliced_operand_ignores_out_of_slice_coordinates() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + + let multipoints = multipoint_column(vec![ + vec![(-100.0, -100.0), (100.0, 100.0)], + vec![(1.0, 2.0), (3.0, 4.0)], + vec![(5.0, 6.0)], + ])?; + let expected = + nullable_rect_column(vec![Some((1.0, 2.0, 3.0, 4.0)), Some((5.0, 6.0, 5.0, 6.0))])?; + assert_arrays_eq!(boxes(multipoints.slice(1..3)?)?, expected, &mut ctx); + Ok(()) + } + /// A null geometry row yields a null box, just like an empty geometry; valid rows keep theirs. #[test] fn null_row_yields_null_box() -> VortexResult<()> { From 080cd81c1880ac97e16760c294973b2efd1fd017 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 21 Jul 2026 15:24:54 -0400 Subject: [PATCH 3/8] refactor(vortex-geo): branch-free envelope reduction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reduce every row unconditionally instead of branching on validity per row: a null or empty row folds its (usually empty) coordinate range into meaningless corners that sit under a null in the output and are never observed. The output validity is computed separately as a bulk mask combine — the operand's null mask AND a non-empty mask built with BitBuffer::collect_bool over the row offsets. Benchmarked against the branchy and hybrid alternatives across geometry shapes and null densities (none / 10% periodic / 50% pseudo-random): the branch-free form is best or tied on list-backed geometries and makes runtime independent of the null distribution. Signed-off-by: Nemo Yu --- vortex-geo/src/extension/mod.rs | 4 +- vortex-geo/src/scalar_fn/envelope.rs | 70 ++++++++++++---------------- 2 files changed, 34 insertions(+), 40 deletions(-) diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index d7c3ee912cd..bcdc0901c70 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -142,8 +142,10 @@ pub(crate) fn flatten_row_offsets( storage: ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult<(Vec, 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 = (0..=storage.len()).collect(); + let mut row_offsets: Vec = (0..=len).collect(); let mut level = storage; while matches!(level.dtype(), DType::List(..)) { let list = list_from_list_view(level.execute::(ctx)?.into_listview(), ctx)?; diff --git a/vortex-geo/src/scalar_fn/envelope.rs b/vortex-geo/src/scalar_fn/envelope.rs index f83c1527745..2213fe68bb1 100644 --- a/vortex-geo/src/scalar_fn/envelope.rs +++ b/vortex-geo/src/scalar_fn/envelope.rs @@ -27,6 +27,7 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -74,50 +75,42 @@ fn output_box_dtype() -> VortexResult> { /// Compute each row's 2-D bounding box: the smallest rectangle covering all of the row's /// coordinates. -/// -/// The output is columnar: four `f64` corner columns in `geoarrow.box` field order -/// (`xmin, ymin, xmax, ymax`), plus the output validity. -/// -/// The output validity narrows the input mask `valid`: a row's box is null when the input row -/// is null, but also when a valid row owns no coordinates: an empty geometry (e.g. -/// `MULTIPOLYGON EMPTY`) has no extent, and a box cannot represent "empty". Null boxes hold -/// placeholder `0.0` corners. -fn row_boxes( - storage: ArrayRef, - valid: &Mask, - ctx: &mut ExecutionCtx, -) -> VortexResult<(Vec, Validity)> { - let len = valid.len(); +fn row_boxes(storage: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<(Vec, Validity)> { + let len = storage.len(); + let valid = storage.validity()?.execute_mask(len, ctx)?; let (row_offsets, coords) = flatten_row_offsets(storage, ctx)?; + + // A row has a box iff it is valid and owns at least one coordinate (an empty geometry has + // no box): the envelope-specific narrowing of the operand's mask, combined word-at-a-time. + let non_empty = Mask::from(BitBuffer::collect_bool(len, |r| { + row_offsets[r] < row_offsets[r + 1] + })); let xs = ordinates(&coords, "x", ctx)?; let ys = ordinates(&coords, "y", ctx)?; - // The output's four columns (xmin, ymin, xmax, ymax), built by transposing each row's - // corners into them (a box column is stored as one flat column per corner field). - let mut corner_columns: [BufferMut; 4] = - std::array::from_fn(|_| BufferMut::with_capacity(len)); - let mut has_boxes = Vec::with_capacity(len); + // The output's four corner columns. + 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]); - // A valid input row with at least one coordinate has a box; a null input row or an - // empty geometry (`start == end`) does not. - let has_box = valid.value(r) && start < end; - let corners = if has_box { - box_corners(&xs[start..end], &ys[start..end]) - } else { - [0.0; 4] - }; - has_boxes.push(has_box); - for (column, corner) in corner_columns.iter_mut().zip(corners) { - column.push(corner); - } + 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); } + Ok(( - corner_columns - .into_iter() - .map(|column| column.freeze().into_array()) - .collect(), - Validity::from_iter(has_boxes), + vec![ + xmins.freeze().into_array(), + ymins.freeze().into_array(), + xmaxs.freeze().into_array(), + ymaxs.freeze().into_array(), + ], + Validity::from_mask(&valid & &non_empty, Nullability::Nullable), )) } @@ -187,12 +180,11 @@ impl ScalarFnVTable for GeoEnvelope { .collect::>>()?; (corners, array.validity()?.into_nullable()) } else { - let valid = array.validity()?.execute_mask(len, ctx)?; - row_boxes(storage, &valid, ctx)? + row_boxes(storage, ctx)? }; // Nullness lives at the box (struct) level: the corner fields stay non-nullable `f64`, - // holding `0.0` placeholders under null rows. + // holding unspecified values under null rows. let storage = StructArray::try_new( FieldNames::from(box_field_names(Dimension::Xy)), corners, From dfbd44363e2d035e2f57e2a8141cca01cc7bb776 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 21 Jul 2026 16:12:20 -0400 Subject: [PATCH 4/8] perf(vortex-geo): write envelope corners into pre-sized buffers BufferMut::push stays an out-of-line call (its grow path defeats inlining), so the envelope row loop paid four calls plus capacity checks per row. Pre-size the four corner buffers and write by index, and iterate offset pairs to drop the per-row bounds check on row_offsets[r + 1]. Signed-off-by: Nemo Yu --- vortex-geo/src/scalar_fn/envelope.rs | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/vortex-geo/src/scalar_fn/envelope.rs b/vortex-geo/src/scalar_fn/envelope.rs index 2213fe68bb1..71f958b8488 100644 --- a/vortex-geo/src/scalar_fn/envelope.rs +++ b/vortex-geo/src/scalar_fn/envelope.rs @@ -89,18 +89,17 @@ fn row_boxes(storage: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<(Vec Date: Tue, 21 Jul 2026 16:13:52 -0400 Subject: [PATCH 5/8] perf(vortex-geo): zero-copy envelope fast path for point columns Point storage is the coordinate struct itself, so each row's box is degenerate: project the x/y ordinate arrays straight into the corner columns (both corners alias the same buffers) and keep the operand's validity lazy, skipping the flatten and reduction entirely. Signed-off-by: Nemo Yu --- vortex-geo/src/scalar_fn/envelope.rs | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/vortex-geo/src/scalar_fn/envelope.rs b/vortex-geo/src/scalar_fn/envelope.rs index 71f958b8488..9a4ff03cab3 100644 --- a/vortex-geo/src/scalar_fn/envelope.rs +++ b/vortex-geo/src/scalar_fn/envelope.rs @@ -178,6 +178,18 @@ impl ScalarFnVTable for GeoEnvelope { .map(|name| coords.unmasked_field_by_name(name).cloned()) .collect::>>()?; (corners, array.validity()?.into_nullable()) + } else if !matches!(storage.dtype(), DType::List(..)) { + // Point storage is the coordinate `Struct` itself: every row owns exactly one + // coordinate, so its box is degenerate and the corner columns are the ordinate + // arrays, zero-copy. No row can be empty, so the output validity is exactly the + // operand's, kept lazy. + let coords = storage.execute::(ctx)?; + let x = coords.unmasked_field_by_name("x")?.clone(); + let y = coords.unmasked_field_by_name("y")?.clone(); + ( + vec![x.clone(), y.clone(), x, y], + array.validity()?.into_nullable(), + ) } else { row_boxes(storage, ctx)? }; @@ -414,6 +426,19 @@ mod tests { Ok(()) } + /// The zero-copy point fast path respects slicing: corners come from the slice window only. + #[test] + fn sliced_point_column_keeps_rows_aligned() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + + let points = point_column(vec![9.0, 1.0, 3.0], vec![8.0, 2.0, 4.0])?; + let expected = + nullable_rect_column(vec![Some((1.0, 2.0, 1.0, 2.0)), Some((3.0, 4.0, 3.0, 4.0))])?; + assert_arrays_eq!(boxes(points.slice(1..3)?)?, expected, &mut ctx); + Ok(()) + } + /// A null geometry row yields a null box, just like an empty geometry; valid rows keep theirs. #[test] fn null_row_yields_null_box() -> VortexResult<()> { From a181a2bbfb32b0c1c981d11e37bc153f85297de2 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 21 Jul 2026 16:30:42 -0400 Subject: [PATCH 6/8] bench(vortex-geo): envelope corner write-path microbenchmark push (BufferMut::with_capacity plus four out-of-line push calls per row, the shape the pre-sized loop replaced) vs the shipped indexed stores over pre-sized buffers: 2x at one vertex per row, 1.6x at four, a wash by 64 vertices, and no size regime where indexed loses (Apple M5 Max). Signed-off-by: Nemo Yu --- Cargo.lock | 1 + vortex-geo/Cargo.toml | 5 ++ vortex-geo/benches/envelope_write.rs | 94 ++++++++++++++++++++++++++++ 3 files changed, 100 insertions(+) create mode 100644 vortex-geo/benches/envelope_write.rs diff --git a/Cargo.lock b/Cargo.lock index 0747e266768..16654b0fed3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10265,6 +10265,7 @@ version = "0.1.0" dependencies = [ "arrow-array 58.3.0", "arrow-schema 58.3.0", + "codspeed-divan-compat", "geo", "geo-traits", "geo-types", diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index 351161a6ec4..b1fa5c2eab7 100644 --- a/vortex-geo/Cargo.toml +++ b/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 diff --git a/vortex-geo/benches/envelope_write.rs b/vortex-geo/benches/envelope_write.rs new file mode 100644 index 00000000000..ee2860f90b2 --- /dev/null +++ b/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, Vec, Vec) { + 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; 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; 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(), + ] + }); +} From 6686db51a8b64b1a191e22de931e44fed440fcdf Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 21 Jul 2026 17:12:23 -0400 Subject: [PATCH 7/8] refactor(vortex-geo): review follow-ups for envelope - Execute list levels straight to ListViewArray instead of Canonical::into_listview (same canonicalization, no panicking downcast), in both flatten walks. - dtype().is_list() over matches!(.., DType::List(..)) (review nit), swept across envelope.rs and extension/mod.rs. - Document why row_boxes combines two masks word-at-a-time rather than folding validity into the collect_bool closure: both the per-index and Mask::iter forms bench 6-33% slower end-to-end. Signed-off-by: Nemo Yu --- vortex-geo/src/extension/mod.rs | 14 +++++--------- vortex-geo/src/scalar_fn/envelope.rs | 5 +++-- 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index bcdc0901c70..99150833ab5 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -40,11 +40,11 @@ pub use point::*; pub use polygon::*; pub use rect::*; use vortex_array::ArrayRef; -use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::ListViewArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::list::ListArrayExt; @@ -110,12 +110,8 @@ pub(crate) fn flatten_coordinates( .execute::(ctx)? .storage_array() .clone(); - while matches!(node.dtype(), DType::List(..)) { - node = node - .execute::(ctx)? - .into_listview() - .elements() - .clone(); + while node.dtype().is_list() { + node = node.execute::(ctx)?.elements().clone(); } node.execute::(ctx) } @@ -147,8 +143,8 @@ pub(crate) fn flatten_row_offsets( // At the outermost level, row `r` starts at element `r`; the extra entry caps the last row. let mut row_offsets: Vec = (0..=len).collect(); let mut level = storage; - while matches!(level.dtype(), DType::List(..)) { - let list = list_from_list_view(level.execute::(ctx)?.into_listview(), ctx)?; + while level.dtype().is_list() { + let list = list_from_list_view(level.execute::(ctx)?, ctx)?; let offsets = list .offsets() .clone() diff --git a/vortex-geo/src/scalar_fn/envelope.rs b/vortex-geo/src/scalar_fn/envelope.rs index 9a4ff03cab3..d04a125147f 100644 --- a/vortex-geo/src/scalar_fn/envelope.rs +++ b/vortex-geo/src/scalar_fn/envelope.rs @@ -81,7 +81,8 @@ fn row_boxes(storage: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<(Vec>>()?; (corners, array.validity()?.into_nullable()) - } else if !matches!(storage.dtype(), DType::List(..)) { + } else if !storage.dtype().is_list() { // Point storage is the coordinate `Struct` itself: every row owns exactly one // coordinate, so its box is degenerate and the corner columns are the ordinate // arrays, zero-copy. No row can be empty, so the output validity is exactly the From 3879f95c09694db95f125c6494111962bf051989 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 21 Jul 2026 17:12:23 -0400 Subject: [PATCH 8/8] bench(vortex-geo): benchmark envelope end-to-end at the scalar-fn level Replace the write-path shape comparison (its numbers are preserved in that commit's message) with an end-to-end benchmark of vortex.geo.envelope: the lazy ScalarFnArray executed to Canonical through a geo session, over harness-built columns. Cases cross nesting depth (Point, MultiPoint 32 pts/row, MultiPolygon 2 polygons x 2 rings x 8 vertices) with validity (non-nullable, ~10% periodic nulls, ~50% pseudo-random nulls). A _test-harness feature plus a self dev-dependency exposes test_harness to benches, the vortex-array pattern. Medians on an Apple M5 Max, 131072 rows per case: point ~0.8us at any validity (the zero-copy fast path makes execute O(1), so this is dispatch overhead), multipoint 1.39ms, multipolygon 2.17ms non-nullable / 2.11ms periodic nulls / 1.37ms random nulls (null rows are empty in storage and cost nothing). Row count scales linearly: throughput per row is identical at 1M rows. Signed-off-by: Nemo Yu --- Cargo.lock | 1 + vortex-geo/Cargo.toml | 7 +- vortex-geo/benches/envelope.rs | 167 +++++++++++++++++++++++++++ vortex-geo/benches/envelope_write.rs | 94 --------------- vortex-geo/src/lib.rs | 4 +- vortex-geo/src/test_harness.rs | 30 +++-- 6 files changed, 189 insertions(+), 114 deletions(-) create mode 100644 vortex-geo/benches/envelope.rs delete mode 100644 vortex-geo/benches/envelope_write.rs diff --git a/Cargo.lock b/Cargo.lock index 16654b0fed3..ab869e730b9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10277,6 +10277,7 @@ dependencies = [ "vortex-arrow", "vortex-buffer", "vortex-error", + "vortex-geo", "vortex-layout", "vortex-mask", "vortex-session", diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index b1fa5c2eab7..eecd2ed7461 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -30,14 +30,19 @@ 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_write" +name = "envelope" harness = false [lints] diff --git a/vortex-geo/benches/envelope.rs b/vortex-geo/benches/envelope.rs new file mode 100644 index 00000000000..6ae5d6404c2 --- /dev/null +++ b/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 = 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::(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)); +} diff --git a/vortex-geo/benches/envelope_write.rs b/vortex-geo/benches/envelope_write.rs deleted file mode 100644 index ee2860f90b2..00000000000 --- a/vortex-geo/benches/envelope_write.rs +++ /dev/null @@ -1,94 +0,0 @@ -// 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, Vec, Vec) { - 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; 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; 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(), - ] - }); -} diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index e823c6301a6..b72df92a3f5 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -30,8 +30,8 @@ pub mod aggregate_fn; pub mod extension; pub mod prune; pub mod scalar_fn; -#[cfg(test)] -mod test_harness; +#[cfg(any(test, feature = "_test-harness"))] +pub mod test_harness; #[cfg(test)] mod tests; diff --git a/vortex-geo/src/test_harness.rs b/vortex-geo/src/test_harness.rs index 593462eb662..7bec0f90ee9 100644 --- a/vortex-geo/src/test_harness.rs +++ b/vortex-geo/src/test_harness.rs @@ -39,7 +39,7 @@ use crate::extension::multipolygon_storage_dtype; use crate::extension::polygon_storage_dtype; /// A fresh session with the geospatial types, functions, and pruning rules registered. -pub(crate) fn geo_session() -> VortexSession { +pub fn geo_session() -> VortexSession { let session = vortex_array::array_session(); crate::initialize(&session); session @@ -106,7 +106,7 @@ fn geo_column + Default>( } /// A `Point` column over the given x/y coordinates, stored as `Struct`. -pub(crate) fn point_column(xs: Vec, ys: Vec) -> VortexResult { +pub fn point_column(xs: Vec, ys: Vec) -> VortexResult { let storage = xy_struct(xs, ys)?; let storage_dtype = storage.dtype().clone(); geo_column::(storage, storage_dtype) @@ -114,7 +114,7 @@ pub(crate) fn point_column(xs: Vec, ys: Vec) -> VortexResult /// A nullable `Point` column: `None` rows are null. Null rows carry placeholder coordinates in /// storage that the geo kernels must never decode (they filter nulls before decoding). -pub(crate) fn nullable_point_column(points: Vec>) -> VortexResult { +pub fn nullable_point_column(points: Vec>) -> VortexResult { let len = points.len(); let valid = points.iter().map(Option::is_some); let xs = points.iter().map(|p| p.map_or(0.0, |(x, _)| x)); @@ -134,7 +134,7 @@ pub(crate) fn nullable_point_column(points: Vec>) -> VortexRe } /// A `LineString` column: each line a list of `(x, y)` vertices, stored as `List>`. -pub(crate) fn linestring_column(lines: Vec>) -> VortexResult { +pub fn linestring_column(lines: Vec>) -> VortexResult { geo_column::( vertex_lists(&lines)?, linestring_storage_dtype(Dimension::Xy, Nullability::NonNullable), @@ -142,7 +142,7 @@ pub(crate) fn linestring_column(lines: Vec>) -> VortexResult>`. -pub(crate) fn multipoint_column(points: Vec>) -> VortexResult { +pub fn multipoint_column(points: Vec>) -> VortexResult { geo_column::( vertex_lists(&points)?, multipoint_storage_dtype(Dimension::Xy, Nullability::NonNullable), @@ -151,7 +151,7 @@ pub(crate) fn multipoint_column(points: Vec>) -> VortexResult>>`. -pub(crate) fn polygon_column(polygons: Vec>>) -> VortexResult { +pub fn polygon_column(polygons: Vec>>) -> VortexResult { geo_column::( vertex_list_lists(&polygons)?, polygon_storage_dtype(Dimension::Xy, Nullability::NonNullable), @@ -159,9 +159,7 @@ pub(crate) fn polygon_column(polygons: Vec>>) -> VortexResul } /// A `MultiLineString` column: each row a list of lines, stored as `List>>`. -pub(crate) fn multilinestring_column( - multilines: Vec>>, -) -> VortexResult { +pub fn multilinestring_column(multilines: Vec>>) -> VortexResult { geo_column::( vertex_list_lists(&multilines)?, multilinestring_storage_dtype(Dimension::Xy, Nullability::NonNullable), @@ -169,10 +167,10 @@ pub(crate) fn multilinestring_column( } /// One multipolygon: polygons → rings → `(x, y)` vertices. -pub(crate) type MultiPolygonRings = Vec>>; +pub type MultiPolygonRings = Vec>>; /// A `MultiPolygon` column, stored as `List>>>`. -pub(crate) fn multipolygon_column(multipolygons: Vec) -> VortexResult { +pub fn multipolygon_column(multipolygons: Vec) -> VortexResult { let polygons: Vec>> = multipolygons.iter().flatten().cloned().collect(); geo_column::( nest(&multipolygons, vertex_list_lists(&polygons)?)?, @@ -181,7 +179,7 @@ pub(crate) fn multipolygon_column(multipolygons: Vec) -> Vort } /// A nullable `MultiPolygon` column: `None` rows are null (an empty-list placeholder in storage). -pub(crate) fn nullable_multipolygon_column( +pub fn nullable_multipolygon_column( multipolygons: Vec>, ) -> VortexResult { let rows: Vec = multipolygons @@ -209,7 +207,7 @@ pub(crate) fn nullable_multipolygon_column( /// A 2D `Rect` (`geoarrow.box`) column over `(xmin, ymin, xmax, ymax)` boxes, stored as /// `Struct`. -pub(crate) fn rect_column(boxes: Vec<(f64, f64, f64, f64)>) -> VortexResult { +pub fn rect_column(boxes: Vec<(f64, f64, f64, f64)>) -> VortexResult { let field = |select: fn(&(f64, f64, f64, f64)) -> f64| { PrimitiveArray::from_iter(boxes.iter().map(select)).into_array() }; @@ -229,9 +227,7 @@ pub(crate) fn rect_column(boxes: Vec<(f64, f64, f64, f64)>) -> VortexResult>, -) -> VortexResult { +pub fn nullable_rect_column(boxes: Vec>) -> VortexResult { let len = boxes.len(); let field = |select: fn(&(f64, f64, f64, f64)) -> f64| { PrimitiveArray::from_iter(boxes.iter().map(|b| b.as_ref().map_or(0.0, select))).into_array() @@ -257,7 +253,7 @@ pub(crate) fn nullable_rect_column( /// Decode a [`Coordinate`] from an extension-typed point scalar (unwrapped to its coordinate /// storage) or a bare coordinate `Struct` scalar — used to read back a single point in assertions. -pub(crate) fn coordinate_from_scalar(scalar: &Scalar) -> VortexResult { +pub fn coordinate_from_scalar(scalar: &Scalar) -> VortexResult { match scalar.as_extension_opt() { Some(ext_scalar) => coordinate_from_struct(&ext_scalar.to_storage_scalar()), None => coordinate_from_struct(scalar),