refactor(vortex-geo): generalize scalar function execution

vortex-geo/src/extension/mod.rs

@@ -61,7 +61,6 @@ use vortex_arrow::FromArrowArray;
 use vortex_buffer::Buffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
-use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
 pub use wkb::*;
 
@@ -78,21 +77,6 @@ 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. The two operands need not share a geometry type — e.g. a `Point` against
-/// a `Polygon` is valid, since distance/containment/intersection across types is meaningful.
-/// 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"
-        );
-    }
-    Ok(())
-}
-
 /// Flatten a native geometry column into a single coordinate `Struct<x, y, ...>` containing
 /// every vertex of every geometry.
 pub(crate) fn flatten_coordinates(

vortex-geo/src/extension/rect.rs

@@ -26,6 +26,7 @@ use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
 use vortex_array::arrays::ExtensionArray;
+use vortex_array::arrays::StructArray;
 use vortex_array::arrays::extension::ExtensionArrayExt;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldNames;
@@ -36,6 +37,7 @@ use vortex_array::dtype::extension::ExtDType;
 use vortex_array::dtype::extension::ExtId;
 use vortex_array::dtype::extension::ExtVTable;
 use vortex_array::scalar::ScalarValue;
+use vortex_array::validity::Validity;
 use vortex_arrow::ArrowExport;
 use vortex_arrow::ArrowExportVTable;
 use vortex_arrow::ArrowImport;
@@ -155,6 +157,24 @@ pub(crate) fn box_dimension(dtype: &DType) -> VortexResult<Dimension> {
         .ok_or_else(|| vortex_err!("not a valid geoarrow.box dimension: {:?}", fields.names()))
 }
 
+/// Build a native [`Rect`] array from canonical min/max ordinate columns.
+pub(crate) fn build_rect_array(
+    ext_dtype: &ExtDType<Rect>,
+    corners: Vec<ArrayRef>,
+    len: usize,
+    validity: Validity,
+) -> VortexResult<ArrayRef> {
+    let dimension = box_dimension(ext_dtype.storage_dtype())?;
+    let storage = StructArray::try_new(
+        FieldNames::from(box_field_names(dimension)),
+        corners,
+        len,
+        validity,
+    )?
+    .into_array();
+    Ok(ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array())
+}
+
 static ARROW_BOX: CachedId = CachedId::new(BoxType::NAME);
 
 /// The `geoarrow.box` extension type for `dimension`.

vortex-geo/src/scalar_fn/contains.rs

@@ -19,11 +19,28 @@ 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;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
-use crate::extension::validate_geometry_operands;
-use crate::scalar_fn::execute::execute_null_propagating;
+use crate::extension::is_native_geometry;
+use crate::scalar_fn::execute::execute_binary_geo_types;
+
+/// Validate the two native geometry operands accepted by `ST_Contains`.
+fn validate_contains_operands(dtypes: &[DType]) -> VortexResult<()> {
+    vortex_ensure!(
+        dtypes.len() == 2,
+        "geo: contains requires exactly two geometry operands, got {}",
+        dtypes.len()
+    );
+    for dtype in dtypes {
+        vortex_ensure!(
+            is_native_geometry(dtype),
+            "geo: contains operand {dtype} is not a native geometry type"
+        );
+    }
+    Ok(())
+}
 
 /// 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
@@ -71,7 +88,7 @@ impl ScalarFnVTable for GeoContains {
     }
 
     fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult<DType> {
-        validate_geometry_operands(dtypes)?;
+        validate_contains_operands(dtypes)?;
         let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable));
         Ok(DType::Bool(nullability))
     }
@@ -87,7 +104,7 @@ impl ScalarFnVTable for GeoContains {
         // Containment is not symmetric: `a` is always the container and `b` the contained. A
         // container's rect must cover the contained's rect (`Rect::contains` is the closed
         // test), so a contained rect poking outside proves the row false.
-        execute_null_propagating(
+        execute_binary_geo_types(
             &a,
             &b,
             |a, b| a.contains(b),

vortex-geo/src/scalar_fn/distance.rs

@@ -21,11 +21,28 @@ 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;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
-use crate::extension::validate_geometry_operands;
-use crate::scalar_fn::execute::execute_null_propagating;
+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,
+        "geo: distance requires exactly two geometry operands, got {}",
+        dtypes.len()
+    );
+    for dtype in dtypes {
+        vortex_ensure!(
+            is_native_geometry(dtype),
+            "geo: distance operand {dtype} is not a native geometry type"
+        );
+    }
+    Ok(())
+}
 
 /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry
 /// operands, each a column or a constant literal.
@@ -72,7 +89,7 @@ impl ScalarFnVTable for GeoDistance {
     }
 
     fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult<DType> {
-        validate_geometry_operands(dtypes)?;
+        validate_distance_operands(dtypes)?;
         let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable));
         Ok(DType::Primitive(PType::F64, nullability))
     }
@@ -86,7 +103,7 @@ impl ScalarFnVTable for GeoDistance {
         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_null_propagating(&a, &b, |x, y| Euclidean.distance(x, y), None, ctx)
+        execute_binary_geo_types(&a, &b, |x, y| Euclidean.distance(x, y), None, ctx)
     }
 
     fn validity(