From 9cb2d3bf682c6c8c965cb8d4622e2eaedaea3a3d Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:39:07 -0400 Subject: [PATCH 1/3] feat(vortex-geo): add intersection scalar function Signed-off-by: Nemo Yu --- vortex-geo/src/extension/multipolygon.rs | 24 ++ vortex-geo/src/lib.rs | 2 + vortex-geo/src/scalar_fn/intersection.rs | 422 +++++++++++++++++++++++ vortex-geo/src/scalar_fn/mod.rs | 1 + 4 files changed, 449 insertions(+) create mode 100644 vortex-geo/src/scalar_fn/intersection.rs 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/lib.rs b/vortex-geo/src/lib.rs index 901b8223947..ebd25fce8bc 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -27,6 +27,7 @@ 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::intersection::GeoIntersection; use crate::scalar_fn::intersects::GeoIntersects; use crate::scalar_fn::length::GeoLength; use crate::scalar_fn::make_line::GeoMakeLine; @@ -76,6 +77,7 @@ pub fn initialize(session: &VortexSession) { session.scalar_fns().register(GeoContains); session.scalar_fns().register(GeoDistance); session.scalar_fns().register(GeoIntersects); + session.scalar_fns().register(GeoIntersection); session.scalar_fns().register(GeoLength); session.scalar_fns().register(GeoMakeLine); diff --git a/vortex-geo/src/scalar_fn/intersection.rs b/vortex-geo/src/scalar_fn/intersection.rs new file mode 100644 index 00000000000..a85e970f7d4 --- /dev/null +++ b/vortex-geo/src/scalar_fn/intersection.rs @@ -0,0 +1,422 @@ +// 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()), + } +} + +/// Resolve the strict native `Polygon x Polygon -> MultiPolygon` overload. +fn intersection_dtype(dtypes: &[DType]) -> VortexResult> { + vortex_ensure!( + dtypes.len() == 2, + "geo: intersection requires exactly two Polygon operands, got {}", + dtypes.len() + ); + for dtype in dtypes { + vortex_ensure!( + dtype + .as_extension_opt() + .is_some_and(|extension| extension.is::()), + "geo: intersection operand {dtype} is not a native Polygon" + ); + } + + let left = dtypes[0].as_extension(); + let right = dtypes[1].as_extension(); + let metadata = intersection_metadata(left.metadata::(), right.metadata::())?; + let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); + ExtDType::try_new( + metadata, + multipolygon_storage_dtype(Dimension::Xy, nullability), + ) +} + +/// Intersect two decoded polygon values. +fn intersect(left: &Geometry, right: &Geometry) -> GeoMultiPolygon { + let (Geometry::Polygon(left), Geometry::Polygon(right)) = (left, right) else { + unreachable!("intersection operands were validated as Polygon") + }; + left.intersection(right) +} + +/// Scatter valid intersection results and build their native MultiPolygon array. +fn build_intersections( + intersections: Vec>, + execution: &Execution<2>, + output_dtype: &ExtDType, +) -> VortexResult { + 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, + ) +} + +/// Execute intersection after shared binary shape and null dispatch. +fn execute_intersection( + execution: Execution<2>, + output_dtype: &ExtDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let intersections = match &execution.operands { + [Operand::Constant(left), Operand::Constant(right)] => { + let intersection = + intersect(&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| intersect(&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| intersect(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)| intersect(left, right)) + .collect() + } + }; + build_intersections(intersections, &execution, output_dtype) +} + +/// Compute the pairwise two-dimensional intersection of native `Polygon` 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 polygon 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::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()) + } + + #[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(()) + } + + #[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_polygon_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/mod.rs b/vortex-geo/src/scalar_fn/mod.rs index 99fe5d28528..ecf3d3ca9a8 100644 --- a/vortex-geo/src/scalar_fn/mod.rs +++ b/vortex-geo/src/scalar_fn/mod.rs @@ -10,6 +10,7 @@ pub mod convex_hull; pub mod distance; pub mod envelope; mod execute; +pub mod intersection; pub mod intersects; pub mod length; pub mod make_line; From 9d09e680b5fb1803578e344bb3ee490ffb9ebcbc Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:39:12 -0400 Subject: [PATCH 2/3] bench(vortex-geo): add intersection benchmark Signed-off-by: Nemo Yu --- vortex-geo/Cargo.toml | 4 ++ vortex-geo/benches/intersection.rs | 107 +++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 vortex-geo/benches/intersection.rs diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index 0b4262fcf90..0d2438634fc 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -78,5 +78,9 @@ harness = false name = "convex_hull" harness = false +[[bench]] +name = "intersection" +harness = false + [lints] workspace = true 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); +} From f2b5e901f8f18a793b69d1b8732424142526f0f0 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:46:51 -0400 Subject: [PATCH 3/3] feat(vortex-geo): intersect multipolygons Signed-off-by: Nemo Yu --- vortex-geo/src/scalar_fn/intersection.rs | 135 +++++++++++++++-------- 1 file changed, 86 insertions(+), 49 deletions(-) diff --git a/vortex-geo/src/scalar_fn/intersection.rs b/vortex-geo/src/scalar_fn/intersection.rs index a85e970f7d4..b23336c524e 100644 --- a/vortex-geo/src/scalar_fn/intersection.rs +++ b/vortex-geo/src/scalar_fn/intersection.rs @@ -58,25 +58,38 @@ fn intersection_metadata(left: &GeoMetadata, right: &GeoMetadata) -> VortexResul } } -/// Resolve the strict native `Polygon x Polygon -> MultiPolygon` overload. +/// 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 Polygon operands, got {}", + "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::()), - "geo: intersection operand {dtype} is not a native Polygon" + dtype.as_extension_opt().is_some_and(|extension| { + extension.is::() || extension.is::() + }), + "geo: intersection operand {dtype} is not a native Polygon or MultiPolygon" ); } - let left = dtypes[0].as_extension(); - let right = dtypes[1].as_extension(); - let metadata = intersection_metadata(left.metadata::(), right.metadata::())?; + 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, @@ -84,36 +97,15 @@ fn intersection_dtype(dtypes: &[DType]) -> VortexResult> ) } -/// Intersect two decoded polygon values. -fn intersect(left: &Geometry, right: &Geometry) -> GeoMultiPolygon { - let (Geometry::Polygon(left), Geometry::Polygon(right)) = (left, right) else { - unreachable!("intersection operands were validated as Polygon") - }; - left.intersection(right) -} - -/// Scatter valid intersection results and build their native MultiPolygon array. -fn build_intersections( - intersections: Vec>, - execution: &Execution<2>, - output_dtype: &ExtDType, -) -> VortexResult { - 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, - ) +/// 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. @@ -122,10 +114,10 @@ fn execute_intersection( output_dtype: &ExtDType, ctx: &mut ExecutionCtx, ) -> VortexResult { - let intersections = match &execution.operands { + let intersections: Vec> = match &execution.operands { [Operand::Constant(left), Operand::Constant(right)] => { let intersection = - intersect(&single_geometry(left, ctx)?, &single_geometry(right, ctx)?); + polygonal_intersection(&single_geometry(left, ctx)?, &single_geometry(right, ctx)?); let one = build_multipolygon_array( &[Some(intersection)], output_dtype.metadata().clone(), @@ -137,14 +129,14 @@ fn execute_intersection( let left = single_geometry(left, ctx)?; geometries(&right.filter(execution.valid.clone())?, ctx)? .iter() - .map(|right| intersect(&left, right)) + .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| intersect(left, &right)) + .map(|left| polygonal_intersection(left, &right)) .collect() } [Operand::Column(left), Operand::Column(right)] => { @@ -152,20 +144,36 @@ fn execute_intersection( let right = geometries(&right.filter(execution.valid.clone())?, ctx)?; left.iter() .zip(&right) - .map(|(left, right)| intersect(left, right)) + .map(|(left, right)| polygonal_intersection(left, right)) .collect() } }; - build_intersections(intersections, &execution, output_dtype) + 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` operands as a native -/// `MultiPolygon`. Disjoint and boundary-only intersections produce an empty `MultiPolygon`. +/// 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 polygon operands by row. + /// 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(), @@ -263,6 +271,7 @@ mod tests { 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; @@ -285,6 +294,14 @@ mod tests { 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![ @@ -319,6 +336,26 @@ mod tests { 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![ @@ -413,7 +450,7 @@ mod tests { } #[test] - fn rejects_non_polygon_input() -> VortexResult<()> { + 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());