From 819b980a4eca17d4aabab9434aa946414763c964 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:39:07 -0400 Subject: [PATCH 1/4] feat(vortex-geo): add intersection scalar function Signed-off-by: Nemo Yu --- vortex-spatial/src/extension/multipolygon.rs | 24 ++ vortex-spatial/src/lib.rs | 2 + vortex-spatial/src/scalar_fn/intersection.rs | 425 +++++++++++++++++++ vortex-spatial/src/scalar_fn/mod.rs | 1 + 4 files changed, 452 insertions(+) create mode 100644 vortex-spatial/src/scalar_fn/intersection.rs diff --git a/vortex-spatial/src/extension/multipolygon.rs b/vortex-spatial/src/extension/multipolygon.rs index 80078a4b07e..2c30c096395 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -13,9 +13,11 @@ use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use geo_traits::to_geo::ToGeoGeometry; use geo_types::Geometry; +use geoarrow::array::GeoArrowArray; use geoarrow::array::GeoArrowArrayAccessor; use geoarrow::array::IntoArrow; use geoarrow::array::MultiPolygonArray; +use geoarrow::array::MultiPolygonBuilder; use geoarrow::datatypes::CoordType; use geoarrow::datatypes::MultiPolygonType; use prost::Message; @@ -24,6 +26,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; @@ -116,6 +119,27 @@ fn multipolygon_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) - MultiPolygonType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) } +/// Build a native 2-D [`MultiPolygon`] array from row-oriented `geo_types` multipolygons. +pub(crate) fn build_multipolygon_array( + multipolygons: &[Option>], + metadata: SpatialMetadata, + nullability: Nullability, +) -> VortexResult { + let multipolygons = MultiPolygonBuilder::from_nullable_multi_polygons( + multipolygons, + multipolygon_type(&metadata, Dimension::Xy), + ) + .finish(); + let storage_dtype = multipolygon_storage_dtype(Dimension::Xy, nullability); + let storage = ArrayRef::from_arrow( + multipolygons.to_array_ref().as_ref(), + nullability == Nullability::Nullable, + )? + .cast(storage_dtype.clone())?; + let ext_dtype = ExtDType::::try_new(metadata, storage_dtype)?; + Ok(ExtensionArray::try_new(ext_dtype.erased(), storage)?.into_array()) +} + /// Decode storage to `geo_types` for the spatial scalar functions (CRS is irrelevant to planar ops). pub(crate) fn multipolygon_geometries( storage: &ArrayRef, diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index 7b90b357e89..0a310b1ddb3 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -27,6 +27,7 @@ use crate::scalar_fn::contains::SpatialContains; use crate::scalar_fn::convex_hull::SpatialConvexHull; use crate::scalar_fn::distance::SpatialDistance; use crate::scalar_fn::envelope::SpatialEnvelope; +use crate::scalar_fn::intersection::SpatialIntersection; use crate::scalar_fn::intersects::SpatialIntersects; use crate::scalar_fn::make_line::SpatialMakeLine; @@ -72,6 +73,7 @@ pub fn initialize(session: &VortexSession) { session.scalar_fns().register(SpatialCollect); session.scalar_fns().register(SpatialConvexHull); session.scalar_fns().register(SpatialEnvelope); + session.scalar_fns().register(SpatialIntersection); session.scalar_fns().register(SpatialContains); session.scalar_fns().register(SpatialDistance); session.scalar_fns().register(SpatialIntersects); diff --git a/vortex-spatial/src/scalar_fn/intersection.rs b/vortex-spatial/src/scalar_fn/intersection.rs new file mode 100644 index 00000000000..44a448d4ff7 --- /dev/null +++ b/vortex-spatial/src/scalar_fn/intersection.rs @@ -0,0 +1,425 @@ +// 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::SpatialMetadata; +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: &SpatialMetadata, + right: &SpatialMetadata, +) -> VortexResult { + match (&left.crs, &right.crs) { + (Some(left_crs), Some(right_crs)) => { + vortex_ensure!( + left_crs == right_crs, + "spatial: intersection operands have different coordinate reference systems: \ + {left_crs} and {right_crs}" + ); + Ok(left.clone()) + } + (Some(_), None) => Ok(left.clone()), + (None, Some(_)) => Ok(right.clone()), + (None, None) => Ok(SpatialMetadata::default()), + } +} + +/// Resolve the strict native `Polygon x Polygon -> MultiPolygon` overload. +fn intersection_dtype(dtypes: &[DType]) -> VortexResult> { + vortex_ensure!( + dtypes.len() == 2, + "spatial: 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::()), + "spatial: 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 SpatialIntersection; + +impl SpatialIntersection { + /// 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(SpatialIntersection, EmptyOptions).erased(), + vec![left, right], + ) + } +} + +impl ScalarFnVTable for SpatialIntersection { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.st.intersection"); + *ID + } + + fn serialize(&self, _: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _: &Self::Options) -> Arity { + Arity::Exact(2) + } + + fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("left"), + 1 => ChildName::from("right"), + _ => unreachable!("intersection has exactly two children"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + Ok(DType::Extension(intersection_dtype(dtypes)?.erased())) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let left = args.get(0)?; + let right = args.get(1)?; + let output_dtype = intersection_dtype(&[left.dtype().clone(), right.dtype().clone()])?; + dispatch_binary( + &left, + &right, + DType::Extension(output_dtype.clone().erased()), + |execution, ctx| execute_intersection(execution, &output_dtype, ctx), + ctx, + ) + } + + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use geo_types::Geometry; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Columnar; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::SpatialIntersection; + use crate::extension::MultiPolygon; + use crate::extension::geometries; + use crate::scalar_fn::area::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 = SpatialIntersection::try_new_array(left, right)?.into_array(); + assert!(intersections.dtype().as_extension().is::()); + + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let decoded = geometries(&intersections, &mut ctx)?; + let polygon_counts = decoded + .iter() + .map(|geometry| match geometry { + Geometry::MultiPolygon(multipolygon) => Ok(multipolygon.0.len()), + other => Err(vortex_err!( + "intersection decoded as {other:?}, expected MultiPolygon" + )), + }) + .collect::>>()?; + assert_eq!(polygon_counts, [1, 0, 0]); + + let areas = 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 = SpatialIntersection::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 = SpatialIntersection::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 = SpatialIntersection::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 = SpatialIntersection::try_new_array(left, right)?.into_array(); + let Columnar::Constant(constant) = result.execute::(&mut ctx)? else { + return Err(vortex_err!( + "intersection of two constants should remain constant" + )); + }; + assert_eq!(constant.len(), 3); + Ok(()) + } + + #[rstest] + #[case::none(0)] + #[case::one(1)] + #[case::three(3)] + fn rejects_wrong_arity(#[case] arity: usize) -> VortexResult<()> { + let dtype = polygon_column(vec![vec![]])?.dtype().clone(); + assert!( + SpatialIntersection + .return_dtype(&EmptyOptions, &vec![dtype; arity]) + .is_err() + ); + Ok(()) + } + + #[test] + fn rejects_non_polygon_input() -> VortexResult<()> { + let polygon = polygon_column(vec![vec![]])?; + let point = point_column(vec![0.0], vec![0.0])?; + assert!(SpatialIntersection::try_new_array(polygon, point).is_err()); + Ok(()) + } +} diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index fcb99e37db0..a0f1be88538 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -10,5 +10,6 @@ pub mod convex_hull; pub mod distance; pub mod envelope; mod execute; +pub mod intersection; pub mod intersects; pub mod make_line; From 824f5e8f4c9ab3ec5c94720b688b4e5c940b746b Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:39:12 -0400 Subject: [PATCH 2/4] bench(vortex-geo): add intersection benchmark Signed-off-by: Nemo Yu --- vortex-spatial/Cargo.toml | 4 + vortex-spatial/benches/intersection.rs | 107 +++++++++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 vortex-spatial/benches/intersection.rs diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 5433ee31e59..951b06ecf0d 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -74,5 +74,9 @@ harness = false name = "convex_hull" harness = false +[[bench]] +name = "intersection" +harness = false + [lints] workspace = true diff --git a/vortex-spatial/benches/intersection.rs b/vortex-spatial/benches/intersection.rs new file mode 100644 index 00000000000..d0ac0f5dbde --- /dev/null +++ b/vortex-spatial/benches/intersection.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for native `ST_Intersection` over polygon pairs. +//! +//! The cases cover simple building-like rectangles, more detailed boundaries, and strict null +//! propagation. Inputs overlap because SpatialBench Q9 prefilters pairs with `ST_Intersects`. +//! +//! Run with `cargo bench -p vortex-spatial --bench intersection`. + +#![expect(clippy::unwrap_used)] + +use std::f64::consts::TAU; +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::validity::Validity; +use vortex_session::VortexSession; +use vortex_spatial::scalar_fn::intersection::SpatialIntersection; +use vortex_spatial::test_harness::polygon_column; +use vortex_spatial::test_harness::spatial_session; + +// Scalar function execution allocates its output inside the timed region, so use the vendored +// allocator instead of measuring glibc differences between CodSpeed runner images. +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(spatial_session); + +const ROWS: usize = 512; + +fn main() { + divan::main(); +} + +fn regular_polygon(cx: f64, cy: f64, radius: f64, vertices: usize) -> Vec<(f64, f64)> { + (0..=vertices) + .map(|vertex| { + let angle = TAU * (vertex % vertices) as f64 / vertices as f64; + (cx + radius * angle.cos(), cy + radius * angle.sin()) + }) + .collect() +} + +fn polygon_pairs(vertices: usize) -> (ArrayRef, ArrayRef) { + let left = polygon_column( + (0..ROWS) + .map(|row| vec![regular_polygon(row as f64, 0.0, 1.0, vertices)]) + .collect(), + ) + .unwrap(); + let right = polygon_column( + (0..ROWS) + .map(|row| vec![regular_polygon(row as f64 + 0.5, 0.0, 1.0, vertices)]) + .collect(), + ) + .unwrap(); + (left, right) +} + +fn intersections(left: &ArrayRef, right: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + SpatialIntersection::try_new_array(left.clone(), right.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +fn bench_intersections(bencher: Bencher, left: ArrayRef, right: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| intersections(&left, &right, &mut ctx)); +} + +#[divan::bench] +fn rectangles(bencher: Bencher) { + let (left, right) = polygon_pairs(4); + bench_intersections(bencher, left, right); +} + +#[divan::bench] +fn thirty_two_vertex_boundaries(bencher: Bencher) { + let (left, right) = polygon_pairs(32); + bench_intersections(bencher, left, right); +} + +#[divan::bench] +fn nullable_rectangles(bencher: Bencher) { + let (left, right) = polygon_pairs(4); + let left = MaskedArray::try_new( + left, + Validity::from_iter((0..ROWS).map(|row| !row.is_multiple_of(8))), + ) + .unwrap() + .into_array(); + bench_intersections(bencher, left, right); +} From 689582ae26223ecb6798e3f61ed9496a6da54d02 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:46:51 -0400 Subject: [PATCH 3/4] feat(vortex-geo): intersect multipolygons Signed-off-by: Nemo Yu --- vortex-spatial/src/scalar_fn/intersection.rs | 147 ++++++++++++------- 1 file changed, 92 insertions(+), 55 deletions(-) diff --git a/vortex-spatial/src/scalar_fn/intersection.rs b/vortex-spatial/src/scalar_fn/intersection.rs index 44a448d4ff7..b7f6cf04b3c 100644 --- a/vortex-spatial/src/scalar_fn/intersection.rs +++ b/vortex-spatial/src/scalar_fn/intersection.rs @@ -29,9 +29,9 @@ use vortex_mask::AllOr; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::SpatialMetadata; use crate::extension::MultiPolygon; use crate::extension::Polygon; +use crate::extension::SpatialMetadata; use crate::extension::build_multipolygon_array; use crate::extension::coordinate::Dimension; use crate::extension::geometries; @@ -61,25 +61,38 @@ fn intersection_metadata( } } -/// Resolve the strict native `Polygon x Polygon -> MultiPolygon` overload. +/// Metadata carried by a validated native polygonal dtype. +fn polygonal_metadata(dtype: &DType) -> &SpatialMetadata { + let extension = dtype.as_extension(); + if extension.is::() { + extension.metadata::() + } else if extension.is::() { + extension.metadata::() + } else { + unreachable!("intersection operand was validated as polygonal") + } +} + +/// Resolve the native polygonal intersection overloads, which always return a MultiPolygon. fn intersection_dtype(dtypes: &[DType]) -> VortexResult> { vortex_ensure!( dtypes.len() == 2, - "spatial: intersection requires exactly two Polygon operands, got {}", + "spatial: intersection requires exactly two polygonal operands, got {}", dtypes.len() ); for dtype in dtypes { vortex_ensure!( - dtype - .as_extension_opt() - .is_some_and(|extension| extension.is::()), - "spatial: intersection operand {dtype} is not a native Polygon" + dtype.as_extension_opt().is_some_and(|extension| { + extension.is::() || extension.is::() + }), + "spatial: 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, @@ -87,36 +100,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. @@ -125,10 +117,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(), @@ -140,14 +132,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)] => { @@ -155,20 +147,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 SpatialIntersection; impl SpatialIntersection { - /// 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(SpatialIntersection, EmptyOptions).erased(), @@ -265,7 +273,8 @@ mod tests { use super::SpatialIntersection; use crate::extension::MultiPolygon; use crate::extension::geometries; - use crate::scalar_fn::area::GeoArea; + use crate::scalar_fn::area::SpatialArea; + use crate::test_harness::multipolygon_column; use crate::test_harness::point_column; use crate::test_harness::polygon_column; @@ -288,6 +297,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![ @@ -316,12 +333,32 @@ mod tests { .collect::>>()?; assert_eq!(polygon_counts, [1, 0, 0]); - let areas = GeoArea::try_new_array(intersections)?.into_array(); + let areas = SpatialArea::try_new_array(intersections)?.into_array(); let expected = PrimitiveArray::from_iter([1.0_f64, 0.0, 0.0]).into_array(); assert_arrays_eq!(areas, expected, &mut ctx); Ok(()) } + #[rstest] + #[case::polygon_polygon(false, false)] + #[case::polygon_multipolygon(false, true)] + #[case::multipolygon_polygon(true, false)] + #[case::multipolygon_multipolygon(true, true)] + fn supports_all_polygonal_combinations( + #[case] left_multi: bool, + #[case] right_multi: bool, + ) -> VortexResult<()> { + let left = polygonal_column(square(0.0, 0.0, 2.0, 2.0), left_multi)?; + let right = polygonal_column(square(1.0, 1.0, 3.0, 3.0), right_multi)?; + let intersections = SpatialIntersection::try_new_array(left, right)?.into_array(); + let areas = SpatialArea::try_new_array(intersections)?.into_array(); + let expected = PrimitiveArray::from_iter([1.0_f64]).into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(areas, expected, &mut ctx); + Ok(()) + } + #[test] fn preserves_holes() -> VortexResult<()> { let left = polygon_column(vec![vec![ @@ -330,7 +367,7 @@ mod tests { ]])?; let right = polygon_column(vec![vec![square(2.0, 0.0, 5.0, 4.0)]])?; let intersections = SpatialIntersection::try_new_array(left, right)?.into_array(); - let areas = GeoArea::try_new_array(intersections)?.into_array(); + let areas = SpatialArea::try_new_array(intersections)?.into_array(); let expected = PrimitiveArray::from_iter([6.0_f64]).into_array(); let mut ctx = vortex_array::array_session().create_execution_ctx(); @@ -353,7 +390,7 @@ mod tests { vec![square(1.0, 1.0, 3.0, 3.0)], ])?; let intersections = SpatialIntersection::try_new_array(left, right)?.into_array(); - let areas = GeoArea::try_new_array(intersections)?.into_array(); + let areas = SpatialArea::try_new_array(intersections)?.into_array(); let expected = PrimitiveArray::new(vec![1.0_f64, 0.0], Validity::from_iter([true, false])) .into_array(); let mut ctx = vortex_array::array_session().create_execution_ctx(); @@ -379,7 +416,7 @@ mod tests { }; let intersections = SpatialIntersection::try_new_array(left, right)?.into_array(); - let areas = GeoArea::try_new_array(intersections)?.into_array(); + let areas = SpatialArea::try_new_array(intersections)?.into_array(); let expected = PrimitiveArray::from_iter([1.0_f64, 0.0]).into_array(); assert_arrays_eq!(areas, expected, &mut ctx); Ok(()) @@ -416,7 +453,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!(SpatialIntersection::try_new_array(polygon, point).is_err()); From 99158a58f9aa3ec20d44eeafd2000ad7af4e72f3 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 7 Aug 2026 14:36:13 -0400 Subject: [PATCH 4/4] refactor(vortex-geo): share structural binary geometry execution Signed-off-by: Nemo Yu --- vortex-spatial/src/extension/multipolygon.rs | 12 +- vortex-spatial/src/scalar_fn/area.rs | 7 +- vortex-spatial/src/scalar_fn/collect.rs | 5 +- vortex-spatial/src/scalar_fn/contains.rs | 3 + vortex-spatial/src/scalar_fn/distance.rs | 12 +- vortex-spatial/src/scalar_fn/execute.rs | 6 +- .../src/scalar_fn/execute/binary.rs | 87 ++++++------ .../src/scalar_fn/execute/geo_types.rs | 128 +++++++++--------- vortex-spatial/src/scalar_fn/execute/unary.rs | 14 +- vortex-spatial/src/scalar_fn/intersection.rs | 75 +--------- vortex-spatial/src/scalar_fn/intersects.rs | 3 + 11 files changed, 152 insertions(+), 200 deletions(-) diff --git a/vortex-spatial/src/extension/multipolygon.rs b/vortex-spatial/src/extension/multipolygon.rs index 2c30c096395..0aecf28038d 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -30,6 +30,7 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtDTypeRef; use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; @@ -122,22 +123,21 @@ fn multipolygon_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) - /// Build a native 2-D [`MultiPolygon`] array from row-oriented `geo_types` multipolygons. pub(crate) fn build_multipolygon_array( multipolygons: &[Option>], - metadata: SpatialMetadata, - nullability: Nullability, + ext_dtype: &ExtDTypeRef, ) -> VortexResult { + let nullability = ext_dtype.storage_dtype().nullability(); let multipolygons = MultiPolygonBuilder::from_nullable_multi_polygons( multipolygons, - multipolygon_type(&metadata, Dimension::Xy), + multipolygon_type(ext_dtype.metadata::(), Dimension::Xy), ) .finish(); - let storage_dtype = multipolygon_storage_dtype(Dimension::Xy, nullability); + let storage_dtype = ext_dtype.storage_dtype(); 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()) + Ok(ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array()) } /// Decode storage to `geo_types` for the spatial scalar functions (CRS is irrelevant to planar ops). diff --git a/vortex-spatial/src/scalar_fn/area.rs b/vortex-spatial/src/scalar_fn/area.rs index e1ea91db06a..73a3951da5f 100644 --- a/vortex-spatial/src/scalar_fn/area.rs +++ b/vortex-spatial/src/scalar_fn/area.rs @@ -97,7 +97,12 @@ impl ScalarFnVTable for SpatialArea { ctx: &mut ExecutionCtx, ) -> VortexResult { let array = args.get(0)?; - execute_unary_geo_types(&array, Area::unsigned_area, ctx) + execute_unary_geo_types( + &array, + DType::Primitive(PType::F64, array.dtype().nullability()), + Area::unsigned_area, + ctx, + ) } fn validity( diff --git a/vortex-spatial/src/scalar_fn/collect.rs b/vortex-spatial/src/scalar_fn/collect.rs index 2f80401904e..b166b3c4d23 100644 --- a/vortex-spatial/src/scalar_fn/collect.rs +++ b/vortex-spatial/src/scalar_fn/collect.rs @@ -202,6 +202,7 @@ fn execute_collect( output_dtype: &ExtDTypeRef, ctx: &mut ExecutionCtx, ) -> VortexResult { + let nullability = output_dtype.storage_dtype().nullability(); match execution.operands { [Operand::Constant(scalar)] => { let one = ConstantArray::new(scalar, 1) @@ -209,7 +210,7 @@ fn execute_collect( .execute::(ctx)?; let collected = collect_list( one, - Validity::from_mask(Mask::new_true(1), execution.nullability), + Validity::from_mask(Mask::new_true(1), nullability), output_dtype, ctx, )?; @@ -219,7 +220,7 @@ fn execute_collect( let valid = execution.valid.execute_mask(execution.len, ctx)?; collect_list( array.execute::(ctx)?, - Validity::from_mask(valid, execution.nullability), + Validity::from_mask(valid, nullability), output_dtype, ctx, ) diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 599c0eee2be..49d6b33b7e2 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -107,6 +107,9 @@ impl ScalarFnVTable for SpatialContains { execute_binary_geo_types( &a, &b, + DType::Bool(Nullability::from( + a.dtype().is_nullable() || b.dtype().is_nullable(), + )), |a, b| a.contains(b), Some(|ra, rb| (!ra.contains(rb)).then_some(false)), ctx, diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index dfd3d09ed23..daf6e73ddd1 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -103,7 +103,17 @@ impl ScalarFnVTable for SpatialDistance { let a = args.get(0)?; let b = args.get(1)?; // Distance is a value, not a verdict: no bounding-rect test can decide it. - execute_binary_geo_types(&a, &b, |x, y| Euclidean.distance(x, y), None, ctx) + execute_binary_geo_types( + &a, + &b, + DType::Primitive( + PType::F64, + Nullability::from(a.dtype().is_nullable() || b.dtype().is_nullable()), + ), + |x, y| Euclidean.distance(x, y), + None, + ctx, + ) } fn validity( diff --git a/vortex-spatial/src/scalar_fn/execute.rs b/vortex-spatial/src/scalar_fn/execute.rs index 3a7494bcb39..105f09983ae 100644 --- a/vortex-spatial/src/scalar_fn/execute.rs +++ b/vortex-spatial/src/scalar_fn/execute.rs @@ -9,8 +9,7 @@ //! //! [`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. +//! `geo_types::Geometry`; the final output is built by a primitive or structural output builder. mod binary; mod geo_types; @@ -21,7 +20,6 @@ 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; use vortex_mask::Mask; @@ -44,6 +42,4 @@ pub(crate) struct Execution { pub(crate) valid: V, /// Number of output rows. pub(crate) len: usize, - /// Output nullability from the scalar function's return dtype. - pub(crate) nullability: Nullability, } diff --git a/vortex-spatial/src/scalar_fn/execute/binary.rs b/vortex-spatial/src/scalar_fn/execute/binary.rs index f2c03bd1beb..cdedb30176c 100644 --- a/vortex-spatial/src/scalar_fn/execute/binary.rs +++ b/vortex-spatial/src/scalar_fn/execute/binary.rs @@ -12,7 +12,6 @@ use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -20,8 +19,7 @@ use vortex_mask::Mask; use super::Execution; use super::Operand; use super::geo_types::GeoTypesOutput; -use super::geo_types::eval_column; -use super::geo_types::eval_column_pair; +use crate::extension::geometries; use crate::extension::single_geometry; /// Dispatch a binary strict geometry kernel over constants and columns. @@ -85,7 +83,6 @@ where operands: [left, right], valid, len, - nullability: output_dtype.nullability(), }, ctx, ) @@ -106,67 +103,61 @@ pub(crate) type BboxPrecheck = fn(&Rect, &Rect) -> Option; pub(crate) fn execute_binary_geo_types( left: &ArrayRef, right: &ArrayRef, + output_dtype: DType, compute: F, bbox_precheck: Option>, ctx: &mut ExecutionCtx, ) -> VortexResult where - T: GeoTypesOutput, F: Fn(&Geometry, &Geometry) -> T + Copy, + T: GeoTypesOutput, { - let nullability = Nullability::from(left.dtype().is_nullable() || right.dtype().is_nullable()); dispatch_binary( left, right, - T::dtype(nullability), + output_dtype.clone(), |execution, ctx| match execution.operands { [Operand::Constant(left), Operand::Constant(right)] => { let left = single_geometry(&left, ctx)?; let right = single_geometry(&right, ctx)?; - Ok(ConstantArray::new( - compute(&left, &right).into_scalar(execution.nullability), - execution.len, - ) - .into_array()) + T::build_constant(compute(&left, &right), execution.len, &output_dtype, ctx) } [Operand::Constant(left), Operand::Column(right)] => { let left = single_geometry(&left, ctx)?; let prescreen = bbox_precheck.zip(left.bounding_rect()); - eval_column( - &right, - &execution.valid, - |right| { + let values = geometries(&right.filter(execution.valid.clone())?, ctx)? + .iter() + .map(|right| { prescreen .and_then(|(precheck, fixed)| precheck(&fixed, &right.bounding_rect()?)) .unwrap_or_else(|| compute(&left, right)) - }, - execution.nullability, - ctx, - ) + }) + .collect(); + T::build_array(execution.len, &execution.valid, values, &output_dtype, ctx) } [Operand::Column(left), Operand::Constant(right)] => { let right = single_geometry(&right, ctx)?; let prescreen = bbox_precheck.zip(right.bounding_rect()); - eval_column( - &left, - &execution.valid, - |left| { + let values = geometries(&left.filter(execution.valid.clone())?, ctx)? + .iter() + .map(|left| { prescreen .and_then(|(precheck, fixed)| precheck(&left.bounding_rect()?, &fixed)) .unwrap_or_else(|| compute(left, &right)) - }, - execution.nullability, - ctx, - ) + }) + .collect(); + T::build_array(execution.len, &execution.valid, values, &output_dtype, ctx) + } + [Operand::Column(left), Operand::Column(right)] => { + let left = geometries(&left.filter(execution.valid.clone())?, ctx)?; + let right = geometries(&right.filter(execution.valid.clone())?, ctx)?; + let values = left + .iter() + .zip(&right) + .map(|(left, right)| compute(left, right)) + .collect(); + T::build_array(execution.len, &execution.valid, values, &output_dtype, ctx) } - [Operand::Column(left), Operand::Column(right)] => eval_column_pair( - &left, - &right, - &execution.valid, - compute, - execution.nullability, - ctx, - ), }, ctx, ) @@ -186,6 +177,7 @@ mod tests { use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; @@ -226,6 +218,7 @@ mod tests { let result = execute_binary_geo_types( &triangle, &probes, + DType::Bool(vortex_array::dtype::Nullability::NonNullable), counting_intersects(&exact_runs), Some(DISJOINT_PRECHECK), &mut ctx, @@ -247,6 +240,7 @@ mod tests { let result = execute_binary_geo_types( &triangle, &probes, + DType::Bool(vortex_array::dtype::Nullability::Nullable), counting_intersects(&exact_runs), Some(DISJOINT_PRECHECK), &mut ctx, @@ -277,6 +271,7 @@ mod tests { let result = execute_binary_geo_types( &probes, &triangle, + DType::Bool(vortex_array::dtype::Nullability::NonNullable), counted, Some(|left, right| (!left.contains(right)).then_some(false)), &mut ctx, @@ -299,6 +294,7 @@ mod tests { let result = execute_binary_geo_types( &empty, &probes, + DType::Bool(vortex_array::dtype::Nullability::NonNullable), counting_intersects(&exact_runs), Some(DISJOINT_PRECHECK), &mut ctx, @@ -324,9 +320,22 @@ mod tests { ])?; let exact = |left: &Geometry, right: &Geometry| left.intersects(right); - let with_precheck = - execute_binary_geo_types(&triangle, &probes, exact, Some(DISJOINT_PRECHECK), &mut ctx)?; - let exact_only = execute_binary_geo_types(&triangle, &probes, exact, None, &mut ctx)?; + let with_precheck = execute_binary_geo_types( + &triangle, + &probes, + DType::Bool(vortex_array::dtype::Nullability::Nullable), + exact, + Some(DISJOINT_PRECHECK), + &mut ctx, + )?; + let exact_only = execute_binary_geo_types( + &triangle, + &probes, + DType::Bool(vortex_array::dtype::Nullability::Nullable), + exact, + None, + &mut ctx, + )?; assert_arrays_eq!(with_precheck, exact_only, &mut ctx); Ok(()) diff --git a/vortex-spatial/src/scalar_fn/execute/geo_types.rs b/vortex-spatial/src/scalar_fn/execute/geo_types.rs index 038aca46502..47f9fad44f9 100644 --- a/vortex-spatial/src/scalar_fn/execute/geo_types.rs +++ b/vortex-spatial/src/scalar_fn/execute/geo_types.rs @@ -7,87 +7,94 @@ //! and return Vortex arrays; they do not expose `geo_types` values as scalar-function outputs. 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::BoolArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_mask::AllOr; use vortex_mask::Mask; +use crate::extension::build_multipolygon_array; use crate::extension::geometries; -/// A primitive result produced after kernel inputs are decoded to `geo_types`. -pub(crate) trait GeoTypesOutput: Copy { - /// The Vortex dtype used to represent this output. - fn dtype(nullability: Nullability) -> DType; - - /// Convert one computed value into a Vortex scalar for constant output. - fn into_scalar(self, nullability: Nullability) -> Scalar; - - /// Scatter values computed for valid rows into a full-length output array. +/// A Vortex representation for values produced by a `geo_types` kernel. +pub(crate) trait GeoTypesOutput: Sized { + /// Build an array from values computed for rows selected by `valid`. fn build_array( len: usize, valid: &Mask, values: Vec, - nullability: Nullability, - ) -> ArrayRef; -} + output_dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult; -impl GeoTypesOutput for f64 { - fn dtype(nullability: Nullability) -> DType { - DType::Primitive(PType::F64, nullability) + /// Build a repeated result for two constant operands. + fn build_constant( + value: Self, + len: usize, + output_dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let one = Self::build_array(1, &Mask::new_true(1), vec![value], output_dtype, ctx)?; + Ok(ConstantArray::new(one.execute_scalar(0, ctx)?, len).into_array()) } +} - fn into_scalar(self, nullability: Nullability) -> Scalar { - Scalar::primitive(self, nullability) +/// Scatter values computed for valid rows into a full-length nullable vector. +pub(crate) fn scatter_valid(len: usize, valid: &Mask, values: Vec) -> Vec> { + match valid.indices() { + AllOr::All => values.into_iter().map(Some).collect(), + AllOr::None => (0..len).map(|_| None).collect(), + AllOr::Some(rows) => { + let mut output = (0..len).map(|_| None).collect::>>(); + for (&row, value) in rows.iter().zip(values) { + output[row] = Some(value); + } + output + } } +} +impl GeoTypesOutput for f64 { fn build_array( len: usize, valid: &Mask, values: Vec, - nullability: Nullability, - ) -> ArrayRef { - let validity = Validity::from_mask(valid.clone(), nullability); - match valid.indices() { + output_dtype: &DType, + _: &mut ExecutionCtx, + ) -> VortexResult { + let validity = Validity::from_mask(valid.clone(), output_dtype.nullability()); + Ok(match valid.indices() { AllOr::All => PrimitiveArray::new(values, validity).into_array(), - AllOr::None => PrimitiveArray::new(vec![0.0f64; len], validity).into_array(), + AllOr::None => PrimitiveArray::new(vec![0.0; len], validity).into_array(), AllOr::Some(rows) => { - let mut data = vec![0.0f64; len]; + let mut data = vec![0.0; len]; for (&row, value) in rows.iter().zip(values) { data[row] = value; } PrimitiveArray::new(data, validity).into_array() } - } + }) } } impl GeoTypesOutput for bool { - fn dtype(nullability: Nullability) -> DType { - DType::Bool(nullability) - } - - fn into_scalar(self, nullability: Nullability) -> Scalar { - Scalar::bool(self, nullability) - } - fn build_array( len: usize, valid: &Mask, values: Vec, - nullability: Nullability, - ) -> ArrayRef { - let validity = Validity::from_mask(valid.clone(), nullability); - match valid.indices() { + output_dtype: &DType, + _: &mut ExecutionCtx, + ) -> VortexResult { + let validity = Validity::from_mask(valid.clone(), output_dtype.nullability()); + Ok(match valid.indices() { AllOr::All => BoolArray::new(BitBuffer::from_iter(values), validity).into_array(), AllOr::None => BoolArray::new(BitBuffer::new_unset(len), validity).into_array(), AllOr::Some(rows) => { @@ -97,7 +104,20 @@ impl GeoTypesOutput for bool { } BoolArray::new(BitBuffer::from_iter(data), validity).into_array() } - } + }) + } +} + +impl GeoTypesOutput for GeoMultiPolygon { + fn build_array( + len: usize, + valid: &Mask, + values: Vec, + output_dtype: &DType, + _: &mut ExecutionCtx, + ) -> VortexResult { + let output_dtype = output_dtype.as_extension(); + build_multipolygon_array(&scatter_valid(len, valid, values), output_dtype) } } @@ -106,7 +126,7 @@ pub(super) fn eval_column( column: &ArrayRef, valid: &Mask, compute: F, - nullability: Nullability, + output_dtype: &DType, ctx: &mut ExecutionCtx, ) -> VortexResult where @@ -116,29 +136,5 @@ where let len = column.len(); let decoded = geometries(&column.filter(valid.clone())?, ctx)?; let values = decoded.iter().map(compute).collect(); - Ok(T::build_array(len, valid, values, nullability)) -} - -/// Evaluate a decoded kernel over rows where both geometry columns are valid. -pub(super) fn eval_column_pair( - left: &ArrayRef, - right: &ArrayRef, - valid: &Mask, - compute: F, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T, -{ - let len = left.len(); - let left = geometries(&left.filter(valid.clone())?, ctx)?; - let right = geometries(&right.filter(valid.clone())?, ctx)?; - let values = left - .iter() - .zip(&right) - .map(|(left, right)| compute(left, right)) - .collect(); - Ok(T::build_array(len, valid, values, nullability)) + T::build_array(len, valid, values, output_dtype, ctx) } diff --git a/vortex-spatial/src/scalar_fn/execute/unary.rs b/vortex-spatial/src/scalar_fn/execute/unary.rs index a8bb74f850c..59c96695136 100644 --- a/vortex-spatial/src/scalar_fn/execute/unary.rs +++ b/vortex-spatial/src/scalar_fn/execute/unary.rs @@ -44,7 +44,6 @@ where operands: [Operand::Constant(constant.scalar().clone())], valid: Validity::AllValid, len, - nullability: output_dtype.nullability(), }, ctx, ); @@ -59,7 +58,6 @@ where operands: [Operand::Column(array.clone())], valid, len, - nullability: output_dtype.nullability(), }, ctx, ) @@ -72,6 +70,7 @@ where /// before broadcast; a column is decoded only for its valid rows. pub(crate) fn execute_unary_geo_types( array: &ArrayRef, + output_dtype: DType, compute: F, ctx: &mut ExecutionCtx, ) -> VortexResult @@ -79,22 +78,17 @@ where T: GeoTypesOutput, F: Fn(&Geometry) -> T, { - let nullability = array.dtype().nullability(); dispatch_unary( array, - T::dtype(nullability), + output_dtype.clone(), |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()) + T::build_constant(compute(&geometry), execution.len, &output_dtype, ctx) } [Operand::Column(array)] => { let valid = execution.valid.execute_mask(execution.len, ctx)?; - eval_column(&array, &valid, compute, execution.nullability, ctx) + eval_column(&array, &valid, compute, &output_dtype, ctx) } }, ctx, diff --git a/vortex-spatial/src/scalar_fn/intersection.rs b/vortex-spatial/src/scalar_fn/intersection.rs index b7f6cf04b3c..92570efddbb 100644 --- a/vortex-spatial/src/scalar_fn/intersection.rs +++ b/vortex-spatial/src/scalar_fn/intersection.rs @@ -8,8 +8,6 @@ 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; @@ -25,21 +23,15 @@ use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::AllOr; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::MultiPolygon; use crate::extension::Polygon; use crate::extension::SpatialMetadata; -use crate::extension::build_multipolygon_array; use crate::extension::coordinate::Dimension; -use crate::extension::geometries; use crate::extension::multipolygon_storage_dtype; -use crate::extension::single_geometry; -use crate::scalar_fn::execute::Execution; -use crate::scalar_fn::execute::Operand; -use crate::scalar_fn::execute::dispatch_binary; +use crate::scalar_fn::execute::execute_binary_geo_types; /// Resolve CRS metadata shared by two polygon operands. fn intersection_metadata( @@ -111,64 +103,6 @@ fn polygonal_intersection(left: &Geometry, right: &Geometry) -> GeoMul } } -/// 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`. @@ -226,11 +160,12 @@ impl ScalarFnVTable for SpatialIntersection { let left = args.get(0)?; let right = args.get(1)?; let output_dtype = intersection_dtype(&[left.dtype().clone(), right.dtype().clone()])?; - dispatch_binary( + execute_binary_geo_types( &left, &right, - DType::Extension(output_dtype.clone().erased()), - |execution, ctx| execute_intersection(execution, &output_dtype, ctx), + DType::Extension(output_dtype.erased()), + polygonal_intersection, + None, ctx, ) } diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index bdabd2b9967..819fa01dd96 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -105,6 +105,9 @@ impl ScalarFnVTable for SpatialIntersects { execute_binary_geo_types( &a, &b, + DType::Bool(Nullability::from( + a.dtype().is_nullable() || b.dtype().is_nullable(), + )), |x, y| x.intersects(y), Some(|ra, rb| (!ra.intersects(rb)).then_some(false)), ctx,