From 72aa947b4164161d69a77c2ec6818a50f8ef540b Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 21 Jul 2026 10:13:40 -0400 Subject: [PATCH 1/3] 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/3] 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/3] 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,