refactor(vortex-spatial): execute distance with RowFn

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.

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.

vortex-array/src/scalar_fn/unstable/row/visitor/check.rs

@@ -33,12 +33,6 @@ const fn assert_input_visit_contract<F: RowFn, Args: ElementTuple>() {
         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<Function, Args, Out>()

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_,

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 }

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,37 @@ impl SpatialDistance {
     }
 }
 
-impl ScalarFnVTable for SpatialDistance {
+impl RowFn for SpatialDistance {
     type Options = EmptyOptions;
 
+    const ARG_NAMES: &'static [&'static str] = &["a", "b"];
+    const INFALLIBLE: bool = true;
+
     fn id(&self) -> ScalarFnId {
         static ID: CachedId = CachedId::new("vortex.st.distance");
         *ID
     }
 
-    fn serialize(&self, _: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
+    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
         Ok(Some(vec![]))
     }
 
-    fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult<Self::Options> {
+    fn deserialize(
+        &self,
+        _metadata: &[u8],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Options> {
         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<DType> {
-        validate_distance_operands(dtypes)?;
-        let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable));
-        Ok(DType::Primitive(PType::F64, nullability))
-    }
-
-    fn execute(
+    fn dispatch<V: RowVisitor>(
         &self,
-        _: &Self::Options,
-        args: &dyn ExecutionArgs,
-        ctx: &mut ExecutionCtx,
-    ) -> VortexResult<ArrayRef> {
-        let a = args.get(0)?;
-        let b = args.get(1)?;
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
         // 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<Option<Expression>> {
-        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 +89,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 +221,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<()> {

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;

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<f64> {
+    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<Geometry<f64>>;
+    type View<'a> = &'a [Geometry<f64>];
+    type Elem<'a> = &'a Geometry<f64>;
+
+    // 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<Self::Column> {
+        geometries(&array, ctx)
+    }
+
+    // Every native geometry type decodes null rows to a placeholder.
+    fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
+        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<Option<Self::Column>> {
+        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::<VortexResult<Vec<_>>>()
+            .map(Some)
+    }
+
+    fn get(column: &Self::Column, index: usize) -> &Geometry<f64> {
+        &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<f64>
+    where
+        Self: 'a,
+    {
+        &view[index]
+    }
+
+    unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry<f64>
+    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) }
+    }
+}