From 2a2d8ae9da281aa39480fdcc61d6cbad430b6924 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 12:59:39 -0400 Subject: [PATCH] Execute spatial distance with RowFn Signed-off-by: Connor Tsui --- vortex-spatial/Cargo.toml | 2 +- vortex-spatial/src/extension/mod.rs | 57 ++++++++++++ vortex-spatial/src/extension/point.rs | 18 ++++ vortex-spatial/src/extension/polygon.rs | 18 ++++ vortex-spatial/src/scalar_fn/distance.rs | 105 +++++++---------------- vortex-spatial/src/scalar_fn/mod.rs | 1 + vortex-spatial/src/scalar_fn/row.rs | 94 ++++++++++++++++++++ 7 files changed, 218 insertions(+), 77 deletions(-) create mode 100644 vortex-spatial/src/scalar_fn/row.rs diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 3be2b2d9d66..cd306089325 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -22,7 +22,7 @@ geo-types = { workspace = true } geoarrow = { workspace = true } geoarrow-cast = { workspace = true } prost = { workspace = true } -vortex-array = { workspace = true } +vortex-array = { workspace = true, features = ["unstable_row_fns"] } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-edition = { workspace = true } diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index f24f31f02aa..8a670bb0e12 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -178,6 +178,63 @@ pub(crate) fn geometries( } } +/// The geometry a null row decodes to under [`geometries_null_tolerant`]. Arbitrary: the caller +/// guarantees null rows are never read. +pub(crate) fn placeholder_geometry() -> Geometry { + Geometry::Point(geo_types::Point::new(0.0, 0.0)) +} + +/// Whether [`geometries_null_tolerant`] supports this array without filtering null rows first. +pub(crate) fn can_decode_geometries_null_tolerant(array: &ArrayRef) -> VortexResult { + if array.validity()?.definitely_no_nulls() { + return Ok(true); + } + + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + + Ok(ext.is::() || ext.is::()) +} + +/// Decode a native geometry column that may contain null rows, writing [`placeholder_geometry`] +/// into their slots. The caller guarantees null rows are never read. +/// +/// `Ok(None)` means this geometry type has no null-tolerant decode yet (`Point` and `Polygon` are +/// covered), and the caller falls back to the filter strategy, which never decodes a null row. A +/// column with definitely no nulls delegates to the ordinary [`geometries`] for any type. +pub(crate) fn geometries_null_tolerant( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>>> { + if array.validity()?.definitely_no_nulls() { + return geometries(array, ctx).map(Some); + } + + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + let storage = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + + if ext.is::() { + point_geometries_null_tolerant(&storage, ctx).map(Some) + } else if ext.is::() { + polygon_geometries_null_tolerant(&storage, ctx).map(Some) + } else { + Ok(None) + } +} + /// Decode a constant operand scalar to one geometry, a constant of any /// supported geometry type is decoded exactly like a column. pub(crate) fn single_geometry( diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index e6a00fe8fea..b774f624c4c 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -50,6 +50,7 @@ use super::coordinate::coordinate_from_struct; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A single location: `geoarrow.point`, stored as `Struct` of non-nullable `f64`. @@ -150,6 +151,23 @@ pub(crate) fn point_geometries( .collect() } +/// Like [`point_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn point_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + point_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + impl ArrowExportVTable for Point { fn arrow_ext_id(&self) -> Id { *ARROW_POINT diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index a4c88b07b22..bce33efe6e5 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -53,6 +53,7 @@ use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A polygon: `geoarrow.polygon`, stored as `List>>` (rings of vertices). @@ -153,6 +154,23 @@ pub(crate) fn polygon_geometries( .collect() } +/// Like [`polygon_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn polygon_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + polygon_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + /// Build a geoarrow `PolygonArray` from a `Polygon`'s `List>` storage. fn polygon_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let polygon_type = polygon_type( diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index dfd3d09ed23..e41999338a6 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -6,43 +6,20 @@ use geo::Distance; use geo::Euclidean; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -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_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; 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_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Distance`. -fn validate_distance_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: distance requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: distance operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry /// operands, each a column or a constant literal. @@ -60,66 +37,41 @@ impl SpatialDistance { } } -impl ScalarFnVTable for SpatialDistance { +impl RowFn for SpatialDistance { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.distance"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &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("a"), - 1 => ChildName::from("b"), - _ => unreachable!("distance has exactly two children"), - } - } - - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_distance_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Primitive(PType::F64, nullability)) - } - - fn execute( + fn deserialize( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - 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) + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn validity( + fn dispatch>( &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 + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(GeometryRow, GeometryRow), UninitElementSink, _>( + |(a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, Euclidean.distance(a, b)) } + }, + ) } } @@ -196,8 +148,9 @@ mod tests { Ok(()) } - /// Distance passes no bounding-rect rejection: a point far outside a constant polygon's - /// bounding rect still gets its true distance, alongside an inside point at distance zero. + /// Distance is a value rather than a verdict, so no bounding-rect rejection may fire for it: a + /// point far outside a constant polygon's rect still gets its true distance. Carried over from + /// #9076, which added the rejection to the predicates but deliberately not to this function. #[test] fn distance_to_constant_polygon_is_exact() -> VortexResult<()> { let session = vortex_array::array_session(); diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index 99fe5d28528..6291075246a 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -13,3 +13,4 @@ mod execute; pub mod intersects; pub mod length; pub mod make_line; +pub(crate) mod row; diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs new file mode 100644 index 00000000000..b7353e10c1c --- /dev/null +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the geo scalar functions add to the row-function machinery: an element type that decodes a +//! native geometry column into `geo_types` geometries. + +use geo_types::Geometry; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::unstable::row::InputElement; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::extension::can_decode_geometries_null_tolerant; +use crate::extension::geometries; +use crate::extension::geometries_null_tolerant; +use crate::extension::is_native_geometry; + +/// Marker for native geometry input elements: accepts any native geometry column and presents each +/// row as a decoded `geo_types` geometry. +/// +/// The two operands of a binary geo function need not share a geometry type, since distance, +/// containment and intersection across types are all meaningful, so this validates only that the +/// column is _some_ native geometry. +pub(crate) struct GeometryRow; + +// SAFETY: [`view`](InputElement::view) returns the decoded geometry slice and +// [`view_len`](InputElement::view_len) reports that slice's exact length. +unsafe impl InputElement for GeometryRow { + type Column = Vec>; + type View<'a> = &'a [Geometry]; + type Elem<'a> = &'a Geometry; + + // A geometry row is decoded from its coordinate storage, which behind a null row holds arbitrary + // coordinates that need not describe a well-formed geometry. + const DENSE_SAFE: bool = false; + // Decoding builds a geometry from stored coordinates, and a malformed one in a *valid* row is a + // domain error rather than an infrastructural failure. + const DECODE_FALLIBLE: bool = true; + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + is_native_geometry(dtype), + "spatial: operand {dtype} is not a native geometry type" + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + geometries(&array, ctx) + } + + fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult { + can_decode_geometries_null_tolerant(array) + } + + fn get(column: &Self::Column, index: usize) -> &Geometry { + &column[index] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column.as_slice() + } + + fn view_len(view: &Self::View<'_>) -> usize { + view.len() + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry + where + Self: 'a, + { + &view[index] + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry + where + Self: 'a, + { + // SAFETY: The caller established that `index` is below the slice length returned by + // `view_len` for this exact view. + unsafe { view.get_unchecked(index) } + } + + /// Null rows decode to a placeholder geometry that the branch-and-skip row loop never reads. + /// `Point` and `Polygon` columns are covered; other geometry types return `Ok(None)` and the + /// batch falls back to the filter strategy. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + geometries_null_tolerant(&array, ctx) + } +}