diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 8d2b8cfe509..951b06ecf0d 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -65,5 +65,18 @@ harness = false [[bench]] name = "area" harness = false + +[[bench]] +name = "collect" +harness = false + +[[bench]] +name = "convex_hull" +harness = false + +[[bench]] +name = "intersection" +harness = false + [lints] workspace = true diff --git a/vortex-spatial/benches/collect.rs b/vortex-spatial/benches/collect.rs new file mode 100644 index 00000000000..e07ed00fde3 --- /dev/null +++ b/vortex-spatial/benches/collect.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for native `ST_Collect` over homogeneous geometry lists. +//! +//! The cases cover each strict overload and the inner-null compaction path. They execute the +//! result to its canonical representation so the full multi-geometry construction is measured. +//! +//! Run with `cargo bench -p vortex-spatial --bench collect`. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_session::VortexSession; +use vortex_spatial::scalar_fn::collect::SpatialCollect; +use vortex_spatial::scalar_fn::envelope::SpatialEnvelope; +use vortex_spatial::test_harness::linestring_column; +use vortex_spatial::test_harness::nullable_point_column; +use vortex_spatial::test_harness::point_column; +use vortex_spatial::test_harness::polygon_column; +use vortex_spatial::test_harness::spatial_session; + +// Scalar function execution allocates its output inside the timed region, so use the vendored +// allocator instead of measuring glibc differences between CodSpeed runner images. +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(spatial_session); + +const ROWS: usize = 512; + +fn main() { + divan::main(); +} + +fn geometry_lists(elements: ArrayRef, elements_per_row: usize) -> ArrayRef { + let offsets = PrimitiveArray::from_iter( + (0..=ROWS).map(|row| u64::try_from(row * elements_per_row).unwrap()), + ) + .into_array(); + ListArray::try_new(elements, offsets, Validity::NonNullable) + .unwrap() + .into_array() +} + +fn point_lists(nullable: bool) -> ArrayRef { + const POINTS_PER_ROW: usize = 8; + let len = ROWS * POINTS_PER_ROW; + let points = if nullable { + nullable_point_column( + (0..len) + .map(|i| (!i.is_multiple_of(8)).then_some((i as f64, (i + 1) as f64))) + .collect(), + ) + .unwrap() + } else { + point_column( + (0..len).map(|i| i as f64).collect(), + (0..len).map(|i| (i + 1) as f64).collect(), + ) + .unwrap() + }; + geometry_lists(points, POINTS_PER_ROW) +} + +fn linestring_lists() -> ArrayRef { + const LINES_PER_ROW: usize = 4; + let lines = linestring_column( + (0..ROWS * LINES_PER_ROW) + .map(|line| { + (0..8) + .map(|vertex| { + let value = (line * 8 + vertex) as f64; + (value, value + 1.0) + }) + .collect() + }) + .collect(), + ) + .unwrap(); + geometry_lists(lines, LINES_PER_ROW) +} + +fn polygon_lists() -> ArrayRef { + const POLYGONS_PER_ROW: usize = 2; + let polygons = polygon_column( + (0..ROWS * POLYGONS_PER_ROW) + .map(|polygon| { + let x = polygon as f64; + vec![vec![ + (x, 0.0), + (x + 1.0, 0.0), + (x + 1.0, 1.0), + (x, 1.0), + (x, 0.0), + ]] + }) + .collect(), + ) + .unwrap(); + geometry_lists(polygons, POLYGONS_PER_ROW) +} + +fn collect(input: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + SpatialCollect::try_new_array(input.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +fn bench_collect(bencher: Bencher, input: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| collect(&input, &mut ctx)); +} + +#[divan::bench] +fn points(bencher: Bencher) { + bench_collect(bencher, point_lists(false)); +} + +#[divan::bench] +fn linestrings(bencher: Bencher) { + bench_collect(bencher, linestring_lists()); +} + +#[divan::bench] +fn polygons(bencher: Bencher) { + bench_collect(bencher, polygon_lists()); +} + +#[divan::bench] +fn nullable_points(bencher: Bencher) { + bench_collect(bencher, point_lists(true)); +} + +/// Collect feeding a consumer that converts the result to a `ListArray`. +/// +/// The cases above stop at [`Canonical`], whose list form is a `ListViewArray`, so they cannot +/// observe whether collect's output still reports itself as zero-copy to a list. `ST_Envelope` +/// reaches that path through `flatten_row_offsets`, and re-gathers the whole payload when the +/// flag is missing. +fn envelope_of_collect(input: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + let collected = SpatialCollect::try_new_array(input.clone()) + .unwrap() + .into_array(); + SpatialEnvelope::try_new_array(collected) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +#[divan::bench] +fn envelope_of_collected_points(bencher: Bencher) { + let input = point_lists(false); + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| envelope_of_collect(&input, &mut ctx)); +} diff --git a/vortex-spatial/benches/convex_hull.rs b/vortex-spatial/benches/convex_hull.rs new file mode 100644 index 00000000000..a3e345a3eb3 --- /dev/null +++ b/vortex-spatial/benches/convex_hull.rs @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for native `ST_ConvexHull` over `MultiPoint` rows. +//! +//! The cases separate ordinary small hulls, larger point sets, and strict null propagation. They +//! execute the result to its canonical polygon representation. +//! +//! Run with `cargo bench -p vortex-spatial --bench convex_hull`. + +#![expect(clippy::unwrap_used)] + +use std::f64::consts::TAU; +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::validity::Validity; +use vortex_session::VortexSession; +use vortex_spatial::scalar_fn::convex_hull::SpatialConvexHull; +use vortex_spatial::test_harness::multipoint_column; +use vortex_spatial::test_harness::spatial_session; + +// Scalar function execution allocates its output inside the timed region, so use the vendored +// allocator instead of measuring glibc differences between CodSpeed runner images. +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(spatial_session); + +const ROWS: usize = 512; + +fn main() { + divan::main(); +} + +fn multipoints(points_per_row: usize) -> ArrayRef { + multipoint_column( + (0..ROWS) + .map(|row| { + (0..points_per_row) + .map(|point| { + let angle = TAU * point as f64 / points_per_row as f64; + let radius = 10.0 + ((row + point) % 7) as f64; + (radius * angle.cos(), radius * angle.sin()) + }) + .collect() + }) + .collect(), + ) + .unwrap() +} + +fn hulls(input: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + SpatialConvexHull::try_new_array(input.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +fn bench_hulls(bencher: Bencher, input: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| hulls(&input, &mut ctx)); +} + +#[divan::bench] +fn eight_points(bencher: Bencher) { + bench_hulls(bencher, multipoints(8)); +} + +#[divan::bench] +fn sixty_four_points(bencher: Bencher) { + bench_hulls(bencher, multipoints(64)); +} + +#[divan::bench] +fn nullable_eight_points(bencher: Bencher) { + let input = MaskedArray::try_new( + multipoints(8), + Validity::from_iter((0..ROWS).map(|row| !row.is_multiple_of(8))), + ) + .unwrap() + .into_array(); + bench_hulls(bencher, input); +} diff --git a/vortex-spatial/benches/intersection.rs b/vortex-spatial/benches/intersection.rs new file mode 100644 index 00000000000..d0ac0f5dbde --- /dev/null +++ b/vortex-spatial/benches/intersection.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for native `ST_Intersection` over polygon pairs. +//! +//! The cases cover simple building-like rectangles, more detailed boundaries, and strict null +//! propagation. Inputs overlap because SpatialBench Q9 prefilters pairs with `ST_Intersects`. +//! +//! Run with `cargo bench -p vortex-spatial --bench intersection`. + +#![expect(clippy::unwrap_used)] + +use std::f64::consts::TAU; +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::validity::Validity; +use vortex_session::VortexSession; +use vortex_spatial::scalar_fn::intersection::SpatialIntersection; +use vortex_spatial::test_harness::polygon_column; +use vortex_spatial::test_harness::spatial_session; + +// Scalar function execution allocates its output inside the timed region, so use the vendored +// allocator instead of measuring glibc differences between CodSpeed runner images. +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(spatial_session); + +const ROWS: usize = 512; + +fn main() { + divan::main(); +} + +fn regular_polygon(cx: f64, cy: f64, radius: f64, vertices: usize) -> Vec<(f64, f64)> { + (0..=vertices) + .map(|vertex| { + let angle = TAU * (vertex % vertices) as f64 / vertices as f64; + (cx + radius * angle.cos(), cy + radius * angle.sin()) + }) + .collect() +} + +fn polygon_pairs(vertices: usize) -> (ArrayRef, ArrayRef) { + let left = polygon_column( + (0..ROWS) + .map(|row| vec![regular_polygon(row as f64, 0.0, 1.0, vertices)]) + .collect(), + ) + .unwrap(); + let right = polygon_column( + (0..ROWS) + .map(|row| vec![regular_polygon(row as f64 + 0.5, 0.0, 1.0, vertices)]) + .collect(), + ) + .unwrap(); + (left, right) +} + +fn intersections(left: &ArrayRef, right: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + SpatialIntersection::try_new_array(left.clone(), right.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +fn bench_intersections(bencher: Bencher, left: ArrayRef, right: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| intersections(&left, &right, &mut ctx)); +} + +#[divan::bench] +fn rectangles(bencher: Bencher) { + let (left, right) = polygon_pairs(4); + bench_intersections(bencher, left, right); +} + +#[divan::bench] +fn thirty_two_vertex_boundaries(bencher: Bencher) { + let (left, right) = polygon_pairs(32); + bench_intersections(bencher, left, right); +} + +#[divan::bench] +fn nullable_rectangles(bencher: Bencher) { + let (left, right) = polygon_pairs(4); + let left = MaskedArray::try_new( + left, + Validity::from_iter((0..ROWS).map(|row| !row.is_multiple_of(8))), + ) + .unwrap() + .into_array(); + bench_intersections(bencher, left, right); +} diff --git a/vortex-spatial/src/extension/multipolygon.rs b/vortex-spatial/src/extension/multipolygon.rs index 80078a4b07e..2c30c096395 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -13,9 +13,11 @@ use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use geo_traits::to_geo::ToGeoGeometry; use geo_types::Geometry; +use geoarrow::array::GeoArrowArray; use geoarrow::array::GeoArrowArrayAccessor; use geoarrow::array::IntoArrow; use geoarrow::array::MultiPolygonArray; +use geoarrow::array::MultiPolygonBuilder; use geoarrow::datatypes::CoordType; use geoarrow::datatypes::MultiPolygonType; use prost::Message; @@ -24,6 +26,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; @@ -116,6 +119,27 @@ fn multipolygon_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) - MultiPolygonType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) } +/// Build a native 2-D [`MultiPolygon`] array from row-oriented `geo_types` multipolygons. +pub(crate) fn build_multipolygon_array( + multipolygons: &[Option>], + metadata: SpatialMetadata, + nullability: Nullability, +) -> VortexResult { + let multipolygons = MultiPolygonBuilder::from_nullable_multi_polygons( + multipolygons, + multipolygon_type(&metadata, Dimension::Xy), + ) + .finish(); + let storage_dtype = multipolygon_storage_dtype(Dimension::Xy, nullability); + let storage = ArrayRef::from_arrow( + multipolygons.to_array_ref().as_ref(), + nullability == Nullability::Nullable, + )? + .cast(storage_dtype.clone())?; + let ext_dtype = ExtDType::::try_new(metadata, storage_dtype)?; + Ok(ExtensionArray::try_new(ext_dtype.erased(), storage)?.into_array()) +} + /// Decode storage to `geo_types` for the spatial scalar functions (CRS is irrelevant to planar ops). pub(crate) fn multipolygon_geometries( storage: &ArrayRef, diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index e727a7ab3cb..d86c2c45c1c 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -13,9 +13,11 @@ use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use geo_traits::to_geo::ToGeoGeometry; use geo_types::Geometry; +use geoarrow::array::GeoArrowArray; use geoarrow::array::GeoArrowArrayAccessor; use geoarrow::array::IntoArrow; use geoarrow::array::PolygonArray; +use geoarrow::array::PolygonBuilder; use geoarrow::datatypes::CoordType; use geoarrow::datatypes::PolygonType; use prost::Message; @@ -24,6 +26,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; @@ -113,6 +116,25 @@ fn polygon_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) -> Pol PolygonType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) } +/// Build a native 2-D [`Polygon`] array from row-oriented `geo_types` polygons. +pub(crate) fn build_polygon_array( + polygons: &[Option>], + metadata: SpatialMetadata, + nullability: Nullability, +) -> VortexResult { + let polygons = + PolygonBuilder::from_nullable_polygons(polygons, polygon_type(&metadata, Dimension::Xy)) + .finish(); + let storage_dtype = polygon_storage_dtype(Dimension::Xy, nullability); + let storage = ArrayRef::from_arrow( + polygons.to_array_ref().as_ref(), + nullability == Nullability::Nullable, + )? + .cast(storage_dtype.clone())?; + let ext_dtype = ExtDType::::try_new(metadata, storage_dtype)?; + Ok(ExtensionArray::try_new(ext_dtype.erased(), storage)?.into_array()) +} + /// Decode `Polygon` storage (`List>`) to `geo_types` polygons, for the spatial scalar /// functions. CRS does not affect planar geometry ops, so default metadata is used. pub(crate) fn polygon_geometries( diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index 6bf96831c48..0a310b1ddb3 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -22,9 +22,12 @@ use crate::extension::WellKnownBinary; use crate::prune::SpatialDistancePrune; use crate::prune::SpatialIntersectsPrune; use crate::scalar_fn::area::SpatialArea; +use crate::scalar_fn::collect::SpatialCollect; use crate::scalar_fn::contains::SpatialContains; +use crate::scalar_fn::convex_hull::SpatialConvexHull; use crate::scalar_fn::distance::SpatialDistance; use crate::scalar_fn::envelope::SpatialEnvelope; +use crate::scalar_fn::intersection::SpatialIntersection; use crate::scalar_fn::intersects::SpatialIntersects; use crate::scalar_fn::make_line::SpatialMakeLine; @@ -67,7 +70,10 @@ pub fn initialize(session: &VortexSession) { // Register the geometry scalar functions. session.scalar_fns().register(SpatialArea); + session.scalar_fns().register(SpatialCollect); + session.scalar_fns().register(SpatialConvexHull); session.scalar_fns().register(SpatialEnvelope); + session.scalar_fns().register(SpatialIntersection); session.scalar_fns().register(SpatialContains); session.scalar_fns().register(SpatialDistance); session.scalar_fns().register(SpatialIntersects); diff --git a/vortex-spatial/src/scalar_fn/collect.rs b/vortex-spatial/src/scalar_fn/collect.rs new file mode 100644 index 00000000000..2f80401904e --- /dev/null +++ b/vortex-spatial/src/scalar_fn/collect.rs @@ -0,0 +1,554 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Collect`: collect homogeneous native geometries into their native multi-geometry type. + +use std::sync::Arc; + +use vortex_array::ArrayRef; +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::ScalarFnArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::arrays::listview::ListViewRebuildMode; +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::ExtDTypeRef; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +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::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::LineString; +use crate::extension::MultiLineString; +use crate::extension::MultiPoint; +use crate::extension::MultiPolygon; +use crate::extension::Point; +use crate::extension::Polygon; +use crate::scalar_fn::execute::Execution; +use crate::scalar_fn::execute::Operand; +use crate::scalar_fn::execute::dispatch_unary; + +/// Resolve the strict homogeneous `ST_Collect` overload for one list operand. +fn collect_dtype(dtypes: &[DType]) -> VortexResult { + vortex_ensure!( + dtypes.len() == 1, + "spatial: collect requires exactly one list operand, got {}", + dtypes.len() + ); + let DType::List(element_dtype, nullability) = &dtypes[0] else { + vortex_bail!("spatial: collect operand {} is not a list", dtypes[0]); + }; + // Multi-geometries cannot contain null components. Null list elements are ignored during + // execution, so their storage is non-nullable in the result. + let multi_storage = |element: &ExtDTypeRef| { + DType::List( + Arc::new(element.storage_dtype().as_nonnullable()), + *nullability, + ) + }; + match element_dtype.as_extension_opt() { + Some(element) if element.is::() => Ok(ExtDType::::try_new( + element.metadata::().clone(), + multi_storage(element), + )? + .erased()), + Some(element) if element.is::() => Ok(ExtDType::::try_new( + element.metadata::().clone(), + multi_storage(element), + )? + .erased()), + Some(element) if element.is::() => Ok(ExtDType::::try_new( + element.metadata::().clone(), + multi_storage(element), + )? + .erased()), + _ => vortex_bail!( + "spatial: collect list element {element_dtype} is not a native Point, LineString, \ + or Polygon" + ), + } +} + +/// Count valid elements in an exact list row without per-element mask lookups. +fn valid_count(mask: &Mask, start: usize, end: usize) -> usize { + match mask.bit_buffer() { + AllOr::All => end - start, + AllOr::None => 0, + AllOr::Some(bits) => bits.count_range(start, end), + } +} + +/// Rewrap a homogeneous geometry list as its corresponding multi-geometry array. +/// +/// The all-valid path reuses the geometry payload and list views. If geometry elements are null, +/// DuckDB semantics require ignoring them; that path first makes the views exact, then compacts the +/// payload and rebuilds the row views. Either way the output carries the input's zero-copy-to-list +/// flag, so a downstream `ListArray` conversion does not re-gather the reused payload. +fn collect_list( + mut list: ListViewArray, + validity: Validity, + output_dtype: &ExtDTypeRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut element_valid = list + .elements() + .validity()? + .execute_mask(list.elements().len(), ctx)?; + if !element_valid.all_true() { + list = list.rebuild(ListViewRebuildMode::MakeExact, ctx)?; + element_valid = list + .elements() + .validity()? + .execute_mask(list.elements().len(), ctx)?; + } + + // Both output paths keep the views exact: reuse forwards `offsets` and `sizes` untouched, and + // compaction rebuilds them as a running sum over the same element order. So the result is + // zero-copy to a `ListArray` exactly when `list` is, which `MakeExact` above has already + // ensured for every list that reaches compaction. + let zero_copy_to_list = list.is_zero_copy_to_list(); + let parts = list.into_data_parts(); + let elements = parts.elements.execute::(ctx)?; + let DType::List(target_element_storage, _) = output_dtype.storage_dtype() else { + unreachable!("collect output storage is always a list") + }; + let target_element_storage = target_element_storage.as_ref().clone(); + + let compact_elements = !element_valid.all_true(); + let element_storage = if compact_elements { + elements + .storage_array() + .filter(element_valid.clone())? + .cast(target_element_storage)? + } else { + elements.storage_array().cast(target_element_storage)? + }; + + let (offsets, sizes) = if compact_elements { + let old_offsets = parts + .offsets + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::>(ctx)?; + let old_sizes = parts + .sizes + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::>(ctx)?; + let mut offsets = BufferMut::::with_capacity(old_offsets.len()); + let mut sizes = BufferMut::::with_capacity(old_sizes.len()); + let mut next_offset = 0_u64; + + for (&old_offset, &old_size) in old_offsets.iter().zip(old_sizes.iter()) { + let start = usize::try_from(old_offset) + .map_err(|_| vortex_err!("spatial: collect element offset exceeds usize"))?; + let size = usize::try_from(old_size) + .map_err(|_| vortex_err!("spatial: collect element count exceeds usize"))?; + let end = start + .checked_add(size) + .ok_or_else(|| vortex_err!("spatial: collect element range overflows usize"))?; + vortex_ensure!( + end <= element_valid.len(), + "spatial: collect element range {start}..{end} exceeds element length {}", + element_valid.len() + ); + let size = u64::try_from(valid_count(&element_valid, start, end)) + .map_err(|_| vortex_err!("spatial: collect valid element count exceeds u64"))?; + offsets.push(next_offset); + sizes.push(size); + next_offset = next_offset + .checked_add(size) + .ok_or_else(|| vortex_err!("spatial: collect output offset exceeds u64"))?; + } + (offsets.into_array(), sizes.into_array()) + } else { + (parts.offsets, parts.sizes) + }; + + let storage = ListViewArray::try_new(element_storage, offsets, sizes, validity)?; + // SAFETY: `zero_copy_to_list` describes views this function either forwarded unchanged or + // replaced with a gapless, non-overlapping running sum over the same elements. Forwarding it + // matters: `list_from_list_view` re-gathers the whole payload for a list view that reports + // `false`, undoing the storage reuse above one operator later. + let storage = unsafe { storage.with_zero_copy_to_list(zero_copy_to_list) }.into_array(); + Ok(ExtensionArray::try_new(output_dtype.clone(), storage)?.into_array()) +} + +/// Execute the structural collect kernel after shared unary shape and null dispatch. +fn execute_collect( + execution: Execution<1, Validity>, + output_dtype: &ExtDTypeRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match execution.operands { + [Operand::Constant(scalar)] => { + let one = ConstantArray::new(scalar, 1) + .into_array() + .execute::(ctx)?; + let collected = collect_list( + one, + Validity::from_mask(Mask::new_true(1), execution.nullability), + output_dtype, + ctx, + )?; + Ok(ConstantArray::new(collected.execute_scalar(0, ctx)?, execution.len).into_array()) + } + [Operand::Column(array)] => { + let valid = execution.valid.execute_mask(execution.len, ctx)?; + collect_list( + array.execute::(ctx)?, + Validity::from_mask(valid, execution.nullability), + output_dtype, + ctx, + ) + } + } +} + +/// Collect a homogeneous list of native `Point`, `LineString`, or `Polygon` values into the +/// corresponding `MultiPoint`, `MultiLineString`, or `MultiPolygon` value. Null geometry elements +/// are ignored. Mixed geometry lists are rejected by the list element dtype rather than represented +/// as a geometry union. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct SpatialCollect; + +impl SpatialCollect { + /// A lazy `ScalarFnArray` collecting each list row into one native multi-geometry value. + pub fn try_new_array(array: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(SpatialCollect, EmptyOptions).erased(), + vec![array], + ) + } +} + +impl ScalarFnVTable for SpatialCollect { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.st.collect"); + *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("geometries"), + _ => unreachable!("collect has exactly one child"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + Ok(DType::Extension(collect_dtype(dtypes)?)) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let output_dtype = collect_dtype(std::slice::from_ref(input.dtype()))?; + dispatch_unary( + &input, + DType::Extension(output_dtype.clone()), + |execution, ctx| execute_collect(execution, &output_dtype, ctx), + ctx, + ) + } + + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Columnar; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::extension::ExtensionArrayExt; + use vortex_array::arrays::listview::ListViewArraySlotsExt; + 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_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::SpatialCollect; + 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_point_column; + use crate::test_harness::point_column; + use crate::test_harness::polygon_column; + + fn list_with_validity( + elements: ArrayRef, + offsets: &[u32], + validity: Validity, + ) -> VortexResult { + Ok(ListArray::try_new( + elements, + PrimitiveArray::from_iter(offsets.iter().copied()).into_array(), + validity, + )? + .into_array()) + } + + fn list(elements: ArrayRef, offsets: &[u32]) -> VortexResult { + list_with_validity(elements, offsets, Validity::NonNullable) + } + + #[test] + fn collects_points_into_multipoints() -> VortexResult<()> { + let points = point_column(vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0])?; + let input = list(points, &[0, 2, 3])?; + let expected = multipoint_column(vec![vec![(0.0, 3.0), (1.0, 4.0)], vec![(2.0, 5.0)]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn all_valid_collect_reuses_geometry_storage() -> VortexResult<()> { + let points = point_column(vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0])?; + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let point_storage = points + .clone() + .execute::(&mut ctx)? + .storage_array() + .clone(); + let input = list(points, &[0, 2, 3])?; + + let result = SpatialCollect::try_new_array(input)? + .into_array() + .execute::(&mut ctx)?; + let result_storage = result + .storage_array() + .clone() + .execute::(&mut ctx)?; + + assert!(ArrayRef::ptr_eq(&point_storage, result_storage.elements())); + Ok(()) + } + + /// A list view that forgets it is zero-copy to a list makes the next + /// `list_from_list_view` re-gather the payload that collect just reused. + #[rstest] + #[case::reused_elements(false)] + #[case::compacted_elements(true)] + fn output_stays_zero_copy_to_list(#[case] null_elements: bool) -> VortexResult<()> { + let points = if null_elements { + nullable_point_column(vec![Some((0.0, 3.0)), None, Some((2.0, 5.0))])? + } else { + point_column(vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0])? + }; + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let input = list(points, &[0, 2, 3])?; + assert!( + input + .clone() + .execute::(&mut ctx)? + .is_zero_copy_to_list(), + "a list column reaches collect as an exact list view" + ); + + let storage = SpatialCollect::try_new_array(input)? + .into_array() + .execute::(&mut ctx)? + .storage_array() + .clone() + .execute::(&mut ctx)?; + + assert!(storage.is_zero_copy_to_list()); + Ok(()) + } + + #[test] + fn collects_linestrings_into_multilinestrings() -> VortexResult<()> { + let line_a = vec![(0.0, 0.0), (1.0, 1.0)]; + let line_b = vec![(2.0, 2.0), (3.0, 3.0)]; + let line_c = vec![(4.0, 4.0), (5.0, 5.0)]; + let input = list( + linestring_column(vec![line_a.clone(), line_b.clone(), line_c.clone()])?, + &[0, 2, 3], + )?; + let expected = multilinestring_column(vec![vec![line_a, line_b], vec![line_c]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn collects_polygons_into_multipolygons() -> VortexResult<()> { + let polygon_a = vec![vec![(0.0, 0.0), (2.0, 0.0), (0.0, 2.0), (0.0, 0.0)]]; + let polygon_b = vec![vec![(3.0, 0.0), (5.0, 0.0), (3.0, 2.0), (3.0, 0.0)]]; + let polygon_c = vec![vec![(6.0, 0.0), (8.0, 0.0), (6.0, 2.0), (6.0, 0.0)]]; + let input = list( + polygon_column(vec![ + polygon_a.clone(), + polygon_b.clone(), + polygon_c.clone(), + ])?, + &[0, 2, 3], + )?; + let expected = multipolygon_column(vec![vec![polygon_a, polygon_b], vec![polygon_c]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn constant_list_remains_constant() -> VortexResult<()> { + let input = list( + nullable_point_column(vec![Some((0.0, 2.0)), None, Some((1.0, 3.0))])?, + &[0, 3], + )?; + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let scalar = input.execute_scalar(0, &mut ctx)?; + let input = ConstantArray::new(scalar, 3).into_array(); + + let result = SpatialCollect::try_new_array(input)?.into_array(); + let Columnar::Constant(constant) = result.clone().execute::(&mut ctx)? else { + return Err(vortex_err!( + "collect of a constant list should remain constant" + )); + }; + assert_eq!(constant.len(), 3); + let expected = multipoint_column(vec![ + vec![(0.0, 2.0), (1.0, 3.0)], + vec![(0.0, 2.0), (1.0, 3.0)], + vec![(0.0, 2.0), (1.0, 3.0)], + ])?; + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn ignores_null_geometry_elements() -> VortexResult<()> { + let points = nullable_point_column(vec![Some((0.0, 2.0)), None, Some((1.0, 3.0)), None])?; + let input = list(points, &[0, 2, 4])?; + let expected = multipoint_column(vec![vec![(0.0, 2.0)], vec![(1.0, 3.0)]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn all_null_geometry_elements_produce_empty_multi_geometry() -> VortexResult<()> { + let input = list(nullable_point_column(vec![None, None])?, &[0, 2])?; + let expected = multipoint_column(vec![vec![]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn propagates_null_list_rows() -> VortexResult<()> { + let input = list_with_validity( + point_column(vec![0.0, 1.0], vec![2.0, 3.0])?, + &[0, 1, 2], + Validity::from_iter([true, false]), + )?; + let expected = MaskedArray::try_new( + multipoint_column(vec![vec![(0.0, 2.0)], vec![(1.0, 3.0)]])?, + Validity::from_iter([true, false]), + )? + .into_array(); + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn rejects_unsupported_inputs() -> VortexResult<()> { + let point = point_column(vec![0.0], vec![0.0])?; + assert!(SpatialCollect::try_new_array(point).is_err()); + + let multipoints = multipoint_column(vec![vec![(0.0, 0.0)]])?; + assert!(SpatialCollect::try_new_array(list(multipoints, &[0, 1])?).is_err()); + + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!( + SpatialCollect + .return_dtype( + &EmptyOptions, + &[DType::List(primitive.into(), Nullability::NonNullable)] + ) + .is_err() + ); + Ok(()) + } +} diff --git a/vortex-spatial/src/scalar_fn/convex_hull.rs b/vortex-spatial/src/scalar_fn/convex_hull.rs new file mode 100644 index 00000000000..ffad2cd0c29 --- /dev/null +++ b/vortex-spatial/src/scalar_fn/convex_hull.rs @@ -0,0 +1,357 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_ConvexHull`: the planar convex hull of each native `MultiPoint`. + +use geo::ConvexHull; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtDTypeRef; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +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_bail; +use vortex_error::vortex_ensure; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::MultiPoint; +use crate::extension::Polygon; +use crate::extension::build_polygon_array; +use crate::extension::coordinate::Dimension; +use crate::extension::geometries; +use crate::extension::polygon_storage_dtype; +use crate::extension::single_geometry; +use crate::scalar_fn::execute::Execution; +use crate::scalar_fn::execute::Operand; +use crate::scalar_fn::execute::dispatch_unary; + +/// Resolve the strict native `MultiPoint -> Polygon` overload. +fn convex_hull_dtype(dtypes: &[DType]) -> VortexResult { + vortex_ensure!( + dtypes.len() == 1, + "spatial: convex_hull requires exactly one MultiPoint operand, got {}", + dtypes.len() + ); + let Some(input) = dtypes[0].as_extension_opt() else { + vortex_bail!( + "spatial: convex_hull operand {} is not a native MultiPoint", + dtypes[0] + ); + }; + vortex_ensure!( + input.is::(), + "spatial: convex_hull operand {} is not a native MultiPoint", + dtypes[0] + ); + + Ok(ExtDType::::try_new( + input.metadata::().clone(), + polygon_storage_dtype(Dimension::Xy, dtypes[0].nullability()), + )? + .erased()) +} + +/// Compute hulls for the valid rows and scatter them into a full-length native polygon array. +fn convex_hull_array( + array: ArrayRef, + valid: &Mask, + output_dtype: &ExtDTypeRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let decoded = geometries(&array.filter(valid.clone())?, ctx)?; + let hulls = decoded.iter().map(ConvexHull::convex_hull); + let polygons = match valid.indices() { + AllOr::All => hulls.map(Some).collect(), + AllOr::None => vec![None; array.len()], + AllOr::Some(rows) => { + let mut polygons = vec![None; array.len()]; + for (&row, hull) in rows.iter().zip(hulls) { + polygons[row] = Some(hull); + } + polygons + } + }; + build_polygon_array( + &polygons, + output_dtype.metadata::().clone(), + output_dtype.nullability(), + ) +} + +/// Execute convex hull after shared unary shape and null dispatch. +fn execute_convex_hull( + execution: Execution<1, Validity>, + output_dtype: &ExtDTypeRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match execution.operands { + [Operand::Constant(scalar)] => { + let hull = single_geometry(&scalar, ctx)?.convex_hull(); + let output = build_polygon_array( + &[Some(hull)], + output_dtype.metadata::().clone(), + output_dtype.nullability(), + )?; + Ok(ConstantArray::new(output.execute_scalar(0, ctx)?, execution.len).into_array()) + } + [Operand::Column(array)] => { + let valid = execution.valid.execute_mask(execution.len, ctx)?; + convex_hull_array(array, &valid, output_dtype, ctx) + } + } +} + +/// Compute the two-dimensional convex hull of each native `MultiPoint` as a native `Polygon`. +/// Empty, single-point, and collinear inputs remain typed polygons with degenerate exterior rings. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct SpatialConvexHull; + +impl SpatialConvexHull { + /// A lazy `ScalarFnArray` computing a polygon hull for each native `MultiPoint` row. + pub fn try_new_array(array: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(SpatialConvexHull, EmptyOptions).erased(), + vec![array], + ) + } +} + +impl ScalarFnVTable for SpatialConvexHull { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.st.convex_hull"); + *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("multipoint"), + _ => unreachable!("convex_hull has exactly one child"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + Ok(DType::Extension(convex_hull_dtype(dtypes)?)) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let output_dtype = convex_hull_dtype(std::slice::from_ref(input.dtype()))?; + dispatch_unary( + &input, + DType::Extension(output_dtype.clone()), + |execution, ctx| execute_convex_hull(execution, &output_dtype, ctx), + ctx, + ) + } + + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Columnar; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::SpatialConvexHull; + use crate::scalar_fn::area::SpatialArea; + use crate::scalar_fn::collect::SpatialCollect; + use crate::test_harness::multipoint_column; + use crate::test_harness::point_column; + use crate::test_harness::polygon_column; + + #[test] + fn computes_polygon_hulls() -> VortexResult<()> { + let input = multipoint_column(vec![vec![ + (0.0, 0.0), + (2.0, 0.0), + (2.0, 2.0), + (0.0, 2.0), + (1.0, 1.0), + ]])?; + let expected = polygon_column(vec![vec![vec![ + (2.0, 0.0), + (2.0, 2.0), + (0.0, 2.0), + (0.0, 0.0), + (2.0, 0.0), + ]]])?; + let result = SpatialConvexHull::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::empty(vec![], vec![])] + #[case::one_point( + vec![(1.0, 2.0)], + vec![vec![(1.0, 2.0), (1.0, 2.0)]] + )] + #[case::collinear( + vec![(0.0, 0.0), (1.0, 1.0), (2.0, 2.0)], + vec![vec![(0.0, 0.0), (2.0, 2.0), (0.0, 0.0)]] + )] + fn degenerate_hulls_remain_polygons( + #[case] points: Vec<(f64, f64)>, + #[case] expected_rings: Vec>, + ) -> VortexResult<()> { + let input = multipoint_column(vec![points])?; + let expected = polygon_column(vec![expected_rings])?; + let result = SpatialConvexHull::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn propagates_nulls() -> VortexResult<()> { + let input = MaskedArray::try_new( + multipoint_column(vec![ + vec![(0.0, 0.0), (1.0, 0.0), (0.0, 1.0)], + vec![(2.0, 2.0)], + ])?, + Validity::from_iter([true, false]), + )? + .into_array(); + let expected = MaskedArray::try_new( + polygon_column(vec![ + vec![vec![(0.0, 0.0), (1.0, 0.0), (0.0, 1.0), (0.0, 0.0)]], + vec![vec![(2.0, 2.0), (2.0, 2.0)]], + ])?, + Validity::from_iter([true, false]), + )? + .into_array(); + let result = SpatialConvexHull::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn constant_remains_constant() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let scalar = multipoint_column(vec![vec![(0.0, 0.0), (1.0, 0.0), (0.0, 1.0)]])? + .execute_scalar(0, &mut ctx)?; + let input = ConstantArray::new(scalar, 3).into_array(); + + let result = SpatialConvexHull::try_new_array(input)?.into_array(); + let Columnar::Constant(constant) = result.clone().execute::(&mut ctx)? else { + return Err(vortex_err!( + "convex_hull of a constant should remain constant" + )); + }; + assert_eq!(constant.len(), 3); + let hull = vec![vec![(0.0, 0.0), (1.0, 0.0), (0.0, 1.0), (0.0, 0.0)]]; + let expected = polygon_column(vec![hull.clone(), hull.clone(), hull])?; + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn collect_hull_area_pipeline() -> VortexResult<()> { + let points = point_column( + vec![0.0, 2.0, 2.0, 0.0, 1.0, 0.0, 1.0, 2.0], + vec![0.0, 0.0, 2.0, 2.0, 1.0, 0.0, 1.0, 2.0], + )?; + let point_lists = ListArray::try_new( + points, + PrimitiveArray::from_iter([0_u32, 5, 8]).into_array(), + Validity::NonNullable, + )? + .into_array(); + + let collected = SpatialCollect::try_new_array(point_lists)?.into_array(); + let hulls = SpatialConvexHull::try_new_array(collected)?.into_array(); + let areas = SpatialArea::try_new_array(hulls)?.into_array(); + let expected = PrimitiveArray::from_iter([4.0_f64, 0.0]).into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(areas, expected, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::none(0)] + #[case::two(2)] + fn rejects_wrong_arity(#[case] arity: usize) -> VortexResult<()> { + let dtype = multipoint_column(vec![vec![]])?.dtype().clone(); + assert!( + SpatialConvexHull + .return_dtype(&EmptyOptions, &vec![dtype; arity]) + .is_err() + ); + Ok(()) + } + + #[test] + fn rejects_non_multipoint_input() -> VortexResult<()> { + let input: ArrayRef = point_column(vec![0.0], vec![0.0])?; + assert!(SpatialConvexHull::try_new_array(input).is_err()); + Ok(()) + } +} diff --git a/vortex-spatial/src/scalar_fn/intersection.rs b/vortex-spatial/src/scalar_fn/intersection.rs new file mode 100644 index 00000000000..b7f6cf04b3c --- /dev/null +++ b/vortex-spatial/src/scalar_fn/intersection.rs @@ -0,0 +1,462 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Intersection`: pairwise planar intersection of native polygons. + +use geo::BooleanOps; +use geo_types::Geometry; +use geo_types::MultiPolygon as GeoMultiPolygon; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +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_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::AllOr; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::MultiPolygon; +use crate::extension::Polygon; +use crate::extension::SpatialMetadata; +use crate::extension::build_multipolygon_array; +use crate::extension::coordinate::Dimension; +use crate::extension::geometries; +use crate::extension::multipolygon_storage_dtype; +use crate::extension::single_geometry; +use crate::scalar_fn::execute::Execution; +use crate::scalar_fn::execute::Operand; +use crate::scalar_fn::execute::dispatch_binary; + +/// Resolve CRS metadata shared by two polygon operands. +fn intersection_metadata( + left: &SpatialMetadata, + right: &SpatialMetadata, +) -> VortexResult { + match (&left.crs, &right.crs) { + (Some(left_crs), Some(right_crs)) => { + vortex_ensure!( + left_crs == right_crs, + "spatial: intersection operands have different coordinate reference systems: \ + {left_crs} and {right_crs}" + ); + Ok(left.clone()) + } + (Some(_), None) => Ok(left.clone()), + (None, Some(_)) => Ok(right.clone()), + (None, None) => Ok(SpatialMetadata::default()), + } +} + +/// Metadata carried by a validated native polygonal dtype. +fn polygonal_metadata(dtype: &DType) -> &SpatialMetadata { + let extension = dtype.as_extension(); + if extension.is::() { + extension.metadata::() + } else if extension.is::() { + extension.metadata::() + } else { + unreachable!("intersection operand was validated as polygonal") + } +} + +/// Resolve the native polygonal intersection overloads, which always return a MultiPolygon. +fn intersection_dtype(dtypes: &[DType]) -> VortexResult> { + vortex_ensure!( + dtypes.len() == 2, + "spatial: intersection requires exactly two polygonal operands, got {}", + dtypes.len() + ); + for dtype in dtypes { + vortex_ensure!( + dtype.as_extension_opt().is_some_and(|extension| { + extension.is::() || extension.is::() + }), + "spatial: intersection operand {dtype} is not a native Polygon or MultiPolygon" + ); + } + + let metadata = intersection_metadata( + polygonal_metadata(&dtypes[0]), + polygonal_metadata(&dtypes[1]), + )?; + let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); + ExtDType::try_new( + metadata, + multipolygon_storage_dtype(Dimension::Xy, nullability), + ) +} + +/// Dispatch decoded geometry enums to `geo`'s concrete polygonal `BooleanOps` implementations. +fn polygonal_intersection(left: &Geometry, right: &Geometry) -> GeoMultiPolygon { + match (left, right) { + (Geometry::Polygon(left), Geometry::Polygon(right)) => left.intersection(right), + (Geometry::Polygon(left), Geometry::MultiPolygon(right)) => left.intersection(right), + (Geometry::MultiPolygon(left), Geometry::Polygon(right)) => left.intersection(right), + (Geometry::MultiPolygon(left), Geometry::MultiPolygon(right)) => left.intersection(right), + _ => unreachable!("intersection operands were validated as polygonal"), + } +} + +/// Execute intersection after shared binary shape and null dispatch. +fn execute_intersection( + execution: Execution<2>, + output_dtype: &ExtDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let intersections: Vec> = match &execution.operands { + [Operand::Constant(left), Operand::Constant(right)] => { + let intersection = + polygonal_intersection(&single_geometry(left, ctx)?, &single_geometry(right, ctx)?); + let one = build_multipolygon_array( + &[Some(intersection)], + output_dtype.metadata().clone(), + execution.nullability, + )?; + return Ok(ConstantArray::new(one.execute_scalar(0, ctx)?, execution.len).into_array()); + } + [Operand::Constant(left), Operand::Column(right)] => { + let left = single_geometry(left, ctx)?; + geometries(&right.filter(execution.valid.clone())?, ctx)? + .iter() + .map(|right| polygonal_intersection(&left, right)) + .collect() + } + [Operand::Column(left), Operand::Constant(right)] => { + let right = single_geometry(right, ctx)?; + geometries(&left.filter(execution.valid.clone())?, ctx)? + .iter() + .map(|left| polygonal_intersection(left, &right)) + .collect() + } + [Operand::Column(left), Operand::Column(right)] => { + let left = geometries(&left.filter(execution.valid.clone())?, ctx)?; + let right = geometries(&right.filter(execution.valid.clone())?, ctx)?; + left.iter() + .zip(&right) + .map(|(left, right)| polygonal_intersection(left, right)) + .collect() + } + }; + let intersections = match execution.valid.indices() { + AllOr::All => intersections.into_iter().map(Some).collect(), + AllOr::None => vec![None; execution.len], + AllOr::Some(rows) => { + let mut output = vec![None; execution.len]; + for (&row, intersection) in rows.iter().zip(intersections) { + output[row] = Some(intersection); + } + output + } + }; + build_multipolygon_array( + &intersections, + output_dtype.metadata().clone(), + execution.nullability, + ) +} + +/// Compute the pairwise two-dimensional intersection of native `Polygon` or `MultiPolygon` +/// operands as a native `MultiPolygon`. Disjoint and boundary-only intersections produce an +/// empty `MultiPolygon`. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct SpatialIntersection; + +impl SpatialIntersection { + /// A lazy `ScalarFnArray` intersecting two native polygonal operands by row. + pub fn try_new_array(left: ArrayRef, right: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(SpatialIntersection, EmptyOptions).erased(), + vec![left, right], + ) + } +} + +impl ScalarFnVTable for SpatialIntersection { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.st.intersection"); + *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(2) + } + + fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("left"), + 1 => ChildName::from("right"), + _ => unreachable!("intersection has exactly two children"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + Ok(DType::Extension(intersection_dtype(dtypes)?.erased())) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let left = args.get(0)?; + let right = args.get(1)?; + let output_dtype = intersection_dtype(&[left.dtype().clone(), right.dtype().clone()])?; + dispatch_binary( + &left, + &right, + DType::Extension(output_dtype.clone().erased()), + |execution, ctx| execute_intersection(execution, &output_dtype, ctx), + ctx, + ) + } + + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use geo_types::Geometry; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Columnar; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::SpatialIntersection; + use crate::extension::MultiPolygon; + use crate::extension::geometries; + use crate::scalar_fn::area::SpatialArea; + use crate::test_harness::multipolygon_column; + use crate::test_harness::point_column; + use crate::test_harness::polygon_column; + + fn square(xmin: f64, ymin: f64, xmax: f64, ymax: f64) -> Vec<(f64, f64)> { + vec![ + (xmin, ymin), + (xmax, ymin), + (xmax, ymax), + (xmin, ymax), + (xmin, ymin), + ] + } + + fn polygon_constant( + ring: Vec<(f64, f64)>, + len: usize, + ctx: &mut vortex_array::ExecutionCtx, + ) -> VortexResult { + let scalar = polygon_column(vec![vec![ring]])?.execute_scalar(0, ctx)?; + Ok(ConstantArray::new(scalar, len).into_array()) + } + + fn polygonal_column(ring: Vec<(f64, f64)>, multi: bool) -> VortexResult { + if multi { + multipolygon_column(vec![vec![vec![ring]]]) + } else { + polygon_column(vec![vec![ring]]) + } + } + + #[test] + fn q9_area_pipeline_handles_overlap_disjoint_and_touching() -> VortexResult<()> { + let left = polygon_column(vec![ + vec![square(0.0, 0.0, 2.0, 2.0)], + vec![square(0.0, 0.0, 1.0, 1.0)], + vec![square(0.0, 0.0, 1.0, 1.0)], + ])?; + let right = polygon_column(vec![ + vec![square(1.0, 1.0, 3.0, 3.0)], + vec![square(2.0, 2.0, 3.0, 3.0)], + vec![square(1.0, 0.0, 2.0, 1.0)], + ])?; + let intersections = SpatialIntersection::try_new_array(left, right)?.into_array(); + assert!(intersections.dtype().as_extension().is::()); + + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let decoded = geometries(&intersections, &mut ctx)?; + let polygon_counts = decoded + .iter() + .map(|geometry| match geometry { + Geometry::MultiPolygon(multipolygon) => Ok(multipolygon.0.len()), + other => Err(vortex_err!( + "intersection decoded as {other:?}, expected MultiPolygon" + )), + }) + .collect::>>()?; + assert_eq!(polygon_counts, [1, 0, 0]); + + let areas = SpatialArea::try_new_array(intersections)?.into_array(); + let expected = PrimitiveArray::from_iter([1.0_f64, 0.0, 0.0]).into_array(); + assert_arrays_eq!(areas, expected, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::polygon_polygon(false, false)] + #[case::polygon_multipolygon(false, true)] + #[case::multipolygon_polygon(true, false)] + #[case::multipolygon_multipolygon(true, true)] + fn supports_all_polygonal_combinations( + #[case] left_multi: bool, + #[case] right_multi: bool, + ) -> VortexResult<()> { + let left = polygonal_column(square(0.0, 0.0, 2.0, 2.0), left_multi)?; + let right = polygonal_column(square(1.0, 1.0, 3.0, 3.0), right_multi)?; + let intersections = SpatialIntersection::try_new_array(left, right)?.into_array(); + let areas = SpatialArea::try_new_array(intersections)?.into_array(); + let expected = PrimitiveArray::from_iter([1.0_f64]).into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(areas, expected, &mut ctx); + Ok(()) + } + + #[test] + fn preserves_holes() -> VortexResult<()> { + let left = polygon_column(vec![vec![ + square(0.0, 0.0, 4.0, 4.0), + square(1.0, 1.0, 3.0, 3.0), + ]])?; + let right = polygon_column(vec![vec![square(2.0, 0.0, 5.0, 4.0)]])?; + let intersections = SpatialIntersection::try_new_array(left, right)?.into_array(); + let areas = SpatialArea::try_new_array(intersections)?.into_array(); + let expected = PrimitiveArray::from_iter([6.0_f64]).into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(areas, expected, &mut ctx); + Ok(()) + } + + #[test] + fn propagates_nulls() -> VortexResult<()> { + let left = MaskedArray::try_new( + polygon_column(vec![ + vec![square(0.0, 0.0, 2.0, 2.0)], + vec![square(0.0, 0.0, 2.0, 2.0)], + ])?, + Validity::from_iter([true, false]), + )? + .into_array(); + let right = polygon_column(vec![ + vec![square(1.0, 1.0, 3.0, 3.0)], + vec![square(1.0, 1.0, 3.0, 3.0)], + ])?; + let intersections = SpatialIntersection::try_new_array(left, right)?.into_array(); + let areas = SpatialArea::try_new_array(intersections)?.into_array(); + let expected = PrimitiveArray::new(vec![1.0_f64, 0.0], Validity::from_iter([true, false])) + .into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(areas, expected, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::constant_left(true)] + #[case::constant_right(false)] + fn pairs_constants_with_columns(#[case] constant_left: bool) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let constant = polygon_constant(square(0.0, 0.0, 2.0, 2.0), 2, &mut ctx)?; + let column = polygon_column(vec![ + vec![square(1.0, 1.0, 3.0, 3.0)], + vec![square(3.0, 3.0, 4.0, 4.0)], + ])?; + let (left, right) = if constant_left { + (constant, column) + } else { + (column, constant) + }; + + let intersections = SpatialIntersection::try_new_array(left, right)?.into_array(); + let areas = SpatialArea::try_new_array(intersections)?.into_array(); + let expected = PrimitiveArray::from_iter([1.0_f64, 0.0]).into_array(); + assert_arrays_eq!(areas, expected, &mut ctx); + Ok(()) + } + + #[test] + fn two_constants_remain_constant() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let left = polygon_constant(square(0.0, 0.0, 2.0, 2.0), 3, &mut ctx)?; + let right = polygon_constant(square(1.0, 1.0, 3.0, 3.0), 3, &mut ctx)?; + + let result = SpatialIntersection::try_new_array(left, right)?.into_array(); + let Columnar::Constant(constant) = result.execute::(&mut ctx)? else { + return Err(vortex_err!( + "intersection of two constants should remain constant" + )); + }; + assert_eq!(constant.len(), 3); + Ok(()) + } + + #[rstest] + #[case::none(0)] + #[case::one(1)] + #[case::three(3)] + fn rejects_wrong_arity(#[case] arity: usize) -> VortexResult<()> { + let dtype = polygon_column(vec![vec![]])?.dtype().clone(); + assert!( + SpatialIntersection + .return_dtype(&EmptyOptions, &vec![dtype; arity]) + .is_err() + ); + Ok(()) + } + + #[test] + fn rejects_non_polygonal_input() -> VortexResult<()> { + let polygon = polygon_column(vec![vec![]])?; + let point = point_column(vec![0.0], vec![0.0])?; + assert!(SpatialIntersection::try_new_array(polygon, point).is_err()); + Ok(()) + } +} diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index 1dcff7d0b95..a0f1be88538 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -4,9 +4,12 @@ //! Geometry scalar functions over the native geometry extension types. pub mod area; +pub mod collect; pub mod contains; +pub mod convex_hull; pub mod distance; pub mod envelope; mod execute; +pub mod intersection; pub mod intersects; pub mod make_line;