ST_Distance and ST_DWithin based on georust/geo

c/sedona-geos/src/st_dwithin.rs

@@ -18,7 +18,7 @@ use std::sync::Arc;
 
 use arrow_array::builder::BooleanBuilder;
 use arrow_schema::DataType;
-use datafusion_common::{error::Result, DataFusionError};
+use datafusion_common::{cast::as_float64_array, error::Result, DataFusionError};
 use datafusion_expr::ColumnarValue;
 use geos::Geom;
 use sedona_expr::scalar_udf::{ScalarKernelRef, SedonaScalarKernel};
@@ -53,26 +53,14 @@ impl SedonaScalarKernel for STDWithin {
         arg_types: &[SedonaType],
         args: &[ColumnarValue],
     ) -> Result<ColumnarValue> {
-        // Extract the constant scalar value before looping over the input geometries
-        let distance: Option<f64>;
         let arg2 = args[2].cast_to(&DataType::Float64, None)?;
-        if let ColumnarValue::Scalar(scalar_arg) = &arg2 {
-            if scalar_arg.is_null() {
-                distance = None;
-            } else {
-                distance = Some(f64::try_from(scalar_arg.clone())?);
-            }
-        } else {
-            return Err(DataFusionError::Execution(format!(
-                "Invalid distance: {:?}",
-                args[2]
-            )));
-        }
-
         let executor = GeosExecutor::new(arg_types, args);
+        let arg2_array = arg2.to_array(executor.num_iterations())?;
+        let arg2_f64_array = as_float64_array(&arg2_array)?;
+        let mut arg2_iter = arg2_f64_array.iter();
         let mut builder = BooleanBuilder::with_capacity(executor.num_iterations());
         executor.execute_wkb_wkb_void(|lhs, rhs| {
-            match (lhs, rhs, distance) {
+            match (lhs, rhs, arg2_iter.next().unwrap()) {
                 (Some(lhs), Some(rhs), Some(distance)) => {
                     builder.append_value(invoke_scalar(lhs, rhs, distance)?);
                 }
@@ -151,9 +139,16 @@ mod tests {
         let expected: ArrayRef = arrow_array!(Boolean, [Some(true), Some(false), None, Some(true)]);
         assert_array_equal(
             &tester
-                .invoke_array_array_scalar(arg1, arg2, distance)
+                .invoke_array_array_scalar(Arc::clone(&arg1), Arc::clone(&arg2), distance)
                 .unwrap(),
             &expected,
         );
+
+        let distance = arrow_array!(Int32, [Some(1), Some(1), Some(1), Some(1)]);
+        let expected: ArrayRef = arrow_array!(Boolean, [Some(true), Some(false), None, Some(true)]);
+        assert_array_equal(
+            &tester.invoke_arrays(vec![arg1, arg2, distance]).unwrap(),
+            &expected,
+        );
     }
 }

rust/sedona-geo/benches/geo-functions.rs

@@ -41,6 +41,30 @@ fn criterion_benchmark(c: &mut Criterion) {
         "st_intersects",
         ArrayScalar(Point, Polygon(500)),
     );
+
+    benchmark::scalar(c, &f, "geo", "st_distance", ArrayScalar(Point, Polygon(10)));
+    benchmark::scalar(
+        c,
+        &f,
+        "geo",
+        "st_distance",
+        ArrayScalar(Point, Polygon(500)),
+    );
+
+    benchmark::scalar(
+        c,
+        &f,
+        "geo",
+        "st_dwithin",
+        ArrayArrayScalar(Polygon(10), Polygon(10), Float64(1.0, 2.0)),
+    );
+    benchmark::scalar(
+        c,
+        &f,
+        "geo",
+        "st_dwithin",
+        ArrayArrayScalar(Polygon(10), Polygon(500), Float64(1.0, 2.0)),
+    );
 }
 
 fn criterion_benchmark_aggr(c: &mut Criterion) {

rust/sedona-geo/src/lib.rs

@@ -18,6 +18,8 @@ pub mod centroid;
 pub mod register;
 mod st_area;
 mod st_centroid;
+mod st_distance;
+mod st_dwithin;
 mod st_intersection_aggr;
 mod st_intersects;
 mod st_length;

rust/sedona-geo/src/register.rs

@@ -21,15 +21,17 @@ use crate::st_intersection_aggr::st_intersection_aggr_impl;
 use crate::st_line_interpolate_point::st_line_interpolate_point_impl;
 use crate::st_union_aggr::st_union_aggr_impl;
 use crate::{
-    st_area::st_area_impl, st_centroid::st_centroid_impl, st_intersects::st_intersects_impl,
-    st_length::st_length_impl,
+    st_area::st_area_impl, st_centroid::st_centroid_impl, st_distance::st_distance_impl,
+    st_dwithin::st_dwithin_impl, st_intersects::st_intersects_impl, st_length::st_length_impl,
 };
 
 pub fn scalar_kernels() -> Vec<(&'static str, ScalarKernelRef)> {
     vec![
         ("st_intersects", st_intersects_impl()),
         ("st_area", st_area_impl()),
         ("st_centroid", st_centroid_impl()),
+        ("st_distance", st_distance_impl()),
+        ("st_dwithin", st_dwithin_impl()),
         ("st_length", st_length_impl()),
         ("st_lineinterpolatepoint", st_line_interpolate_point_impl()),
     ]