Execute spatial distance with RowFn

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/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<f64> {
+    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<bool> {
+    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::<Point>() || ext.is::<Polygon>())
+}
+
+/// 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<Option<Vec<Geometry<f64>>>> {
+    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::<ExtensionArray>(ctx)?
+        .storage_array()
+        .clone();
+
+    if ext.is::<Point>() {
+        point_geometries_null_tolerant(&storage, ctx).map(Some)
+    } else if ext.is::<Polygon>() {
+        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(

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<x, y[, z][, m]>` 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<Vec<Geometry<f64>>> {
+    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

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<List<Struct<x, y[, z][, m]>>>` (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<Vec<Geometry<f64>>> {
+    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<List<coordinate>>` storage.
 fn polygon_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<PolygonArray> {
     let polygon_type = polygon_type(

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<Option<Vec<u8>>> {
+    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
         Ok(Some(vec![]))
     }
 
-    fn deserialize(&self, _: &[u8], _: &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 deserialize(
         &self,
-        _: &Self::Options,
-        args: &dyn ExecutionArgs,
-        ctx: &mut ExecutionCtx,
-    ) -> VortexResult<ArrayRef> {
-        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<Self::Options> {
+        Ok(EmptyOptions)
     }
 
-    fn validity(
+    fn dispatch<V: RowVisitor<Self::Options>>(
         &self,
-        _: &Self::Options,
-        expression: &Expression,
-    ) -> VortexResult<Option<Expression>> {
-        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<V::VisitResult> {
+        visitor.visit_into::<(GeometryRow, GeometryRow), UninitElementSink<f64>, _>(
+            |(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();

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,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<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_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<Self::Column> {
+        geometries(&array, ctx)
+    }
+
+    fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult<bool> {
+        can_decode_geometries_null_tolerant(array)
+    }
+
+    fn get(column: &Self::Column, index: usize) -> &Geometry<f64> {
+        &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<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 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<Option<Self::Column>> {
+        geometries_null_tolerant(&array, ctx)
+    }
+}