diff --git a/Cargo.lock b/Cargo.lock index 0747e266768..ab869e730b9 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", @@ -10276,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 351161a6ec4..eecd2ed7461 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -30,10 +30,20 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } wkb = { workspace = true } +[features] +# Exposes the `test_harness` module to benches; not part of the public API. +_test-harness = [] + [dev-dependencies] +divan = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-geo = { path = ".", features = ["_test-harness"] } vortex-layout = { workspace = true } +[[bench]] +name = "envelope" +harness = false + [lints] workspace = true 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/src/aggregate_fn/aabb.rs b/vortex-geo/src/aggregate_fn/aabb.rs index 81f4450a674..8ac5c9e2452 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,8 @@ use crate::extension::GeoMetadata; use crate::extension::Rect; use crate::extension::box_storage_dtype; use crate::extension::coordinate::Dimension; +use crate::extension::coordinate::box_corners; +use crate::extension::coordinate::ordinates; use crate::extension::flatten_coordinates; use crate::extension::is_native_geometry; @@ -78,18 +78,10 @@ fn aabb_storage_dtype() -> DType { /// The AABB of the raw `x`/`y` slices, or `None` when empty. fn aabb_of(xs: &[f64], ys: &[f64]) -> Option> { - if xs.is_empty() { - return None; - } - let min_max = |vals: &[f64]| { - vals.iter() - .fold((f64::INFINITY, f64::NEG_INFINITY), |(lo, hi), &v| { - (lo.min(v), hi.max(v)) - }) - }; - let (xmin, xmax) = min_max(xs); - let (ymin, ymax) = min_max(ys); - Some(GeoRect::new((xmin, ymin), (xmax, ymax))) + (!xs.is_empty()).then(|| { + let [xmin, ymin, xmax, ymax] = box_corners(xs, ys); + GeoRect::new((xmin, ymin), (xmax, ymax)) + }) } /// Read an AABB stat scalar (a nullable native `geoarrow.box`) into a [`GeoRect`], or `None` when @@ -217,15 +209,9 @@ impl AggregateFnVTable for GeometryAabb { // `unmasked_field_by_name` reads are therefore safe. Min/max the raw x/y buffers directly: // cheap, and avoids `to_geometry`'s panic on empty points (which decoding would hit). let coords = flatten_coordinates(&array, ctx)?; - let xs = coords - .unmasked_field_by_name("x")? - .clone() - .execute::(ctx)?; - let ys = coords - .unmasked_field_by_name("y")? - .clone() - .execute::(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 b6e324cd457..dc3537cfb30 100644 --- a/vortex-geo/src/extension/coordinate.rs +++ b/vortex-geo/src/extension/coordinate.rs @@ -15,12 +15,16 @@ use std::fmt::Display; use std::fmt::Formatter; use geoarrow::datatypes::Dimension as GeoArrowDimension; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; use vortex_array::scalar::Scalar; +use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -189,6 +193,33 @@ pub(crate) fn coordinate_from_struct(scalar: &Scalar) -> VortexResult VortexResult> { + coords + .unmasked_field_by_name(name)? + .clone() + .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)] mod tests { use rstest::rstest; diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 1a3db5a7b3e..99150833ab5 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -40,19 +40,25 @@ 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; 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; @@ -104,16 +110,55 @@ 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) } +/// 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)> { + 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..=len).collect(); + let mut level = storage; + while level.dtype().is_list() { + let list = list_from_list_view(level.execute::(ctx)?, 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/lib.rs b/vortex-geo/src/lib.rs index 3a0c52f1eb9..b72df92a3f5 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -23,14 +23,15 @@ 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; 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; @@ -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..d04a125147f --- /dev/null +++ b/vortex-geo/src/scalar_fn/envelope.rs @@ -0,0 +1,512 @@ +// 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::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +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::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_buffer::BitBuffer; +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::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 +/// 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 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), + ) +} + +/// Compute each row's 2-D bounding box: the smallest rectangle covering all of the row's +/// coordinates. +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). Two masks combined word-at-a-time: folding `valid` into the closure instead — + // per-index or via `Mask::iter` — benches 6-33% slower end-to-end. + 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 corner columns. + 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; + } + + Ok(( + 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), + )) +} + +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_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, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let array = args.get(0)?; + 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 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 + // 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)? + }; + + // Nullness lives at the box (struct) level: the corner fields stay non-nullable `f64`, + // holding unspecified values 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> { + // 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::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; + 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(()) + } + + /// 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] + 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(()) + } + + /// 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] + 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 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(()) + } + + /// 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<()> { + 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..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)?)?, @@ -180,9 +178,36 @@ pub(crate) fn multipolygon_column(multipolygons: Vec) -> Vort ) } +/// A nullable `MultiPolygon` column: `None` rows are null (an empty-list placeholder in storage). +pub 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 { +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() }; @@ -199,9 +224,36 @@ 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 { +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),