feat(vortex-geo): null-propagating scalar functions

Make the binary geo scalar functions (ST_Distance, ST_Intersects,
ST_Contains) null-propagating, matching SQL/OGC semantics and the other
Vortex binary kernels: a nullable geometry operand is now allowed, and any
row whose geometry input is null yields a null result (output nullable iff an
operand is nullable).

A geo kernel decodes each operand into a geo_types geometry, and a null row
has no geometry to decode, so it can't compute over every row and mask the
nulls afterwards the way numeric kernels do. The shared execute_null_propagating
dispatch in scalar_fn/mod.rs instead filters the null rows out up front,
computes over the rows valid in both operands, and scatters the results back
under the combined null mask.

- validate_geometry_operands no longer rejects nullable operands
- return_dtype mirrors operand nullability
- tests for nullable columns, constant-null operands, and column/column nulls

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

Author: HarukiMoriarty

Cargo.lock

@@ -24,7 +24,9 @@ geoarrow-cast = { workspace = true }
 prost = { workspace = true }
 vortex-array = { workspace = true }
 vortex-arrow = { workspace = true }
+vortex-buffer = { workspace = true }
 vortex-error = { workspace = true }
+vortex-mask = { workspace = true }
 vortex-session = { workspace = true }
 wkb = { workspace = true }
 

vortex-geo/Cargo.toml

@@ -72,18 +72,15 @@ pub(crate) fn is_native_geometry(dtype: &DType) -> bool {
     })
 }
 
-/// Validate the operands of a geo scalar function: each must be a native geometry type (so the
-/// kernel can decode it) and non-nullable (geometry arrays never carry nulls).
+/// Validate the operands of a geo scalar function: each must be a native geometry type so the
+/// kernel can decode it. Nullable operands are allowed; the kernels propagate nulls (a null
+/// geometry input yields a null result) rather than decoding null rows.
 pub(crate) fn validate_geometry_operands(dtypes: &[DType]) -> VortexResult<()> {
     for dtype in dtypes {
         vortex_ensure!(
             is_native_geometry(dtype),
             "geo: operand {dtype} is not a native geometry type"
         );
-        vortex_ensure!(
-            !dtype.is_nullable(),
-            "geo: nullable operand {dtype} is unsupported"
-        );
     }
     Ok(())
 }

vortex-geo/src/extension/mod.rs

@@ -6,14 +6,9 @@
 use geo::Contains;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
-use vortex_array::IntoArray;
-use vortex_array::arrays::BoolArray;
-use vortex_array::arrays::Constant;
-use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::ScalarFnArray;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
-use vortex_array::scalar::Scalar;
 use vortex_array::scalar_fn::Arity;
 use vortex_array::scalar_fn::ChildName;
 use vortex_array::scalar_fn::EmptyOptions;
@@ -22,13 +17,11 @@ 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_eq;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
-use crate::extension::geometries;
-use crate::extension::single_geometry;
 use crate::extension::validate_geometry_operands;
+use crate::scalar_fn::execute_null_propagating;
 
 /// OGC `ST_Contains` between two native geometry operands, each a column or a constant
 /// literal: true where operand `b` lies completely inside operand `a` (boundary contact alone
@@ -77,7 +70,8 @@ impl ScalarFnVTable for GeoContains {
 
     fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult<DType> {
         validate_geometry_operands(dtypes)?;
-        Ok(DType::Bool(Nullability::NonNullable))
+        let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable));
+        Ok(DType::Bool(nullability))
     }
 
     fn execute(
@@ -89,59 +83,10 @@ impl ScalarFnVTable for GeoContains {
         let a = args.get(0)?;
         let b = args.get(1)?;
         // Containment is not symmetric: `a` is always the container and `b` the contained.
-        match (a.as_opt::<Constant>(), b.as_opt::<Constant>()) {
-            (Some(qa), Some(qb)) => {
-                let ga = single_geometry(qa.scalar(), ctx)?;
-                let gb = single_geometry(qb.scalar(), ctx)?;
-                Ok(ConstantArray::new(
-                    Scalar::bool(ga.contains(&gb), Nullability::NonNullable),
-                    a.len(),
-                )
-                .into_array())
-            }
-            (Some(qa), None) => constant_contains_column(qa.scalar(), &b, ctx),
-            (None, Some(qb)) => column_contains_constant(&a, qb.scalar(), ctx),
-            (None, None) => {
-                vortex_ensure_eq!(
-                    a.len(),
-                    b.len(),
-                    "geo contains: operand length mismatch {} vs {}",
-                    a.len(),
-                    b.len()
-                );
-                let ag = geometries(&a, ctx)?;
-                let bg = geometries(&b, ctx)?;
-                let hits = ag.iter().zip(&bg).map(|(x, y)| x.contains(y));
-                Ok(BoolArray::from_iter(hits).into_array())
-            }
-        }
+        execute_null_propagating(&a, &b, |a, b| a.contains(b), ctx)
     }
 }
 
-/// Whether the constant `container` contains each row of `contained`.
-fn constant_contains_column(
-    container: &Scalar,
-    contained: &ArrayRef,
-    ctx: &mut ExecutionCtx,
-) -> VortexResult<ArrayRef> {
-    let container = single_geometry(container, ctx)?;
-    let geoms = geometries(contained, ctx)?;
-    let hits = geoms.iter().map(|g| container.contains(g));
-    Ok(BoolArray::from_iter(hits).into_array())
-}
-
-/// Whether each row of `container` contains the constant `contained`.
-fn column_contains_constant(
-    container: &ArrayRef,
-    contained: &Scalar,
-    ctx: &mut ExecutionCtx,
-) -> VortexResult<ArrayRef> {
-    let contained = single_geometry(contained, ctx)?;
-    let geoms = geometries(container, ctx)?;
-    let hits = geoms.iter().map(|g| g.contains(&contained));
-    Ok(BoolArray::from_iter(hits).into_array())
-}
-
 #[cfg(test)]
 mod tests {
     use geo_types::Geometry;
@@ -160,13 +105,17 @@ mod tests {
     use vortex_array::dtype::DType;
     use vortex_array::dtype::Nullability;
     use vortex_array::dtype::PType;
+    use vortex_array::scalar::Scalar;
     use vortex_array::scalar_fn::EmptyOptions;
     use vortex_array::scalar_fn::ScalarFnVTable;
+    use vortex_array::validity::Validity;
+    use vortex_buffer::BitBuffer;
     use vortex_error::VortexResult;
     use vortex_error::vortex_err;
     use wkb::writer::WriteOptions;
 
     use super::GeoContains;
+    use crate::test_harness::nullable_point_column;
     use crate::test_harness::point_column;
 
     /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes.
@@ -286,12 +235,53 @@ mod tests {
         assert_contains(polygons, points, [true, false])
     }
 
-    /// Geometry arrays are never nullable, so a nullable operand dtype is rejected.
+    /// Output nullability mirrors the operands: nullable if any operand is nullable, otherwise
+    /// non-nullable.
     #[test]
-    fn nullable_operand_is_rejected() -> VortexResult<()> {
+    fn output_nullability_mirrors_operands() -> VortexResult<()> {
         let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone();
-        let result = GeoContains.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]);
-        assert!(result.is_err());
+        let non_nullable =
+            GeoContains.return_dtype(&EmptyOptions, &[dtype.clone(), dtype.clone()])?;
+        assert!(!non_nullable.is_nullable());
+        let nullable = GeoContains.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype])?;
+        assert!(nullable.is_nullable());
+        Ok(())
+    }
+
+    /// A null row in the contained operand yields a null verdict; valid rows keep their verdict
+    /// (a strictly interior point is contained, an outside point is not).
+    #[test]
+    fn contains_propagates_null_rows() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?;
+        let points = nullable_point_column(vec![Some((2.0, 2.0)), None, Some((10.0, 10.0))])?;
+        let contains = GeoContains::try_new_array(container, points)?.into_array();
+
+        let expected = BoolArray::new(
+            BitBuffer::from_iter([true, false, false]),
+            Validity::from_iter([true, false, true]),
+        )
+        .into_array();
+        assert_arrays_eq!(contains, expected, &mut ctx);
+        Ok(())
+    }
+
+    /// A constant-null operand produces an all-null output.
+    #[test]
+    fn contains_constant_null_is_all_null() -> VortexResult<()> {
+        let session = vortex_array::array_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable();
+        let null_const = ConstantArray::new(Scalar::null(point_dtype), 2).into_array();
+        let points = point_column(vec![2.0, 10.0], vec![2.0, 10.0])?;
+        let contains = GeoContains::try_new_array(null_const, points)?.into_array();
+
+        let expected =
+            BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array();
+        assert_arrays_eq!(contains, expected, &mut ctx);
         Ok(())
     }