diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index fde5b161dfd..a051f55ec5d 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -7,10 +7,10 @@ use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use super::super::Interleave; use super::super::InterleaveArrayExt; +use super::validate_selectors; use crate::array::Array; use crate::arrays::Bool; use crate::arrays::BoolArray; @@ -71,46 +71,18 @@ fn gather, R: AsPrimitive>( branches: &[A], rows: &[R], ) -> VortexResult { - let len = validate_selectors(value_bits, branches, rows)?; + let len = validate_selectors( + value_bits.len(), + |branch| value_bits[branch].len(), + branches, + rows, + )?; // SAFETY: `validate_selectors` proved `branches.len() == rows.len() == len`, and for every // `i < len` that `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()`. Ok(unsafe { gather_bits(len, value_bits, branches, rows) }) } -/// Validates the per-row selector bounds, returning the output length (`branches.len()`). -/// -/// On success, `rows.len() == branches.len() == len` and, for every `i < len`, -/// `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()` — exactly the -/// preconditions of [`gather_bits`]. Errors (rather than panics) on any out-of-bounds selector. -fn validate_selectors, R: AsPrimitive>( - value_bits: &[BitBuffer], - branches: &[A], - rows: &[R], -) -> VortexResult { - // The two selectors are validated to equal length at construction, which is the output length. - let len = branches.len(); - vortex_ensure!( - rows.len() == len, - "interleave selectors differ in length: array_indices {len}, row_indices {}", - rows.len() - ); - - for i in 0..len { - let branch = branches[i].as_(); - vortex_ensure!( - branch < value_bits.len(), - "interleave array index out of bounds" - ); - vortex_ensure!( - rows[i].as_() < value_bits[branch].len(), - "interleave row index out of bounds" - ); - } - - Ok(len) -} - /// Gathers one bit per output from `bits[branches[i]]` at position `rows[i]`, packing 64 results per /// word with [`BitBufferMut::collect_bool`]. /// diff --git a/vortex-array/src/arrays/interleave/execute/mod.rs b/vortex-array/src/arrays/interleave/execute/mod.rs index 05dcd161f62..a8047968136 100644 --- a/vortex-array/src/arrays/interleave/execute/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/mod.rs @@ -5,14 +5,16 @@ //! //! All values share a type (validated in [`Interleave::check`]), so the //! physical gather kernel is chosen from the first value. The selector types are an orthogonal -//! concern handled within each kernel. Only boolean values are implemented today (see the [`bool`] module). +//! concern handled within each kernel. //! //! [`Interleave::check`]: super::Interleave::check -//! [`bool`]: module@crate::arrays::interleave::execute::bool mod bool; +mod primitive; +use num_traits::AsPrimitive; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use super::Interleave; @@ -28,12 +30,43 @@ pub(super) fn execute( ) -> VortexResult { if array.value(0).dtype().is_boolean() { bool::execute(array, ctx) + } else if array.value(0).dtype().is_primitive() { + primitive::execute(array, ctx) } else { let value_dtype = array.value(0).dtype().clone(); vortex_panic!( - "interleave execution is only implemented for boolean values; value dtype {} is not \ - yet supported", + "interleave execution is not implemented for value dtype {}", value_dtype ) } } + +/// Validate selector lengths and bounds, returning the output length. +fn validate_selectors( + num_values: usize, + value_len: F, + branches: &[A], + rows: &[R], +) -> VortexResult +where + A: AsPrimitive, + R: AsPrimitive, + F: Fn(usize) -> usize, +{ + let len = branches.len(); + vortex_ensure!( + rows.len() == len, + "interleave selectors differ in length: array_indices {len}, row_indices {}", + rows.len() + ); + + for i in 0..len { + let branch = branches[i].as_(); + vortex_ensure!(branch < num_values, "interleave array index out of bounds"); + vortex_ensure!( + rows[i].as_() < value_len(branch), + "interleave row index out of bounds" + ); + } + Ok(len) +} diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs new file mode 100644 index 00000000000..40e59b5847d --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution for primitive [`Interleave`] values. + +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use super::super::Interleave; +use super::super::InterleaveArrayExt; +use super::validate_selectors; +use crate::array::Array; +use crate::array::ArrayView; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::dtype::NativePType; +use crate::executor::ExecutionCtx; +use crate::executor::ExecutionResult; +use crate::match_each_native_ptype; +use crate::match_each_unsigned_integer_ptype; +use crate::require_child; + +pub(super) fn execute( + array: Array, + _ctx: &mut ExecutionCtx, +) -> VortexResult { + let num_values = array.num_values(); + let mut array = array; + array = require_child!(array, array.array_indices(), 0 => Primitive); + array = require_child!(array, array.row_indices(), 1 => Primitive); + for i in 0..num_values { + array = require_child!(array, array.value(i), i + 2 => Primitive); + } + + let validity = array.as_ref().validity()?; + let output = match_each_native_ptype!(array.value(0).as_::().ptype(), |T| { + let values = gather_values::(&array)?; + VortexResult::Ok(PrimitiveArray::new(values, validity)) + })?; + + Ok(ExecutionResult::done(output)) +} + +fn gather_values(array: &Array) -> VortexResult> { + let buffers = (0..array.num_values()) + .map(|i| array.value(i).as_::().to_buffer::()) + .collect::>(); + let branches = array.array_indices().as_::(); + let rows = array.row_indices().as_::(); + + match_each_unsigned_integer_ptype!(branches.ptype(), |A| { + gather_rows::(&buffers, branches.as_slice::(), rows) + }) +} + +fn gather_rows( + values: &[Buffer], + branches: &[A], + rows: ArrayView<'_, Primitive>, +) -> VortexResult> +where + T: NativePType, + A: AsPrimitive, +{ + match_each_unsigned_integer_ptype!(rows.ptype(), |R| { + gather(values, branches, rows.as_slice::()) + }) +} + +fn gather(values: &[Buffer], branches: &[A], rows: &[R]) -> VortexResult> +where + T: NativePType, + A: AsPrimitive, + R: AsPrimitive, +{ + let len = validate_selectors(values.len(), |branch| values[branch].len(), branches, rows)?; + let mut output = BufferMut::with_capacity(len); + for i in 0..len { + output.push(values[branches[i].as_()][rows[i].as_()]); + } + Ok(output.freeze()) +} diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index bff03ab055f..9b63b721c3f 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -719,17 +719,18 @@ mod tests { } #[test] - #[should_panic(expected = "only implemented for boolean values")] - fn non_boolean_value_execution_panics() { - // Execution dispatches on the value type: primitive values have no kernel yet. - let v0 = PrimitiveArray::from_iter([1u32]).into_array(); - let v1 = PrimitiveArray::from_iter([2u32]).into_array(); - let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array(); - let row_indices = PrimitiveArray::from_iter([0u32, 0]).into_array(); - let interleaved = InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices) - .vortex_expect("primitive values should construct") - .into_array(); + fn executes_primitive_values() -> VortexResult<()> { + let v0 = PrimitiveArray::from_iter([1.0f64, 2.0]).into_array(); + let v1 = PrimitiveArray::from_option_iter([Some(10.0f64), None]).into_array(); + let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array(); + let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices)?.into_array(); + let expected = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(10.0), Some(2.0), None]) + .into_array(); let mut ctx = array_session().create_execution_ctx(); - interleaved.execute::(&mut ctx).ok(); + assert_arrays_eq!(interleaved, expected, &mut ctx); + Ok(()) } } diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index 6eddab411d9..c8a0e8a318d 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -58,5 +58,33 @@ harness = false name = "distance" harness = false +[[bench]] +name = "make_line" +harness = false + +[[bench]] +name = "length" +harness = false + +[[bench]] +name = "area" +harness = false + +[[bench]] +name = "collect" +harness = false + +[[bench]] +name = "convex_hull" +harness = false + +[[bench]] +name = "intersection" +harness = false + +[[bench]] +name = "hilbert" +harness = false + [lints] workspace = true diff --git a/vortex-geo/benches/area.rs b/vortex-geo/benches/area.rs new file mode 100644 index 00000000000..38b1bc138e2 --- /dev/null +++ b/vortex-geo/benches/area.rs @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for native `ST_Area` over polygons and multipolygons. +//! +//! The cases separate the costs of vertex traversal, interior rings, nested polygons, and strict +//! null propagation. They execute through the scalar function and materialize the `f64` result. +//! +//! Run with `cargo bench -p vortex-geo --bench area`. + +#![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::MaskedArray; +use vortex_array::validity::Validity; +use vortex_geo::scalar_fn::area::GeoArea; +use vortex_geo::test_harness::geo_session; +use vortex_geo::test_harness::multipolygon_column; +use vortex_geo::test_harness::polygon_column; +use vortex_session::VortexSession; + +// 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(geo_session); + +const ROWS: usize = 512; + +fn main() { + divan::main(); +} + +/// A closed square ring centered at `(cx, cy)`. +fn square(cx: f64, cy: f64, radius: f64) -> Vec<(f64, f64)> { + vec![ + (cx - radius, cy - radius), + (cx + radius, cy - radius), + (cx + radius, cy + radius), + (cx - radius, cy + radius), + (cx - radius, cy - radius), + ] +} + +fn simple_polygons() -> ArrayRef { + polygon_column( + (0..ROWS) + .map(|row| vec![square(row as f64, row as f64, 10.0)]) + .collect(), + ) + .unwrap() +} + +fn polygons_with_holes() -> ArrayRef { + polygon_column( + (0..ROWS) + .map(|row| { + let center = row as f64; + vec![ + square(center, center, 10.0), + square(center - 4.0, center, 1.0), + square(center + 4.0, center, 1.0), + ] + }) + .collect(), + ) + .unwrap() +} + +fn multipolygons() -> ArrayRef { + multipolygon_column( + (0..ROWS) + .map(|row| { + let center = row as f64; + vec![ + vec![square(center - 12.0, center, 5.0)], + vec![square(center, center, 5.0)], + vec![square(center + 12.0, center, 5.0)], + ] + }) + .collect(), + ) + .unwrap() +} + +fn areas(geometry: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + GeoArea::try_new_array(geometry.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +fn bench_area(bencher: Bencher, geometry: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| areas(&geometry, &mut ctx)); +} + +#[divan::bench] +fn simple_polygon(bencher: Bencher) { + bench_area(bencher, simple_polygons()); +} + +#[divan::bench] +fn polygon_with_holes(bencher: Bencher) { + bench_area(bencher, polygons_with_holes()); +} + +#[divan::bench] +fn multipolygon(bencher: Bencher) { + bench_area(bencher, multipolygons()); +} + +#[divan::bench] +fn nullable_polygon(bencher: Bencher) { + let geometry = MaskedArray::try_new( + simple_polygons(), + Validity::from_iter((0..ROWS).map(|row| !row.is_multiple_of(8))), + ) + .unwrap() + .into_array(); + bench_area(bencher, geometry); +} diff --git a/vortex-geo/benches/collect.rs b/vortex-geo/benches/collect.rs new file mode 100644 index 00000000000..d9b0f52acc5 --- /dev/null +++ b/vortex-geo/benches/collect.rs @@ -0,0 +1,149 @@ +// 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-geo --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_geo::scalar_fn::collect::GeoCollect; +use vortex_geo::test_harness::geo_session; +use vortex_geo::test_harness::linestring_column; +use vortex_geo::test_harness::nullable_point_column; +use vortex_geo::test_harness::point_column; +use vortex_geo::test_harness::polygon_column; +use vortex_session::VortexSession; + +// 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(geo_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 { + GeoCollect::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)); +} diff --git a/vortex-geo/benches/convex_hull.rs b/vortex-geo/benches/convex_hull.rs new file mode 100644 index 00000000000..0c3d5fdc684 --- /dev/null +++ b/vortex-geo/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-geo --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_geo::scalar_fn::convex_hull::GeoConvexHull; +use vortex_geo::test_harness::geo_session; +use vortex_geo::test_harness::multipoint_column; +use vortex_session::VortexSession; + +// 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(geo_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 { + GeoConvexHull::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-geo/benches/hilbert.rs b/vortex-geo/benches/hilbert.rs new file mode 100644 index 00000000000..a9477bc9e68 --- /dev/null +++ b/vortex-geo/benches/hilbert.rs @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for `vortex.geo.hilbert` with constant `Rect` bounds. +//! +//! Run with `cargo bench -p vortex-geo --bench hilbert`. + +#![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::scalar::Scalar; +use vortex_geo::scalar_fn::hilbert::GeoHilbert; +use vortex_geo::test_harness::MultiPolygonRings; +use vortex_geo::test_harness::geo_session; +use vortex_geo::test_harness::multipolygon_column; +use vortex_geo::test_harness::point_column; +use vortex_geo::test_harness::rect_column; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(geo_session); + +const ROWS: usize = 1 << 9; + +fn main() { + divan::main(); +} + +fn ordinate(i: usize) -> f64 { + (i.wrapping_mul(2_654_435_761) % 1_000) as f64 +} + +fn bounds(ctx: &mut ExecutionCtx) -> Scalar { + rect_column(vec![(0.0, 0.0, 1_000.0, 1_000.0)]) + .unwrap() + .execute_scalar(0, ctx) + .unwrap() +} + +fn hilbert(column: &ArrayRef, bounds: &Scalar, ctx: &mut ExecutionCtx) -> ArrayRef { + GeoHilbert::try_new_array(column.clone(), bounds.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +fn bench_hilbert(bencher: Bencher, column: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + let bounds = bounds(&mut ctx); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| hilbert(&column, &bounds, &mut ctx)); +} + +#[divan::bench] +fn points(bencher: Bencher) { + let column = point_column( + (0..ROWS).map(ordinate).collect(), + (0..ROWS).map(|row| ordinate(row + 1)).collect(), + ) + .unwrap(); + bench_hilbert(bencher, column); +} + +fn multipolygon(row: usize) -> MultiPolygonRings { + let ring = |offset: usize| { + (0..8) + .map(|vertex| { + ( + ordinate(row + offset + vertex), + ordinate(row + offset + vertex + 1), + ) + }) + .collect() + }; + vec![vec![ring(0), ring(8)], vec![ring(16), ring(24)]] +} + +#[divan::bench] +fn multipolygons(bencher: Bencher) { + let column = multipolygon_column((0..ROWS).map(multipolygon).collect()).unwrap(); + bench_hilbert(bencher, column); +} diff --git a/vortex-geo/benches/intersection.rs b/vortex-geo/benches/intersection.rs new file mode 100644 index 00000000000..b776a49f538 --- /dev/null +++ b/vortex-geo/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-geo --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_geo::scalar_fn::intersection::GeoIntersection; +use vortex_geo::test_harness::geo_session; +use vortex_geo::test_harness::polygon_column; +use vortex_session::VortexSession; + +// 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(geo_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 { + GeoIntersection::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-geo/benches/length.rs b/vortex-geo/benches/length.rs new file mode 100644 index 00000000000..5d267eba600 --- /dev/null +++ b/vortex-geo/benches/length.rs @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for native `ST_Length` over LineStrings. +//! +//! The two-vertex case tracks ordinary route segments, while the longer-line case captures the +//! per-vertex traversal cost. The nullable case measures strict null propagation separately. +//! +//! Run with `cargo bench -p vortex-geo --bench length`. + +#![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::MaskedArray; +use vortex_array::validity::Validity; +use vortex_geo::scalar_fn::length::GeoLength; +use vortex_geo::test_harness::geo_session; +use vortex_geo::test_harness::linestring_column; +use vortex_session::VortexSession; + +// 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(geo_session); + +const ROWS: usize = 512; + +fn main() { + divan::main(); +} + +/// A deterministic vertex ordinate. +fn ordinate(i: usize) -> f64 { + (i.wrapping_mul(2_654_435_761) % 10_000) as f64 / 100.0 +} + +fn linestrings(vertices: usize) -> ArrayRef { + linestring_column( + (0..ROWS) + .map(|row| { + (0..vertices) + .map(|vertex| (ordinate(row + vertex), ordinate(row + vertex + 1))) + .collect() + }) + .collect(), + ) + .unwrap() +} + +fn lengths(lines: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + GeoLength::try_new_array(lines.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +#[divan::bench] +fn two_vertex_lines(bencher: Bencher) { + let lines = linestrings(2); + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| lengths(&lines, &mut ctx)); +} + +#[divan::bench] +fn sixteen_vertex_lines(bencher: Bencher) { + let lines = linestrings(16); + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| lengths(&lines, &mut ctx)); +} + +#[divan::bench] +fn nullable_two_vertex_lines(bencher: Bencher) { + let lines = MaskedArray::try_new( + linestrings(2), + Validity::from_iter((0..ROWS).map(|i| !i.is_multiple_of(8))), + ) + .unwrap() + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| lengths(&lines, &mut ctx)); +} diff --git a/vortex-geo/benches/make_line.rs b/vortex-geo/benches/make_line.rs new file mode 100644 index 00000000000..5f87b7e0f2c --- /dev/null +++ b/vortex-geo/benches/make_line.rs @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for native `ST_MakeLine`. +//! +//! The cases cover the normal paired-column operation, a broadcast point constant, and strict +//! null propagation. They execute the result to its canonical representation so the benchmark +//! includes construction of the two-vertex line storage. +//! +//! Run with `cargo bench -p vortex-geo --bench make_line`. + +#![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::ConstantArray; +use vortex_geo::scalar_fn::make_line::GeoMakeLine; +use vortex_geo::test_harness::geo_session; +use vortex_geo::test_harness::nullable_point_column; +use vortex_geo::test_harness::point_column; +use vortex_session::VortexSession; + +// 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(geo_session); + +const ROWS: usize = 512; + +fn main() { + divan::main(); +} + +/// Deterministic pseudo-random value in `[0, 1)`. +fn unit(i: usize) -> f64 { + ((i.wrapping_mul(2_654_435_761) >> 8) % 10_000) as f64 / 10_000.0 +} + +fn points(offset: usize) -> ArrayRef { + let xs = (0..ROWS) + .map(|i| 300.0 * unit(i + offset) - 150.0) + .collect(); + let ys = (0..ROWS) + .map(|i| 300.0 * unit(i + offset + 1) - 150.0) + .collect(); + point_column(xs, ys).unwrap() +} + +fn nullable_points(offset: usize, null_every: usize) -> ArrayRef { + nullable_point_column( + (0..ROWS) + .map(|i| { + (!i.is_multiple_of(null_every)).then(|| { + ( + 300.0 * unit(i + offset) - 150.0, + 300.0 * unit(i + offset + 1) - 150.0, + ) + }) + }) + .collect(), + ) + .unwrap() +} + +fn point_constant(ctx: &mut ExecutionCtx) -> ArrayRef { + let scalar = point_column(vec![0.0], vec![0.0]) + .unwrap() + .execute_scalar(0, ctx) + .unwrap(); + ConstantArray::new(scalar, ROWS).into_array() +} + +fn make_lines(starts: &ArrayRef, ends: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + GeoMakeLine::try_new_array(starts.clone(), ends.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +#[divan::bench] +fn column_x_column(bencher: Bencher) { + let starts = points(0); + let ends = points(97); + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| make_lines(&starts, &ends, &mut ctx)); +} + +#[divan::bench] +fn column_x_constant(bencher: Bencher) { + let starts = points(0); + let mut ctx = SESSION.create_execution_ctx(); + let end = point_constant(&mut ctx); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| make_lines(&starts, &end, &mut ctx)); +} + +#[divan::bench] +fn nullable_columns(bencher: Bencher) { + let starts = nullable_points(0, 8); + let ends = nullable_points(97, 11); + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| make_lines(&starts, &ends, &mut ctx)); +} diff --git a/vortex-geo/src/extension/coordinate.rs b/vortex-geo/src/extension/coordinate.rs index dc3537cfb30..43599a705b7 100644 --- a/vortex-geo/src/extension/coordinate.rs +++ b/vortex-geo/src/extension/coordinate.rs @@ -72,6 +72,21 @@ impl Dimension { Dimension::Xyzm => &["x", "y", "z", "m"], } } + + /// Promote two coordinate dimensions to the smallest dimension that represents both. + /// + /// Missing `z`/`m` ordinates are materialized as zero when values are converted to this + /// dimension, matching DuckDB Spatial's `ST_MakeLine` promotion. + pub(crate) fn promote(self, other: Self) -> Self { + match (self, other) { + (Self::Xyzm, _) | (_, Self::Xyzm) | (Self::Xyz, Self::Xym) | (Self::Xym, Self::Xyz) => { + Self::Xyzm + } + (Self::Xyz, _) | (_, Self::Xyz) => Self::Xyz, + (Self::Xym, _) | (_, Self::Xym) => Self::Xym, + (Self::Xy, Self::Xy) => Self::Xy, + } + } } impl From for Dimension { diff --git a/vortex-geo/src/extension/linestring.rs b/vortex-geo/src/extension/linestring.rs index 274d20f28ba..b627bcafc83 100644 --- a/vortex-geo/src/extension/linestring.rs +++ b/vortex-geo/src/extension/linestring.rs @@ -22,14 +22,22 @@ use prost::Message; 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::InterleaveArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +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::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; +use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; @@ -38,10 +46,12 @@ use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::FromArrowArray; use vortex_arrow::FromArrowType; +use vortex_buffer::Buffer; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_session::registry::CachedId; use vortex_session::registry::Id; @@ -103,6 +113,75 @@ pub(crate) fn linestring_dimension(dtype: &DType) -> VortexResult { coordinate_dimension(coords) } +/// Build one native [`LineString`] per corresponding pair of point coordinate rows. +pub(crate) fn linestring_array_from_point_pairs( + ext_dtype: &ExtDType, + starts: &StructArray, + ends: &StructArray, + validity: Validity, +) -> VortexResult { + let len = starts.len(); + vortex_ensure_eq!( + len, + ends.len(), + "geo: line string point columns must have equal lengths" + ); + let vertex_count = len + .checked_mul(2) + .ok_or_else(|| vortex_err!("geo: two-vertex line string length overflow"))?; + let dimension = linestring_dimension(ext_dtype.storage_dtype())?; + let start_dimension = coordinate_dimension(starts.dtype())?; + let end_dimension = coordinate_dimension(ends.dtype())?; + + let array_indices = PrimitiveArray::from_iter((0..len).flat_map(|_| [0u8, 1])).into_array(); + let rows = (0..len) + .flat_map(|row| [row, row]) + .map(|row| { + u64::try_from(row).map_err(|_| vortex_err!("geo: line string row index overflow")) + }) + .collect::>>()?; + let row_indices = Buffer::from(rows).into_array(); + let ordinate = + |points: &StructArray, point_dimension: Dimension, name: &str| -> VortexResult { + if point_dimension.field_names().contains(&name) { + points.unmasked_field_by_name(name).cloned() + } else { + Ok(ConstantArray::new(0.0f64, len).into_array()) + } + }; + let ordinates = dimension + .field_names() + .iter() + .map(|name| { + Ok(InterleaveArray::try_new( + vec![ + ordinate(starts, start_dimension, name)?, + ordinate(ends, end_dimension, name)?, + ], + array_indices.clone(), + row_indices.clone(), + )? + .into_array()) + }) + .collect::>>()?; + let vertices = StructArray::try_new( + FieldNames::from(dimension.field_names()), + ordinates, + vertex_count, + Validity::NonNullable, + )? + .into_array(); + let offsets = (0..=len) + .map(|row| { + i32::try_from(row * 2) + .map_err(|_| vortex_err!("geo: two-vertex line string offset overflow")) + }) + .collect::>>()?; + let storage = + ListArray::try_new(vertices, Buffer::from(offsets).into_array(), validity)?.into_array(); + Ok(ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array()) +} + static ARROW_LINESTRING: CachedId = CachedId::new(LineStringType::NAME); /// The `geoarrow.linestring` extension type for `dimension`, with separated (struct) coordinates diff --git a/vortex-geo/src/extension/multipolygon.rs b/vortex-geo/src/extension/multipolygon.rs index 524e470749c..124b545a5da 100644 --- a/vortex-geo/src/extension/multipolygon.rs +++ b/vortex-geo/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; @@ -117,6 +120,27 @@ fn multipolygon_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> MultiP MultiPolygonType::new(dimension.into(), geoarrow_metadata(geo_metadata)) } +/// Build a native 2-D [`MultiPolygon`] array from row-oriented `geo_types` multipolygons. +pub(crate) fn build_multipolygon_array( + multipolygons: &[Option>], + metadata: GeoMetadata, + 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 geo scalar functions (CRS is irrelevant to planar ops). pub(crate) fn multipolygon_geometries( storage: &ArrayRef, diff --git a/vortex-geo/src/extension/polygon.rs b/vortex-geo/src/extension/polygon.rs index 9a74c3ce7b3..9b825dfb1bf 100644 --- a/vortex-geo/src/extension/polygon.rs +++ b/vortex-geo/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; @@ -114,6 +117,25 @@ fn polygon_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> PolygonType PolygonType::new(dimension.into(), geoarrow_metadata(geo_metadata)) } +/// Build a native 2-D [`Polygon`] array from row-oriented `geo_types` polygons. +pub(crate) fn build_polygon_array( + polygons: &[Option>], + metadata: GeoMetadata, + 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 geo scalar /// functions. CRS does not affect planar geometry ops, so default metadata is used. pub(crate) fn polygon_geometries( diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index b72df92a3f5..6b0e526cbfd 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -21,10 +21,17 @@ use crate::extension::Rect; use crate::extension::WellKnownBinary; use crate::prune::GeoDistancePrune; use crate::prune::GeoIntersectsPrune; +use crate::scalar_fn::area::GeoArea; +use crate::scalar_fn::collect::GeoCollect; use crate::scalar_fn::contains::GeoContains; +use crate::scalar_fn::convex_hull::GeoConvexHull; use crate::scalar_fn::distance::GeoDistance; use crate::scalar_fn::envelope::GeoEnvelope; +use crate::scalar_fn::hilbert::GeoHilbert; +use crate::scalar_fn::intersection::GeoIntersection; use crate::scalar_fn::intersects::GeoIntersects; +use crate::scalar_fn::length::GeoLength; +use crate::scalar_fn::make_line::GeoMakeLine; pub mod aggregate_fn; pub mod extension; @@ -64,10 +71,17 @@ pub fn initialize(session: &VortexSession) { session.arrow().register_importer(Arc::new(Rect)); // Register the geometry scalar functions. + session.scalar_fns().register(GeoArea); + session.scalar_fns().register(GeoCollect); + session.scalar_fns().register(GeoConvexHull); session.scalar_fns().register(GeoEnvelope); session.scalar_fns().register(GeoContains); session.scalar_fns().register(GeoDistance); + session.scalar_fns().register(GeoHilbert); session.scalar_fns().register(GeoIntersects); + session.scalar_fns().register(GeoIntersection); + session.scalar_fns().register(GeoLength); + session.scalar_fns().register(GeoMakeLine); // The axis-aligned bounding-box (AABB) aggregate; self-declares as a per-chunk zone stat for // geometry columns. diff --git a/vortex-geo/src/scalar_fn/area.rs b/vortex-geo/src/scalar_fn/area.rs new file mode 100644 index 00000000000..ba2cff58b2a --- /dev/null +++ b/vortex-geo/src/scalar_fn/area.rs @@ -0,0 +1,246 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Area`: unsigned planar area of native geometries. + +use geo::Area; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +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_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::is_native_geometry; +use crate::scalar_fn::execute::execute_unary_geo_types; + +/// Validate the native geometry operand accepted by `ST_Area`. +fn validate_area_operand(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure!( + dtypes.len() == 1, + "geo: area requires exactly one geometry operand, got {}", + dtypes.len() + ); + vortex_ensure!( + is_native_geometry(&dtypes[0]), + "geo: area operand {} is not a native geometry", + dtypes[0] + ); + Ok(()) +} + +/// Unsigned planar `ST_Area` of native geometries. +/// +/// Points and line strings have zero area, polygons and multipolygons use their two-dimensional +/// coordinates, and rectangles use width times height. Higher coordinate dimensions are ignored. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct GeoArea; + +impl GeoArea { + /// A lazy `ScalarFnArray` computing the per-row area of a native geometry operand. + pub fn try_new_array(array: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(GeoArea, EmptyOptions).erased(), + vec![array], + ) + } +} + +impl ScalarFnVTable for GeoArea { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.area"); + *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!("area has exactly one child"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_area_operand(dtypes)?; + Ok(DType::Primitive(PType::F64, dtypes[0].nullability())) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let array = args.get(0)?; + execute_unary_geo_types(&array, Area::unsigned_area, 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::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + 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 super::GeoArea; + 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::point_column; + use crate::test_harness::polygon_column; + use crate::test_harness::rect_column; + + #[rstest] + #[case::point(point_column(vec![1.0], vec![2.0]), &[0.0])] + #[case::line_string( + linestring_column(vec![vec![(0.0, 0.0), (3.0, 4.0)]]), + &[0.0] + )] + #[case::multi_point( + multipoint_column(vec![vec![(0.0, 0.0), (1.0, 1.0)]]), + &[0.0] + )] + #[case::multi_line_string( + multilinestring_column(vec![vec![ + vec![(0.0, 0.0), (1.0, 1.0)], + vec![(2.0, 2.0), (3.0, 3.0)], + ]]), + &[0.0] + )] + #[case::polygon( + polygon_column(vec![ + vec![ + vec![(0.0, 0.0), (4.0, 0.0), (4.0, 3.0), (0.0, 3.0), (0.0, 0.0)], + vec![(1.0, 1.0), (2.0, 1.0), (2.0, 2.0), (1.0, 2.0), (1.0, 1.0)], + ], + vec![], + ]), + &[11.0, 0.0] + )] + #[case::multi_polygon( + multipolygon_column(vec![vec![ + vec![vec![ + (0.0, 0.0), + (2.0, 0.0), + (2.0, 2.0), + (0.0, 2.0), + (0.0, 0.0), + ]], + vec![vec![ + (3.0, 0.0), + (6.0, 0.0), + (6.0, 3.0), + (3.0, 3.0), + (3.0, 0.0), + ]], + ]]), + &[13.0] + )] + #[case::rect(rect_column(vec![(0.0, 0.0, 5.0, 3.0)]), &[15.0])] + fn measures_native_geometries( + #[case] geometry: VortexResult, + #[case] expected: &[f64], + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let areas = GeoArea::try_new_array(geometry?)?.into_array(); + let expected = PrimitiveArray::from_iter(expected.iter().copied()).into_array(); + assert_arrays_eq!(areas, expected, &mut ctx); + Ok(()) + } + + #[test] + fn propagates_nulls() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let multipolygons = nullable_multipolygon_column(vec![ + Some(vec![vec![vec![ + (0.0, 0.0), + (2.0, 0.0), + (2.0, 2.0), + (0.0, 2.0), + (0.0, 0.0), + ]]]), + None, + ])?; + let areas = GeoArea::try_new_array(multipolygons)?.into_array(); + let expected = + PrimitiveArray::new(vec![4.0f64, 0.0], Validity::from_iter([true, false])).into_array(); + + 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 = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + assert!( + GeoArea + .return_dtype(&EmptyOptions, &vec![dtype; arity]) + .is_err() + ); + Ok(()) + } + + #[test] + fn rejects_non_geometry_dtype() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(GeoArea.return_dtype(&EmptyOptions, &[primitive]).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/collect.rs b/vortex-geo/src/scalar_fn/collect.rs new file mode 100644 index 00000000000..c1c6aa070d0 --- /dev/null +++ b/vortex-geo/src/scalar_fn/collect.rs @@ -0,0 +1,508 @@ +// 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, + "geo: collect requires exactly one list operand, got {}", + dtypes.len() + ); + let DType::List(element_dtype, nullability) = &dtypes[0] else { + vortex_bail!("geo: collect operand {} is not a list", dtypes[0]); + }; + let Some(element) = element_dtype.as_extension_opt() else { + vortex_bail!( + "geo: collect list element {} is not a native Point, LineString, or Polygon", + element_dtype + ); + }; + // Multi-geometries cannot contain null components. Null list elements are ignored during + // execution, so their storage is non-nullable in the result. + let storage = DType::List( + Arc::new(element.storage_dtype().as_nonnullable()), + *nullability, + ); + let output = if element.is::() { + ExtDType::::try_new(element.metadata::().clone(), storage)?.erased() + } else if element.is::() { + ExtDType::::try_new(element.metadata::().clone(), storage)? + .erased() + } else if element.is::() { + ExtDType::::try_new(element.metadata::().clone(), storage)?.erased() + } else { + vortex_bail!( + "geo: collect list element {} is not a native Point, LineString, or Polygon", + element_dtype + ); + }; + Ok(DType::Extension(output)) +} + +/// 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. +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)?; + } + + 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!("geo: collect element offset exceeds usize"))?; + let size = usize::try_from(old_size) + .map_err(|_| vortex_err!("geo: collect element count exceeds usize"))?; + let end = start + .checked_add(size) + .ok_or_else(|| vortex_err!("geo: collect element range overflows usize"))?; + vortex_ensure!( + end <= element_valid.len(), + "geo: 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!("geo: 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!("geo: 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)?.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), output_dtype.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, output_dtype.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 GeoCollect; + +impl GeoCollect { + /// 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(GeoCollect, EmptyOptions).erased(), + vec![array], + ) + } +} + +impl ScalarFnVTable for GeoCollect { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.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 { + 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()))?; + let output = output_dtype.as_extension().clone(); + dispatch_unary( + &input, + output_dtype, + |execution, ctx| execute_collect(execution, &output, 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 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::GeoCollect; + 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 = GeoCollect::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 = GeoCollect::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(()) + } + + #[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 = GeoCollect::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 = GeoCollect::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 = GeoCollect::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 = GeoCollect::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 = GeoCollect::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 = GeoCollect::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!(GeoCollect::try_new_array(point).is_err()); + + let multipoints = multipoint_column(vec![vec![(0.0, 0.0)]])?; + assert!(GeoCollect::try_new_array(list(multipoints, &[0, 1])?).is_err()); + + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!( + GeoCollect + .return_dtype( + &EmptyOptions, + &[DType::List(primitive.into(), Nullability::NonNullable)] + ) + .is_err() + ); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/convex_hull.rs b/vortex-geo/src/scalar_fn/convex_hull.rs new file mode 100644 index 00000000000..653fb8a3fe7 --- /dev/null +++ b/vortex-geo/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, + "geo: convex_hull requires exactly one MultiPoint operand, got {}", + dtypes.len() + ); + let Some(input) = dtypes[0].as_extension_opt() else { + vortex_bail!( + "geo: convex_hull operand {} is not a native MultiPoint", + dtypes[0] + ); + }; + vortex_ensure!( + input.is::(), + "geo: 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 GeoConvexHull; + +impl GeoConvexHull { + /// 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(GeoConvexHull, EmptyOptions).erased(), + vec![array], + ) + } +} + +impl ScalarFnVTable for GeoConvexHull { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.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::GeoConvexHull; + use crate::scalar_fn::area::GeoArea; + use crate::scalar_fn::collect::GeoCollect; + 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 = GeoConvexHull::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 = GeoConvexHull::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 = GeoConvexHull::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 = GeoConvexHull::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 = GeoCollect::try_new_array(point_lists)?.into_array(); + let hulls = GeoConvexHull::try_new_array(collected)?.into_array(); + let areas = GeoArea::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!( + GeoConvexHull + .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!(GeoConvexHull::try_new_array(input).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/execute.rs b/vortex-geo/src/scalar_fn/execute.rs index ca5b4018249..3a7494bcb39 100644 --- a/vortex-geo/src/scalar_fn/execute.rs +++ b/vortex-geo/src/scalar_fn/execute.rs @@ -3,20 +3,23 @@ //! Shared execution for native geometry scalar functions. //! -//! [`dispatch_unary`] and the binary dispatcher handle constant/column operands and strict null +//! [`dispatch_unary`] and [`dispatch_binary`] handle constant/column operands and strict null //! propagation without prescribing how a kernel represents geometries or builds its output. -//! Native columnar kernels such as `ST_Envelope` use the unary dispatcher directly. +//! Native columnar kernels such as `ST_MakeLine` use these dispatchers directly. //! -//! [`execute_binary_geo_types`] adapts row-oriented algorithms from the `geo` ecosystem. It decodes -//! valid inputs into `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`], such -//! as an `f64` or boolean array. +//! [`execute_unary_geo_types`] and [`execute_binary_geo_types`] are convenience adapters for +//! row-oriented algorithms from the `geo` ecosystem. They decode valid inputs into +//! `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`], such as an `f64` or +//! boolean array. mod binary; mod geo_types; mod unary; +pub(crate) use binary::dispatch_binary; pub(crate) use binary::execute_binary_geo_types; pub(crate) use unary::dispatch_unary; +pub(crate) use unary::execute_unary_geo_types; use vortex_array::ArrayRef; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; diff --git a/vortex-geo/src/scalar_fn/execute/unary.rs b/vortex-geo/src/scalar_fn/execute/unary.rs index bdbbd0b33ac..fd5c3abe179 100644 --- a/vortex-geo/src/scalar_fn/execute/unary.rs +++ b/vortex-geo/src/scalar_fn/execute/unary.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Unary operand dispatch for native geometry kernels. +//! Unary operand dispatch, plus an adapter for row-oriented `geo_types` kernels. +use geo_types::Geometry; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -15,6 +16,9 @@ use vortex_error::VortexResult; use super::Execution; use super::Operand; +use super::geo_types::GeoTypesOutput; +use super::geo_types::eval_column; +use crate::extension::single_geometry; /// Dispatch a unary strict geometry kernel over a constant or column. /// @@ -60,3 +64,46 @@ where ctx, ) } + +/// Run a unary row-oriented kernel whose input is decoded to `geo_types::Geometry`. +/// +/// The `geo_types` name describes the value passed to `compute`, not the output. `T` is converted +/// into a Vortex array before this function returns. A constant is decoded and computed once +/// before broadcast; a column is decoded only for its valid rows. +pub(crate) fn execute_unary_geo_types( + array: &ArrayRef, + compute: F, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: GeoTypesOutput, + F: Fn(&Geometry) -> T, +{ + let nullability = array.dtype().nullability(); + dispatch_unary( + array, + T::dtype(nullability), + |execution, ctx| match execution.operands { + [Operand::Constant(scalar)] => { + let geometry = single_geometry(&scalar, ctx)?; + Ok(ConstantArray::new( + compute(&geometry).into_scalar(execution.nullability), + execution.len, + ) + .into_array()) + } + [Operand::Column(array)] => { + let valid = execution.valid.execute_mask(execution.len, ctx)?; + if execution.len != 0 && valid.all_false() { + return Ok(ConstantArray::new( + Scalar::null(T::dtype(execution.nullability)), + execution.len, + ) + .into_array()); + } + eval_column(&array, &valid, compute, execution.nullability, ctx) + } + }, + ctx, + ) +} diff --git a/vortex-geo/src/scalar_fn/hilbert.rs b/vortex-geo/src/scalar_fn/hilbert.rs new file mode 100644 index 00000000000..69881e45628 --- /dev/null +++ b/vortex-geo/src/scalar_fn/hilbert.rs @@ -0,0 +1,579 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `hilbert`: a locality-preserving `u32` key for spatially clustering a native geometry column. +//! +//! The caller supplies one dataset-wide [`Rect`](crate::extension::Rect) bound. Each geometry is +//! represented by the center of its XY envelope, quantized to 16 bits per axis within that bound, +//! and encoded on a Hilbert curve. The scalar function only computes keys; sorting the complete +//! row set and writing it in key order is an ingestion-layer responsibility. + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::scalar::Scalar; +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_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::Rect; +use crate::extension::coordinate::ordinates; +use crate::extension::is_native_geometry; +use crate::scalar_fn::envelope::GeoEnvelope; +use crate::scalar_fn::execute::Execution; +use crate::scalar_fn::execute::Operand; +use crate::scalar_fn::execute::dispatch_unary; + +const AXIS_MAX: f64 = u16::MAX as f64; + +#[expect( + clippy::cast_possible_truncation, + reason = "Hilbert quantization intentionally truncates onto a 16-bit grid" +)] +fn quantize(value: f64, min: f64, max: f64) -> Option { + if !value.is_finite() { + return None; + } + if min == max { + return Some(0); + } + let normalized = (value - min) / (max - min); + normalized + .is_finite() + .then(|| (normalized * AXIS_MAX).clamp(0.0, AXIS_MAX) as u32) +} + +fn hilbert_key([xmin, ymin, xmax, ymax]: [f64; 4], x: f64, y: f64) -> Option { + Some(hilbert_encode_16( + quantize(x, xmin, xmax)?, + quantize(y, ymin, ymax)?, + )) +} + +/// Return the midpoint without overflowing when finite endpoints have opposite signs or are near +/// the edge of the `f64` range. +#[inline] +fn midpoint(min: f64, max: f64) -> f64 { + min / 2.0 + max / 2.0 +} + +#[inline] +fn envelope_center([xmin, ymin, xmax, ymax]: [f64; 4]) -> (f64, f64) { + (midpoint(xmin, xmax), midpoint(ymin, ymax)) +} + +/// Interleave the low 16 bits of `value` with zero bits. +#[inline] +fn hilbert_interleave(mut value: u32) -> u32 { + value = (value | (value << 8)) & 0x00ff_00ff; + value = (value | (value << 4)) & 0x0f0f_0f0f; + value = (value | (value << 2)) & 0x3333_3333; + (value | (value << 1)) & 0x5555_5555 +} + +/// Encode a 16-bit-per-axis point as a 32-bit Hilbert index. +/// +/// This is the public-domain prefix-scan algorithm from +/// . +#[inline] +fn hilbert_encode_16(x: u32, y: u32) -> u32 { + debug_assert!(x <= u32::from(u16::MAX)); + debug_assert!(y <= u32::from(u16::MAX)); + + let input_x = x; + let input_y = y; + let mut state_a = x ^ y; + let mut state_b = 0xffff ^ state_a; + let mut state_c = 0xffff ^ (x | y); + let mut state_d = x & (y ^ 0xffff); + let mut next_a = state_a | (state_b >> 1); + let mut next_b = (state_a >> 1) ^ state_a; + let mut next_c = ((state_c >> 1) ^ (state_b & (state_d >> 1))) ^ state_c; + let mut next_d = ((state_a & (state_c >> 1)) ^ (state_d >> 1)) ^ state_d; + + state_a = next_a; + state_b = next_b; + state_c = next_c; + state_d = next_d; + next_a = (state_a & (state_a >> 2)) ^ (state_b & (state_b >> 2)); + next_b = (state_a & (state_b >> 2)) ^ (state_b & ((state_a ^ state_b) >> 2)); + next_c ^= (state_a & (state_c >> 2)) ^ (state_b & (state_d >> 2)); + next_d ^= (state_b & (state_c >> 2)) ^ ((state_a ^ state_b) & (state_d >> 2)); + + state_a = next_a; + state_b = next_b; + state_c = next_c; + state_d = next_d; + next_a = (state_a & (state_a >> 4)) ^ (state_b & (state_b >> 4)); + next_b = (state_a & (state_b >> 4)) ^ (state_b & ((state_a ^ state_b) >> 4)); + next_c ^= (state_a & (state_c >> 4)) ^ (state_b & (state_d >> 4)); + next_d ^= (state_b & (state_c >> 4)) ^ ((state_a ^ state_b) & (state_d >> 4)); + + state_a = next_a; + state_b = next_b; + state_c = next_c; + state_d = next_d; + next_c ^= (state_a & (state_c >> 8)) ^ (state_b & (state_d >> 8)); + next_d ^= (state_b & (state_c >> 8)) ^ ((state_a ^ state_b) & (state_d >> 8)); + + let state_a = next_c ^ (next_c >> 1); + let state_b = next_d ^ (next_d >> 1); + let i0 = input_x ^ input_y; + let i1 = state_b | (0xffff ^ (i0 | state_a)); + + (hilbert_interleave(i1) << 1) | hilbert_interleave(i0) +} + +/// Read the four coordinates from the constant [`Rect`] bounds. +fn rect_bounds(scalar: &Scalar) -> VortexResult<[f64; 4]> { + let storage = scalar.as_extension().to_storage_scalar(); + let fields = storage.as_struct(); + let read = |name: &str| -> VortexResult { + f64::try_from( + &fields + .field(name) + .ok_or_else(|| vortex_err!("geo: hilbert bounds missing {name}"))?, + ) + }; + let bounds = [read("xmin")?, read("ymin")?, read("xmax")?, read("ymax")?]; + let [xmin, ymin, xmax, ymax] = bounds; + vortex_ensure!( + bounds.into_iter().all(f64::is_finite), + "geo: hilbert bounds must be finite" + ); + vortex_ensure!( + xmin <= xmax && ymin <= ymax, + "geo: hilbert bounds must satisfy xmin <= xmax and ymin <= ymax" + ); + Ok(bounds) +} + +fn geometry_keys( + geometry: ArrayRef, + bounds: [f64; 4], + ctx: &mut ExecutionCtx, +) -> VortexResult { + let envelopes = GeoEnvelope::try_new_array(geometry)? + .into_array() + .execute::(ctx)?; + let storage = envelopes.storage_array().clone(); + let valid = storage.validity()?.execute_mask(storage.len(), ctx)?; + let boxes = storage.execute::(ctx)?; + let xmins = ordinates(&boxes, "xmin", ctx)?; + let ymins = ordinates(&boxes, "ymin", ctx)?; + let xmaxs = ordinates(&boxes, "xmax", ctx)?; + let ymaxs = ordinates(&boxes, "ymax", ctx)?; + let mut keys = BufferMut::zeroed(xmins.len()); + let mut key_valid = vec![false; xmins.len()]; + for row in 0..keys.len() { + let (x, y) = envelope_center([xmins[row], ymins[row], xmaxs[row], ymaxs[row]]); + if let Some(key) = hilbert_key(bounds, x, y) { + keys[row] = key; + key_valid[row] = true; + } + } + let key_valid = Mask::from(BitBuffer::from_iter(key_valid)); + let validity = Validity::from_mask(&valid & &key_valid, Nullability::Nullable); + Ok(PrimitiveArray::new(keys.freeze(), validity).into_array()) +} + +fn execute_hilbert( + execution: Execution<1, Validity>, + bounds: [f64; 4], + ctx: &mut ExecutionCtx, +) -> VortexResult { + match execution.operands { + [Operand::Constant(geometry)] => { + let one = ConstantArray::new(geometry, 1).into_array(); + let key = geometry_keys(one, bounds, ctx)?.execute_scalar(0, ctx)?; + Ok(ConstantArray::new(key, execution.len).into_array()) + } + [Operand::Column(geometry)] => geometry_keys(geometry, bounds, ctx), + } +} + +fn validate_hilbert_operands(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure!( + dtypes.len() == 2, + "geo: hilbert requires a geometry and one whole-column bound, got {} operands", + dtypes.len() + ); + vortex_ensure!( + is_native_geometry(&dtypes[0]), + "geo: hilbert operand {} is not a native geometry type", + dtypes[0] + ); + vortex_ensure!( + dtypes[1] + .as_extension_opt() + .is_some_and(|ext| ext.is::()), + "geo: hilbert bounds {} are not a geometry box", + dtypes[1] + ); + Ok(()) +} + +/// A locality-preserving `u32` key for the center of each geometry's XY envelope. +/// +/// One common bounds scalar must cover the complete column being clustered. Null, empty, or +/// non-finite geometries yield null keys. This function computes keys only; it does not sort rows. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct GeoHilbert; + +impl GeoHilbert { + /// A lazy Hilbert-key column using one dataset-wide `bounds` scalar for every geometry row. + pub fn try_new_array(geometry: ArrayRef, bounds: Scalar) -> VortexResult { + let len = geometry.len(); + ScalarFnArray::try_new( + TypedScalarFnInstance::new(GeoHilbert, EmptyOptions).erased(), + vec![geometry, ConstantArray::new(bounds, len).into_array()], + ) + } +} + +impl ScalarFnVTable for GeoHilbert { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.hilbert"); + *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("geometry"), + 1 => ChildName::from("bounds"), + _ => unreachable!("hilbert has exactly two children"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_hilbert_operands(dtypes)?; + Ok(DType::Primitive(PType::U32, Nullability::Nullable)) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let geometry = args.get(0)?; + let bounds = args.get(1)?; + let bounds = bounds.as_::(); + if bounds.scalar().is_null() { + return Ok(ConstantArray::new( + Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)), + geometry.len(), + ) + .into_array()); + } + let bounds = rect_bounds(bounds.scalar())?; + dispatch_unary( + &geometry, + DType::Primitive(PType::U32, Nullability::Nullable), + |execution, ctx| execute_hilbert(execution, bounds, ctx), + ctx, + ) + } + + fn validity(&self, _: &Self::Options, _: &Expression) -> VortexResult> { + // Empty and non-finite geometries yield null even when both inputs are valid. + Ok(None) + } + + fn is_strict(&self, _: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _: &Self::Options) -> bool { + true + } +} + +#[cfg(test)] +mod tests { + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::aggregate_fn::EmptyOptions as AggregateEmptyOptions; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + + use super::GeoHilbert; + use super::hilbert_encode_16; + use crate::aggregate_fn::GeometryAabb; + 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::point_column; + use crate::test_harness::rect_column; + + fn bounds( + corners: (f64, f64, f64, f64), + ctx: &mut vortex_array::ExecutionCtx, + ) -> VortexResult { + rect_column(vec![corners])?.execute_scalar(0, ctx) + } + + fn keys(geometry: ArrayRef, bounds: Scalar) -> VortexResult { + Ok(GeoHilbert::try_new_array(geometry, bounds)?.into_array()) + } + + /// Reference values lock down scaling, truncation, orientation, and bit ordering. + #[test] + fn matches_reference_vector() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + let points = point_column(vec![0.25, 0.50, 0.75], vec![0.25, 0.50, 0.75])?; + let expected = PrimitiveArray::new( + vec![178_956_970_u32, 715_827_882, 2_326_440_618], + Validity::from_iter([true, true, true]), + ) + .into_array(); + + assert_arrays_eq!( + keys(points, bounds((0.0, 0.0, 1.0, 1.0), &mut ctx)?)?, + expected, + &mut ctx + ); + Ok(()) + } + + /// A point and geometries whose envelopes have the same center receive the same key under one + /// common domain. + #[test] + fn keys_the_envelope_center() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + let bounds = bounds((0.0, 0.0, 10.0, 10.0), &mut ctx)?; + let point = point_column(vec![2.0], vec![3.0])?; + let multipoint = multipoint_column(vec![vec![(0.0, 1.0), (4.0, 5.0)]])?; + let rect = rect_column(vec![(0.0, 1.0, 4.0, 5.0)])?; + let expected = PrimitiveArray::new( + vec![hilbert_encode_16(13_107, 19_660)], + Validity::from_iter([true]), + ) + .into_array(); + + for geometry in [point, multipoint, rect] { + assert_arrays_eq!(keys(geometry, bounds.clone())?, expected.clone(), &mut ctx); + } + Ok(()) + } + + /// The existing whole-column AABB aggregate produces exactly the scalar accepted by Hilbert. + #[test] + fn consumes_geometry_aabb_result() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + let points = point_column(vec![0.0, 0.5, 1.0], vec![0.0, 0.5, 1.0])?; + let mut aabb = + Accumulator::try_new(GeometryAabb, AggregateEmptyOptions, points.dtype().clone())?; + aabb.accumulate(&points, &mut ctx)?; + + let expected = PrimitiveArray::new( + vec![0_u32, 715_827_882, 2_863_311_530], + Validity::from_iter([true, true, true]), + ) + .into_array(); + assert_arrays_eq!(keys(points, aabb.finish()?)?, expected, &mut ctx); + Ok(()) + } + + #[test] + fn constant_geometry_is_broadcast() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + let point = point_column(vec![1.0], vec![1.0])?.execute_scalar(0, &mut ctx)?; + let geometries = ConstantArray::new(point, 3).into_array(); + let expected = PrimitiveArray::new( + vec![715_827_882_u32; 3], + Validity::from_iter([true, true, true]), + ) + .into_array(); + + assert_arrays_eq!( + keys(geometries, bounds((0.0, 0.0, 2.0, 2.0), &mut ctx)?)?, + expected, + &mut ctx + ); + Ok(()) + } + + /// Null and empty geometries retain their row positions as null keys. + #[test] + fn empty_and_null_geometries_are_null() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + let geometries = nullable_multipolygon_column(vec![ + Some(vec![vec![vec![(0.0, 0.0), (2.0, 2.0)]]]), + Some(vec![]), + None, + ])?; + let expected = PrimitiveArray::new( + vec![715_827_882_u32, 0, 0], + Validity::from_iter([true, false, false]), + ) + .into_array(); + + assert_arrays_eq!( + keys(geometries, bounds((0.0, 0.0, 2.0, 2.0), &mut ctx)?)?, + expected, + &mut ctx + ); + Ok(()) + } + + /// Non-finite point coordinates do not become arbitrary sortable keys. + #[test] + fn non_finite_geometry_is_null() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + let points = point_column(vec![f64::NAN, 1.0], vec![f64::NAN, 1.0])?; + let expected = + PrimitiveArray::new(vec![0_u32, 715_827_882], Validity::from_iter([false, true])) + .into_array(); + + assert_arrays_eq!( + keys(points, bounds((0.0, 0.0, 2.0, 2.0), &mut ctx)?)?, + expected, + &mut ctx + ); + Ok(()) + } + + /// A zero-width dataset axis is valid: it contributes no ordering information while the + /// other axis remains sortable. + #[test] + fn degenerate_axis_is_supported() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + let points = point_column(vec![1.0, 1.0], vec![0.0, 2.0])?; + let expected = PrimitiveArray::new( + vec![0_u32, 1_431_655_765], + Validity::from_iter([true, true]), + ) + .into_array(); + + assert_arrays_eq!( + keys(points, bounds((1.0, 0.0, 1.0, 2.0), &mut ctx)?)?, + expected, + &mut ctx + ); + Ok(()) + } + + /// Slicing still attributes nested coordinates to the correct outer geometry row. + #[test] + fn sliced_nested_geometry_keeps_rows_aligned() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + let geometries = multipolygon_column(vec![ + vec![vec![vec![(-100.0, -100.0), (100.0, 100.0)]]], + vec![vec![vec![(0.0, 0.0), (2.0, 2.0)]]], + vec![vec![vec![(4.0, 4.0), (6.0, 6.0)]]], + ])?; + let expected = PrimitiveArray::new( + vec![ + hilbert_encode_16(8_191, 8_191), + hilbert_encode_16(40_959, 40_959), + ], + Validity::from_iter([true, true]), + ) + .into_array(); + + assert_arrays_eq!( + keys( + geometries.slice(1..3)?, + bounds((0.0, 0.0, 8.0, 8.0), &mut ctx)? + )?, + expected, + &mut ctx + ); + Ok(()) + } + + /// Planning rejects non-geometry operands and non-box bounds before execution. + #[test] + fn validates_operand_types() -> VortexResult<()> { + let numeric = DType::Primitive(PType::I32, Nullability::NonNullable); + let geometry = point_column(vec![0.0], vec![0.0])?.dtype().clone(); + assert!( + GeoHilbert + .return_dtype(&EmptyOptions, &[numeric.clone(), geometry.clone()]) + .is_err() + ); + assert!( + GeoHilbert + .return_dtype(&EmptyOptions, &[geometry, numeric]) + .is_err() + ); + Ok(()) + } + + /// Execution materializes a canonical nullable `u32` primitive array. + #[test] + fn result_is_nullable_u32() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + let result = keys( + nullable_point_column(vec![Some((1.0, 1.0)), None])?, + bounds((0.0, 0.0, 2.0, 2.0), &mut ctx)?, + )? + .execute::(&mut ctx)? + .into_primitive(); + assert_eq!(result.ptype(), PType::U32); + assert!(result.dtype().is_nullable()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/intersection.rs b/vortex-geo/src/scalar_fn/intersection.rs new file mode 100644 index 00000000000..b23336c524e --- /dev/null +++ b/vortex-geo/src/scalar_fn/intersection.rs @@ -0,0 +1,459 @@ +// 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::GeoMetadata; +use crate::extension::MultiPolygon; +use crate::extension::Polygon; +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: &GeoMetadata, right: &GeoMetadata) -> VortexResult { + match (&left.crs, &right.crs) { + (Some(left_crs), Some(right_crs)) => { + vortex_ensure!( + left_crs == right_crs, + "geo: 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(GeoMetadata::default()), + } +} + +/// Metadata carried by a validated native polygonal dtype. +fn polygonal_metadata(dtype: &DType) -> &GeoMetadata { + 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, + "geo: 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::() + }), + "geo: 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 GeoIntersection; + +impl GeoIntersection { + /// 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(GeoIntersection, EmptyOptions).erased(), + vec![left, right], + ) + } +} + +impl ScalarFnVTable for GeoIntersection { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.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::GeoIntersection; + use crate::extension::MultiPolygon; + use crate::extension::geometries; + use crate::scalar_fn::area::GeoArea; + 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 = GeoIntersection::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 = GeoArea::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 = GeoIntersection::try_new_array(left, right)?.into_array(); + let areas = GeoArea::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 = GeoIntersection::try_new_array(left, right)?.into_array(); + let areas = GeoArea::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 = GeoIntersection::try_new_array(left, right)?.into_array(); + let areas = GeoArea::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 = GeoIntersection::try_new_array(left, right)?.into_array(); + let areas = GeoArea::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 = GeoIntersection::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!( + GeoIntersection + .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!(GeoIntersection::try_new_array(polygon, point).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/length.rs b/vortex-geo/src/scalar_fn/length.rs new file mode 100644 index 00000000000..38f7e639bea --- /dev/null +++ b/vortex-geo/src/scalar_fn/length.rs @@ -0,0 +1,254 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Length`: planar (Euclidean) length of native line strings. + +use geo::Euclidean; +use geo::Length; +use geo_types::Geometry; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +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_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::scalar_fn::execute::execute_unary_geo_types; + +/// Validate the native line-string operand accepted by `ST_Length`. +fn validate_length_operands(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure!( + dtypes.len() == 1, + "geo: length requires exactly one line string operand, got {}", + dtypes.len() + ); + vortex_ensure!( + dtypes[0] + .as_extension_opt() + .is_some_and(|extension| extension.is::()), + "geo: length operand {} is not a native line string", + dtypes[0] + ); + Ok(()) +} + +/// Planar (Euclidean) `ST_Length` (no geodesic correction) of native line strings. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct GeoLength; + +impl GeoLength { + /// A lazy `ScalarFnArray` computing the per-row length of a line string operand. + pub fn try_new_array(array: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(GeoLength, EmptyOptions).erased(), + vec![array], + ) + } +} + +impl ScalarFnVTable for GeoLength { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.length"); + *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!("length has exactly one child"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + validate_length_operands(dtypes)?; + Ok(DType::Primitive(PType::F64, dtypes[0].nullability())) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let array = args.get(0)?; + execute_unary_geo_types( + &array, + |geometry| match geometry { + Geometry::LineString(line) => Euclidean.length(line), + _ => unreachable!("length input is validated as a line string"), + }, + 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 vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::Columnar; + use vortex_array::ExecutionCtx; + 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::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::GeoLength; + use crate::test_harness::linestring_column; + use crate::test_harness::point_column; + + fn line_constant( + line: Vec<(f64, f64)>, + len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let scalar = linestring_column(vec![line])?.execute_scalar(0, ctx)?; + Ok(ConstantArray::new(scalar, len).into_array()) + } + + #[test] + fn measures_each_linestring() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let lines = linestring_column(vec![ + vec![(0.0, 0.0), (3.0, 4.0)], + vec![(0.0, 0.0), (3.0, 4.0), (3.0, 8.0)], + vec![], + ])?; + + let lengths = GeoLength::try_new_array(lines)? + .into_array() + .execute::(&mut ctx)? + .into_primitive(); + + assert_eq!(lengths.as_slice::(), &[5.0, 9.0, 0.0]); + Ok(()) + } + + #[test] + fn constant_is_computed_once_and_remains_constant() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let lines = line_constant(vec![(0.0, 0.0), (3.0, 4.0)], 3, &mut ctx)?; + + let result = GeoLength::try_new_array(lines)? + .into_array() + .execute::(&mut ctx)?; + let Columnar::Constant(lengths) = result else { + return Err(vortex_err!("length of a constant should remain constant")); + }; + assert_eq!(lengths.len(), 3); + assert_eq!(f64::try_from(lengths.scalar())?, 5.0); + Ok(()) + } + + #[test] + fn null_constant_is_all_null() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let dtype = linestring_column(vec![vec![]])?.dtype().as_nullable(); + let lines = ConstantArray::new(Scalar::null(dtype), 2).into_array(); + + let result = GeoLength::try_new_array(lines)? + .into_array() + .execute::(&mut ctx)?; + let Columnar::Constant(lengths) = result else { + return Err(vortex_err!( + "length of a null constant should remain constant" + )); + }; + assert_eq!(lengths.len(), 2); + assert!(lengths.scalar().is_null()); + Ok(()) + } + + #[test] + fn nullable_rows_propagate_without_decoding_null_geometries() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let lines = MaskedArray::try_new( + linestring_column(vec![ + vec![(0.0, 0.0), (3.0, 4.0)], + vec![(0.0, 0.0), (1.0, 1.0)], + vec![(0.0, 0.0), (0.0, 4.0)], + ])?, + Validity::from_iter([true, false, true]), + )? + .into_array(); + + let expected = PrimitiveArray::new( + vec![5.0, 0.0, 4.0], + Validity::from_iter([true, false, true]), + ) + .into_array(); + let lengths = GeoLength::try_new_array(lines)?.into_array(); + assert_arrays_eq!(lengths, expected, &mut ctx); + Ok(()) + } + + #[test] + fn rejects_non_linestring_dtype() -> VortexResult<()> { + let point = point_column(vec![0.0], vec![0.0])?; + assert!( + GeoLength + .return_dtype(&EmptyOptions, std::slice::from_ref(point.dtype())) + .is_err() + ); + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(GeoLength.return_dtype(&EmptyOptions, &[primitive]).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/make_line.rs b/vortex-geo/src/scalar_fn/make_line.rs new file mode 100644 index 00000000000..e48569c3d50 --- /dev/null +++ b/vortex-geo/src/scalar_fn/make_line.rs @@ -0,0 +1,540 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_MakeLine`: construct a native line string between two native points. + +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_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::GeoMetadata; +use crate::extension::LineString; +use crate::extension::Point; +use crate::extension::coordinate::coordinate_dimension; +use crate::extension::flatten_coordinates; +use crate::extension::linestring_array_from_point_pairs; +use crate::extension::linestring_storage_dtype; +use crate::scalar_fn::execute::Execution; +use crate::scalar_fn::execute::Operand; +use crate::scalar_fn::execute::dispatch_binary; + +/// Validate the two point operands accepted by `ST_MakeLine`. +fn validate_make_line_operands(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure!( + dtypes.len() == 2, + "geo: make_line requires exactly two point operands, got {}", + dtypes.len() + ); + for dtype in dtypes { + vortex_ensure!( + dtype + .as_extension_opt() + .is_some_and(|extension| extension.is::()), + "geo: make_line operand {dtype} is not a native point" + ); + } + Ok(()) +} + +/// Resolve DuckDB's `ST_MakeLine` CRS propagation for two geometry operands. +fn make_line_metadata(left: &GeoMetadata, right: &GeoMetadata) -> VortexResult { + match (&left.crs, &right.crs) { + (Some(left_crs), Some(right_crs)) => { + vortex_ensure!( + left_crs == right_crs, + "geo: make_line 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(GeoMetadata::default()), + } +} + +/// The native `LineString` dtype emitted by `ST_MakeLine`. +fn make_line_dtype(dtypes: &[DType]) -> VortexResult> { + validate_make_line_operands(dtypes)?; + let left = dtypes[0].as_extension(); + let right = dtypes[1].as_extension(); + let dimension = coordinate_dimension(left.storage_dtype())? + .promote(coordinate_dimension(right.storage_dtype())?); + let metadata = make_line_metadata(left.metadata::(), right.metadata::())?; + let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); + ExtDType::try_new(metadata, linestring_storage_dtype(dimension, nullability)) +} + +/// Build a native line-string column from dispatched point operands. +fn build_make_lines( + operands: [Operand; 2], + len: usize, + valid: Mask, + output_dtype: &ExtDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let [start, end] = operands.map(|operand| match operand { + Operand::Constant(point) => ConstantArray::new(point, len).into_array(), + Operand::Column(points) => points, + }); + let starts = flatten_coordinates(&start, ctx)?; + let ends = flatten_coordinates(&end, ctx)?; + linestring_array_from_point_pairs( + output_dtype, + &starts, + &ends, + Validity::from_mask(valid, output_dtype.storage_dtype().nullability()), + ) +} + +/// Execute `ST_MakeLine` after shared constant/column and null dispatch. +fn execute_make_line( + execution: Execution<2>, + output_dtype: &ExtDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match execution.operands { + [Operand::Constant(start), Operand::Constant(end)] => { + let one = build_make_lines( + [Operand::Constant(start), Operand::Constant(end)], + 1, + Mask::new_true(1), + output_dtype, + ctx, + )?; + Ok(ConstantArray::new(one.execute_scalar(0, ctx)?, execution.len).into_array()) + } + operands => build_make_lines(operands, execution.len, execution.valid, output_dtype, ctx), + } +} + +/// Construct `LineString`s from paired native point operands. The output's vertices preserve all +/// coordinate ordinates (`x`, `y`, and any `z`/`m`) and appear in operand order. When the points +/// have different dimensions, it promotes to their union and fills absent `z`/`m` ordinates with +/// zero. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct GeoMakeLine; + +impl GeoMakeLine { + /// A lazy `ScalarFnArray` constructing one two-vertex line string per pair of point operands. + pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(GeoMakeLine, EmptyOptions).erased(), + vec![a, b], + ) + } +} + +impl ScalarFnVTable for GeoMakeLine { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.geo.make_line"); + *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("start"), + 1 => ChildName::from("end"), + _ => unreachable!("make_line has exactly two children"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + Ok(DType::Extension(make_line_dtype(dtypes)?.erased())) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let a = args.get(0)?; + let b = args.get(1)?; + let output_dtype = make_line_dtype(&[a.dtype().clone(), b.dtype().clone()])?; + dispatch_binary( + &a, + &b, + DType::Extension(output_dtype.clone().erased()), + |execution, ctx| execute_make_line(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::Coord; + use geo_types::Geometry; + use geo_types::LineString as GeoLineString; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Columnar; + use vortex_array::ExecutionCtx; + 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::extension::ExtDType; + use vortex_array::scalar::Scalar; + 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::GeoMakeLine; + use crate::extension::GeoMetadata; + use crate::extension::LineString; + use crate::extension::Point; + use crate::extension::coordinate::Dimension; + use crate::extension::coordinate::coordinate_dimension; + use crate::extension::coordinate::ordinates; + use crate::extension::flatten_coordinates; + use crate::extension::geometries; + use crate::scalar_fn::length::GeoLength; + use crate::test_harness::nullable_point_column; + use crate::test_harness::point_column; + + fn dimensional_point( + dimension: Dimension, + coordinate: [f64; 4], + crs: Option<&str>, + ) -> VortexResult { + let mut fields = vec![ + ("x", PrimitiveArray::from_iter([coordinate[0]]).into_array()), + ("y", PrimitiveArray::from_iter([coordinate[1]]).into_array()), + ]; + if matches!(dimension, Dimension::Xyz | Dimension::Xyzm) { + fields.push(("z", PrimitiveArray::from_iter([coordinate[2]]).into_array())); + } + if matches!(dimension, Dimension::Xym | Dimension::Xyzm) { + fields.push(("m", PrimitiveArray::from_iter([coordinate[3]]).into_array())); + } + let storage = StructArray::from_fields(&fields)?.into_array(); + let dtype = ExtDType::::try_new( + GeoMetadata { + crs: crs.map(str::to_owned), + }, + storage.dtype().clone(), + )?; + Ok(ExtensionArray::try_new(dtype.erased(), storage)?.into_array()) + } + + fn point_constant( + x: f64, + y: f64, + len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let scalar = point_column(vec![x], vec![y])?.execute_scalar(0, ctx)?; + Ok(ConstantArray::new(scalar, len).into_array()) + } + + #[test] + fn connects_paired_points_in_operand_order() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let starts = point_column(vec![0.0, 3.0], vec![0.0, 4.0])?; + let ends = point_column(vec![3.0, 0.0], vec![4.0, 0.0])?; + + let lines = GeoMakeLine::try_new_array(starts, ends)?.into_array(); + assert!(lines.dtype().as_extension().is::()); + assert_eq!( + geometries(&lines, &mut ctx)?, + vec![ + Geometry::LineString(GeoLineString::new(vec![ + Coord { x: 0.0, y: 0.0 }, + Coord { x: 3.0, y: 4.0 }, + ])), + Geometry::LineString(GeoLineString::new(vec![ + Coord { x: 3.0, y: 4.0 }, + Coord { x: 0.0, y: 0.0 }, + ])), + ] + ); + Ok(()) + } + + #[test] + fn two_constants_are_built_once_and_remain_constant() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let starts = point_constant(0.0, 0.0, 3, &mut ctx)?; + let ends = point_constant(3.0, 4.0, 3, &mut ctx)?; + + let result = GeoMakeLine::try_new_array(starts, ends)? + .into_array() + .execute::(&mut ctx)?; + let Columnar::Constant(lines) = result else { + return Err(vortex_err!( + "make_line of two constants should remain constant" + )); + }; + assert_eq!(lines.len(), 3); + assert_eq!( + geometries(&lines.into_array(), &mut ctx)?, + vec![ + Geometry::LineString(GeoLineString::new(vec![ + Coord { x: 0.0, y: 0.0 }, + Coord { x: 3.0, y: 4.0 }, + ])); + 3 + ] + ); + Ok(()) + } + + #[rstest] + #[case::constant_start(true)] + #[case::constant_end(false)] + fn constant_and_column_are_paired_by_row(#[case] constant_start: bool) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let constant = point_constant(0.0, 0.0, 2, &mut ctx)?; + let column = point_column(vec![3.0, 6.0], vec![4.0, 8.0])?; + let (starts, ends) = if constant_start { + (constant, column) + } else { + (column, constant) + }; + + let lines = GeoMakeLine::try_new_array(starts, ends)?.into_array(); + let endpoints = [(3.0, 4.0), (6.0, 8.0)]; + let expected = endpoints + .into_iter() + .map(|(x, y)| { + let constant = Coord { x: 0.0, y: 0.0 }; + let column = Coord { x, y }; + Geometry::LineString(GeoLineString::new(if constant_start { + vec![constant, column] + } else { + vec![column, constant] + })) + }) + .collect::>(); + assert_eq!(geometries(&lines, &mut ctx)?, expected); + Ok(()) + } + + #[test] + fn null_constant_is_all_null() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable(); + let starts = ConstantArray::new(Scalar::null(point_dtype), 2).into_array(); + let ends = point_column(vec![3.0, 6.0], vec![4.0, 8.0])?; + + let result = GeoMakeLine::try_new_array(starts, ends)? + .into_array() + .execute::(&mut ctx)?; + let Columnar::Constant(lines) = result else { + return Err(vortex_err!( + "make_line with a null constant should remain constant" + )); + }; + assert_eq!(lines.len(), 2); + assert!(lines.scalar().is_null()); + assert!(lines.dtype().as_extension().is::()); + Ok(()) + } + + /// A null endpoint produces a null line, which in turn produces a null length when the two + /// scalar functions are composed. + #[test] + fn propagates_endpoint_nulls_through_length() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let starts = nullable_point_column(vec![Some((0.0, 0.0)), None, Some((0.0, 0.0))])?; + let ends = nullable_point_column(vec![Some((3.0, 4.0)), Some((1.0, 1.0)), None])?; + + let lines = GeoMakeLine::try_new_array(starts, ends)?.into_array(); + let lengths = GeoLength::try_new_array(lines)?.into_array(); + let expected = PrimitiveArray::new( + vec![5.0f64, 0.0, 0.0], + Validity::from_iter([true, false, false]), + ) + .into_array(); + + assert_arrays_eq!(lengths, expected, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::xy_xyz( + Dimension::Xy, + [1.0, 2.0, 0.0, 0.0], + Dimension::Xyz, + [3.0, 4.0, 5.0, 0.0], + Dimension::Xyz, + Some([0.0, 5.0]), + None + )] + #[case::xyz_xy( + Dimension::Xyz, + [1.0, 2.0, 5.0, 0.0], + Dimension::Xy, + [3.0, 4.0, 0.0, 0.0], + Dimension::Xyz, + Some([5.0, 0.0]), + None + )] + #[case::xym_xyz( + Dimension::Xym, + [1.0, 2.0, 0.0, 6.0], + Dimension::Xyz, + [3.0, 4.0, 5.0, 0.0], + Dimension::Xyzm, + Some([0.0, 5.0]), + Some([6.0, 0.0]) + )] + #[case::xyz_xym( + Dimension::Xyz, + [1.0, 2.0, 5.0, 0.0], + Dimension::Xym, + [3.0, 4.0, 0.0, 6.0], + Dimension::Xyzm, + Some([5.0, 0.0]), + Some([0.0, 6.0]) + )] + fn promotes_mixed_point_dimensions( + #[case] start_dimension: Dimension, + #[case] start: [f64; 4], + #[case] end_dimension: Dimension, + #[case] end: [f64; 4], + #[case] expected_dimension: Dimension, + #[case] expected_z: Option<[f64; 2]>, + #[case] expected_m: Option<[f64; 2]>, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let start = dimensional_point(start_dimension, start, None)?; + let end = dimensional_point(end_dimension, end, None)?; + + let lines = GeoMakeLine::try_new_array(start, end)?.into_array(); + let vertices = flatten_coordinates(&lines, &mut ctx)?; + assert_eq!(coordinate_dimension(vertices.dtype())?, expected_dimension); + for (name, expected) in [("z", expected_z), ("m", expected_m)] { + if let Some(expected) = expected { + assert_eq!( + ordinates(&vertices, name, &mut ctx)? + .iter() + .copied() + .collect::>(), + expected + ); + } + } + Ok(()) + } + + #[rstest] + #[case::matching(Some("EPSG:4326"), Some("EPSG:4326"), Some("EPSG:4326"))] + #[case::left_only(Some("EPSG:4326"), None, Some("EPSG:4326"))] + #[case::right_only(None, Some("EPSG:3857"), Some("EPSG:3857"))] + #[case::unreferenced(None, None, None)] + fn propagates_compatible_crs( + #[case] left_crs: Option<&str>, + #[case] right_crs: Option<&str>, + #[case] expected: Option<&str>, + ) -> VortexResult<()> { + let left = dimensional_point(Dimension::Xy, [0.0; 4], left_crs)?; + let right = dimensional_point(Dimension::Xy, [1.0; 4], right_crs)?; + + let dtype = GeoMakeLine.return_dtype( + &EmptyOptions, + &[left.dtype().clone(), right.dtype().clone()], + )?; + assert_eq!( + dtype.as_extension().metadata::().crs.as_deref(), + expected + ); + Ok(()) + } + + #[test] + fn rejects_mismatched_crs() -> VortexResult<()> { + let left = dimensional_point(Dimension::Xy, [0.0; 4], Some("EPSG:4326"))?; + let right = dimensional_point(Dimension::Xy, [1.0; 4], Some("EPSG:3857"))?; + assert!( + GeoMakeLine + .return_dtype( + &EmptyOptions, + &[left.dtype().clone(), right.dtype().clone()] + ) + .is_err() + ); + Ok(()) + } + + #[test] + fn rejects_non_point_dtype() -> VortexResult<()> { + let points = point_column(vec![0.0], vec![0.0])?; + assert!( + GeoMakeLine + .return_dtype(&EmptyOptions, std::slice::from_ref(points.dtype())) + .is_err() + ); + let non_point = DType::Bool(vortex_array::dtype::Nullability::NonNullable); + assert!( + GeoMakeLine + .return_dtype(&EmptyOptions, &[points.dtype().clone(), non_point]) + .is_err() + ); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/mod.rs b/vortex-geo/src/scalar_fn/mod.rs index e6770be4fff..239ea5e5519 100644 --- a/vortex-geo/src/scalar_fn/mod.rs +++ b/vortex-geo/src/scalar_fn/mod.rs @@ -3,8 +3,15 @@ //! 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 hilbert; +pub mod intersection; pub mod intersects; +pub mod length; +pub mod make_line;