From 70571c4ca85175420b1d40a3fa7f36e6d9d59bbf Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 25 Aug 2026 11:06:25 -0400 Subject: [PATCH 1/2] refactor(vortex-spatial): execute distance with RowFn Signed-off-by: Nemo Yu --- vortex-spatial/Cargo.toml | 2 +- vortex-spatial/src/scalar_fn/distance.rs | 128 ++++++++++------------- vortex-spatial/src/scalar_fn/mod.rs | 1 + vortex-spatial/src/scalar_fn/row.rs | 116 ++++++++++++++++++++ 4 files changed, 173 insertions(+), 74 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/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index 7f01b9a716f..ad01a42f725 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -6,43 +6,18 @@ 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::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; 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 +35,38 @@ impl SpatialDistance { } } -impl ScalarFnVTable for SpatialDistance { +impl RowFn for SpatialDistance { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + // The per-row distance cannot fail, but decoding a geometry operand can. + const INFALLIBLE: bool = false; + 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 { + fn deserialize( + &self, + _metadata: &[u8], + _session: &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 dispatch( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { // 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) - } - - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } - - fn is_strict(&self, _: &Self::Options) -> bool { - true - } - - fn is_infallible(&self, _: &Self::Options) -> bool { - true + visitor.visit::<(GeometryRow, GeometryRow), f64>(|(a, b)| Euclidean.distance(a, b)) } } @@ -143,6 +90,7 @@ mod tests { use vortex_error::VortexResult; use super::SpatialDistance; + use crate::test_harness::nullable_multipolygon_column; use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; use crate::test_harness::polygon_column; @@ -274,6 +222,40 @@ mod tests { Ok(()) } + /// A nullable column of a list-storage geometry type (`MultiPolygon`): nulls propagate and + /// valid rows get exact distances, with the null row's storage never decoded. + #[test] + fn distance_propagates_nulls_for_multipolygon() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let near = vec![vec![vec![ + (0.0, 0.0), + (4.0, 0.0), + (4.0, 4.0), + (0.0, 4.0), + (0.0, 0.0), + ]]]; + let far = vec![vec![vec![ + (10.0, 6.0), + (14.0, 6.0), + (14.0, 10.0), + (10.0, 10.0), + (10.0, 6.0), + ]]]; + let a = nullable_multipolygon_column(vec![Some(near), None, Some(far)])?; + let b = point_constant(7.0, 2.0, 3, &mut ctx)?; + let distance = SpatialDistance::try_new(a, b)?.into_array(); + + let expected = PrimitiveArray::new( + vec![3.0f64, 0.0, 5.0], + Validity::from_iter([true, false, true]), + ) + .into_array(); + assert_arrays_eq!(distance, expected, &mut ctx); + Ok(()) + } + /// Both operands nullable: a row is null if either operand is null there. #[test] fn distance_propagates_column_pair_nulls() -> VortexResult<()> { 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..24558ceffba --- /dev/null +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -0,0 +1,116 @@ +// 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 vortex_error::vortex_err; + +use crate::extension::geometries; +use crate::extension::is_native_geometry; + +/// The geometry a null row decodes to. Arbitrary: valid-row execution never reads null slots. A +/// point is the cheapest well-formed variant (no allocation). +fn placeholder_geometry() -> Geometry { + Geometry::Point(geo_types::Point::new(0.0, 0.0)) +} + +/// 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: The row-loop view is a geometry slice. Its `ViewLen` is the slice length, so every index +// below that length is valid for both checked and unchecked access. +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_INFALLIBLE: bool = false; + + 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) + } + + // Every native geometry type decodes null rows to a placeholder. + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + + /// Null rows decode to [`placeholder_geometry`], which the branch-and-skip row loop never + /// reads. Null rows are filtered out before decoding, so their storage — which can hold + /// arbitrary payloads — is never interpreted; the placeholders keep the column indexed by row. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let valid = array.validity()?.execute_mask(array.len(), ctx)?; + if valid.all_true() { + return geometries(&array, ctx).map(Some); + } + + let mut decoded = geometries(&array.filter(valid.clone())?, ctx)?.into_iter(); + valid + .to_bit_buffer() + .iter() + .map(|is_valid| { + if is_valid { + decoded.next().ok_or_else(|| { + vortex_err!("spatial: filtered geometry column dropped a valid row") + }) + } else { + Ok(placeholder_geometry()) + } + }) + .collect::>>() + .map(Some) + } + + fn get(column: &Self::Column, index: usize) -> &Geometry { + &column[index] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column.as_slice() + } + + 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 `ViewLen::len` for this view, which + // is the geometry slice length. + unsafe { view.get_unchecked(index) } + } +} From 4193db32e1d6fdd0eda8c9c913df30b7b2b3de62 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 25 Aug 2026 11:46:03 -0400 Subject: [PATCH 2/2] fix(vortex-array): decouple row operation fallibility Signed-off-by: Nemo Yu --- vortex-array/src/scalar_fn/unstable/row/row_fn.rs | 7 +++++-- .../src/scalar_fn/unstable/row/types/element/input.rs | 3 +++ vortex-array/src/scalar_fn/unstable/row/visitor/check.rs | 6 ------ .../src/scalar_fn/unstable/row/visitor/row_visitor.rs | 3 ++- vortex-spatial/src/scalar_fn/distance.rs | 3 +-- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index 55087124b72..24c4ce5c0e9 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -47,12 +47,15 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// The arguments in display order. Its length is the function's exact arity. const ARG_NAMES: &'static [&'static str]; - /// Whether every dispatch is infallible. + /// Whether the row operation selected by every dispatch is infallible. /// /// See [`ScalarFnVTable::is_infallible`](crate::scalar_fn::ScalarFnVTable::is_infallible) for /// a more detailed explanation of semantic errors. /// - /// The framework checks dispatched element and result types. A conservative `false` is allowed. + /// Input decoding fallibility is declared separately by + /// [`InputElement::DECODE_INFALLIBLE`](crate::scalar_fn::unstable::row::InputElement::DECODE_INFALLIBLE) + /// and does not affect this flag. The framework checks dispatched result types. A conservative + /// `false` is allowed. const INFALLIBLE: bool; /// Returns the ID of the scalar function. diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 148f784bcba..e6e89fff220 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -47,6 +47,9 @@ pub unsafe trait InputElement: 'static { /// Whether [`decode`](Self::decode) is infallible for _legal_ input data. /// /// This excludes infrastructural failures such as IO or allocation. + /// It is independent of + /// [`RowFn::INFALLIBLE`](crate::scalar_fn::unstable::row::RowFn::INFALLIBLE), which describes + /// the row operation rather than input decoding. const DECODE_INFALLIBLE: bool; /// Validate that `dtype` is an acceptable input column dtype for this element type. diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs index cf5bf7ad44f..1089e71401c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -33,12 +33,6 @@ const fn assert_input_visit_contract() { Args::ARITY == F::ARG_NAMES.len(), "the visited argument tuple must have the arity declared by RowFn::ARG_NAMES", ); - // Dictionary push-down can evaluate values that no input row references. Every dispatch must - // therefore match the function-wide fallibility declaration. - assert!( - Args::DECODE_INFALLIBLE || !F::INFALLIBLE, - "RowFn::INFALLIBLE must be false when input decoding can fail", - ); } pub(super) const fn assert_owned_visit_contract() diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 4de9e34b330..288fea8415b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -22,7 +22,8 @@ use crate::scalar_fn::unstable::row::SinkResult; /// /// Only the framework implements this trait. The `visit_prepared*` methods derive shared state /// from constant arguments before visiting any rows. Every visit verifies that the argument tuple -/// matches [`RowFn::ARG_NAMES`] and that fallible decoding agrees with [`RowFn::INFALLIBLE`]. +/// matches [`RowFn::ARG_NAMES`] and that row-result fallibility agrees with +/// [`RowFn::INFALLIBLE`]. Input decoding declares its fallibility independently. /// /// A visit selects the _storage dtype_, the dtype the chosen [`OutputElement`] or [`OutputSink`] /// physically builds. [`with_output_dtype`](Self::with_output_dtype) declares the _output dtype_, diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index ad01a42f725..c78c8f89262 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -39,8 +39,7 @@ impl RowFn for SpatialDistance { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["a", "b"]; - // The per-row distance cannot fail, but decoding a geometry operand can. - const INFALLIBLE: bool = false; + const INFALLIBLE: bool = true; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.distance");