perf(rust): point-coordinate fast paths for ST_Distance/ST_DWithin/ST_Azimuth

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

@@ -206,6 +206,20 @@ fn criterion_benchmark(c: &mut Criterion) {
         "st_azimuth",
         BenchmarkArgs::ArrayArray(Point, Point),
     );
+    benchmark::scalar(
+        c,
+        &f,
+        "native",
+        "st_azimuth",
+        BenchmarkArgs::ArrayScalar(Point, Point),
+    );
+    benchmark::scalar(
+        c,
+        &f,
+        "native",
+        "st_azimuth",
+        BenchmarkArgs::ScalarArray(Point, Point),
+    );
 
     benchmark::aggregate(c, &f, "native", "st_envelope_agg", Point);
     benchmark::aggregate(c, &f, "native", "st_envelope_agg", LineString(10));

rust/sedona-functions/src/executor.rs

@@ -23,8 +23,9 @@ use datafusion_common::error::Result;
 use datafusion_common::{DataFusionError, ScalarValue};
 use datafusion_expr::ColumnarValue;
 use sedona_common::sedona_internal_err;
+use sedona_geometry::wkb_header::WkbPointLayout;
 use sedona_schema::datatypes::SedonaType;
-use wkb::reader::Wkb;
+use wkb::reader::{read_wkb, Wkb};
 
 /// Helper for writing general kernel implementations with geometry
 ///
@@ -269,6 +270,51 @@ impl GeometryFactory for WkbBytesFactory {
 /// the [WkbExecutor].
 pub type WkbBytesExecutor<'a, 'b> = GenericExecutor<'a, 'b, WkbBytesFactory, WkbBytesFactory>;
 
+/// A finite Point's `(x, y)`, or the parsed geometry for anything else.
+///
+/// Produced by [PointXYGeometryFactory] for kernels whose hot path only needs a
+/// Point's coordinates (ST_X, ST_Distance, ST_DWithin, …): a Point is decoded
+/// straight from its fixed WKB offset (no full parse), while every other
+/// geometry — including `POINT EMPTY` — falls back to a parsed [Wkb] so the
+/// kernel can handle it with the general path.
+pub enum PointOrWkb<'a> {
+    /// A finite Point, decoded from its fixed WKB offset without a full parse.
+    Point(f64, f64),
+    /// Any non-Point (or `POINT EMPTY`), fully parsed.
+    Other(Wkb<'a>),
+}
+
+/// A [GeometryFactory] that decodes a finite Point's `(x, y)` from its fixed WKB
+/// offset and parses anything else into a [Wkb].
+///
+/// Because [GenericExecutor] parses each scalar argument once (and each array
+/// element per row), a kernel built on this factory gets parse-once-for-scalars
+/// for free and only pays a coordinate read where the Point fast path applies.
+#[derive(Default)]
+pub struct PointXYGeometryFactory {}
+
+impl GeometryFactory for PointXYGeometryFactory {
+    type Geom<'a> = PointOrWkb<'a>;
+
+    #[inline]
+    fn try_from_wkb<'a>(&self, wkb_bytes: &'a [u8]) -> Result<Self::Geom<'a>> {
+        // A finite Point: read its coordinate from the fixed offset, no parse.
+        // POINT EMPTY (NaN/NaN) and any header error fall through to the parser.
+        if let Ok(Some(layout)) = WkbPointLayout::try_from_wkb(wkb_bytes) {
+            if let Ok(Some((x, y))) = layout.read_xy(wkb_bytes) {
+                return Ok(PointOrWkb::Point(x, y));
+            }
+        }
+        let wkb = read_wkb(wkb_bytes).map_err(|e| DataFusionError::External(Box::new(e)))?;
+        Ok(PointOrWkb::Other(wkb))
+    }
+}
+
+/// Alias for an executor that iterates geometries as [PointOrWkb] — the Point
+/// fast path for point-coordinate kernels.
+pub type PointXYExecutor<'a, 'b> =
+    GenericExecutor<'a, 'b, PointXYGeometryFactory, PointXYGeometryFactory>;
+
 /// Trait for iterating over a container type as geometry scalars
 ///
 /// Currently the only scalar type supported is [Wkb]; however, for future

rust/sedona-functions/src/st_azimuth.rs

@@ -18,16 +18,15 @@ use arrow_array::builder::Float64Builder;
 use arrow_schema::DataType;
 use datafusion_common::{error::Result, exec_err};
 use datafusion_expr::{ColumnarValue, Volatility};
-use geo_traits::{CoordTrait, GeometryTrait, GeometryType, PointTrait};
 use sedona_expr::{
     item_crs::ItemCrsKernel,
     scalar_udf::{SedonaScalarKernel, SedonaScalarUDF},
 };
+use sedona_geometry::wkb_header::read_point_xy;
 use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher};
 use std::sync::Arc;
-use wkb::reader::Wkb;
 
-use crate::executor::WkbExecutor;
+use crate::executor::WkbBytesExecutor;
 
 /// ST_Azimuth() scalar UDF
 ///
@@ -58,14 +57,17 @@ impl SedonaScalarKernel for STAzimuth {
         arg_types: &[SedonaType],
         args: &[ColumnarValue],
     ) -> Result<ColumnarValue> {
-        let executor = WkbExecutor::new(arg_types, args);
+        // ST_Azimuth is defined only for two non-empty POINTs, so it never needs
+        // the full geometry: read each operand's (x, y) straight from the WKB
+        // offset via `read_point_xy` (no parse, no enum). POINT EMPTY or any
+        // non-Point maps to the same error the general path raised.
+        let executor = WkbBytesExecutor::new(arg_types, args);
         let mut builder = Float64Builder::with_capacity(executor.num_iterations());
         executor.execute_wkb_wkb_void(|maybe_start, maybe_end| {
             match (maybe_start, maybe_end) {
-                (Some(start), Some(end)) => match invoke_scalar(start, end)? {
-                    Some(angle) => builder.append_value(angle),
-                    None => builder.append_null(),
-                },
+                (Some(start), Some(end)) => {
+                    builder.append_option(azimuth_from_bytes(start, end)?);
+                }
                 _ => builder.append_null(),
             }
 
@@ -76,22 +78,13 @@ impl SedonaScalarKernel for STAzimuth {
     }
 }
 
-fn invoke_scalar(start: &Wkb, end: &Wkb) -> Result<Option<f64>> {
-    match (start.as_type(), end.as_type()) {
-        (GeometryType::Point(start_point), GeometryType::Point(end_point)) => {
-            match (start_point.coord(), end_point.coord()) {
-                // If both geometries are non-empty points, calculate the angle
-                (Some(start_coord), Some(end_coord)) => Ok(calc_azimuth(
-                    start_coord.x(),
-                    start_coord.y(),
-                    end_coord.x(),
-                    end_coord.y(),
-                )),
-                // If either of the points is empty, raise an error.
-                _ => {
-                    exec_err!("ST_Azimuth expects both arguments to be non-empty POINT geometries")
-                }
-            }
+/// Azimuth between two operands read as raw WKB. Both must be non-empty POINTs;
+/// `read_point_xy` returns `Ok(None)` for POINT EMPTY and `Err` for non-Point
+/// input, both of which map to the same error the general parser raised.
+fn azimuth_from_bytes(start: &[u8], end: &[u8]) -> Result<Option<f64>> {
+    match (read_point_xy(start), read_point_xy(end)) {
+        (Ok(Some((start_x, start_y))), Ok(Some((end_x, end_y)))) => {
+            Ok(calc_azimuth(start_x, start_y, end_x, end_y))
         }
         _ => exec_err!("ST_Azimuth expects both arguments to be non-empty POINT geometries"),
     }

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

@@ -60,7 +60,12 @@ fn criterion_benchmark(c: &mut Criterion) {
         ArrayScalar(Point, Polygon(500)),
     );
 
+    benchmark::scalar(c, &f, "geo", "st_distance", ArrayArray(Point, Point));
+    benchmark::scalar(c, &f, "geo", "st_distance", ArrayScalar(Point, Point));
+    benchmark::scalar(c, &f, "geo", "st_distance", ScalarArray(Point, Point));
     benchmark::scalar(c, &f, "geo", "st_distance", ArrayScalar(Point, Polygon(10)));
+    benchmark::scalar(c, &f, "geo", "st_distance", ArrayScalar(Polygon(10), Point));
+    benchmark::scalar(c, &f, "geo", "st_distance", ArrayArray(Point, Polygon(10)));
     benchmark::scalar(
         c,
         &f,
@@ -83,6 +88,21 @@ fn criterion_benchmark(c: &mut Criterion) {
         "st_dwithin",
         ArrayArrayScalar(Polygon(10), Polygon(500), Float64(1.0, 2.0)),
     );
+    // Point/point — the fast-path case (and the spatial-join shape).
+    benchmark::scalar(
+        c,
+        &f,
+        "geo",
+        "st_dwithin",
+        ArrayArrayScalar(Point, Point, Float64(1.0, 2.0)),
+    );
+    benchmark::scalar(
+        c,
+        &f,
+        "geo",
+        "st_dwithin",
+        ArrayArrayScalar(Point, Polygon(10), Float64(1.0, 2.0)),
+    );
 }
 
 fn criterion_benchmark_aggr(c: &mut Criterion) {

rust/sedona-geo/src/st_distance.rs

@@ -20,14 +20,14 @@ use arrow_array::builder::Float64Builder;
 use arrow_schema::DataType;
 use datafusion_common::error::Result;
 use datafusion_expr::ColumnarValue;
+use geo_types::Point;
 use sedona_expr::{
     item_crs::ItemCrsKernel,
     scalar_udf::{ScalarKernelRef, SedonaScalarKernel},
 };
-use sedona_functions::executor::WkbExecutor;
+use sedona_functions::executor::{PointOrWkb, PointXYExecutor};
 use sedona_geo_generic_alg::line_measures::DistanceExt;
 use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher};
-use wkb::reader::Wkb;
 
 /// ST_Distance() implementation using [DistanceExt]
 pub fn st_distance_impl() -> Vec<ScalarKernelRef> {
@@ -52,42 +52,60 @@ impl SedonaScalarKernel for STDistance {
         arg_types: &[SedonaType],
         args: &[ColumnarValue],
     ) -> Result<ColumnarValue> {
-        let executor = WkbExecutor::new(arg_types, args);
+        // The PointXY executor decodes each operand as a `PointOrWkb`: a finite
+        // Point yields its `(x, y)` straight from the WKB offset (no parse),
+        // anything else is parsed. Scalars are decoded once and reused per row.
+        let executor = PointXYExecutor::new(arg_types, args);
         let mut builder = Float64Builder::with_capacity(executor.num_iterations());
-        executor.execute_wkb_wkb_void(|maybe_wkb0, maybe_wkb1| {
-            match (maybe_wkb0, maybe_wkb1) {
-                (Some(wkb0), Some(wkb1)) => {
-                    builder.append_value(invoke_scalar(wkb0, wkb1)?);
-                }
+
+        executor.execute_wkb_wkb_void(|maybe0, maybe1| {
+            match (maybe0, maybe1) {
+                (Some(a), Some(b)) => builder.append_value(point_or_wkb_distance(a, b)),
                 _ => builder.append_null(),
             }
-
             Ok(())
         })?;
 
         executor.finish(Arc::new(builder.finish()))
     }
 }
 
-fn invoke_scalar(wkb_a: &Wkb, wkb_b: &Wkb) -> Result<f64> {
-    Ok(wkb_a.distance_ext(wkb_b))
+/// Euclidean distance between two operands. Point-to-Point is computed directly
+/// from the coordinates; any other combination defers to the generic
+/// [DistanceExt] (a Point operand becomes a [`geo_types::Point`] so the mixed
+/// Point/geometry case never re-parses the Point). Shared with `ST_DWithin`.
+pub(crate) fn point_or_wkb_distance(a: &PointOrWkb, b: &PointOrWkb) -> f64 {
+    match (a, b) {
+        (PointOrWkb::Point(ax, ay), PointOrWkb::Point(bx, by)) => {
+            let dx = ax - bx;
+            let dy = ay - by;
+            (dx * dx + dy * dy).sqrt()
+        }
+        (PointOrWkb::Point(x, y), PointOrWkb::Other(wkb)) => Point::new(*x, *y).distance_ext(wkb),
+        (PointOrWkb::Other(wkb), PointOrWkb::Point(x, y)) => wkb.distance_ext(&Point::new(*x, *y)),
+        (PointOrWkb::Other(wkb0), PointOrWkb::Other(wkb1)) => wkb0.distance_ext(wkb1),
+    }
 }
 
 #[cfg(test)]
 mod tests {
+    use arrow_array::{ArrayRef, Float64Array};
     use datafusion_common::scalar::ScalarValue;
     use rstest::rstest;
     use sedona_expr::scalar_udf::SedonaScalarUDF;
     use sedona_schema::datatypes::{WKB_GEOMETRY, WKB_GEOMETRY_ITEM_CRS, WKB_VIEW_GEOMETRY};
-    use sedona_testing::create::create_scalar;
+    use sedona_testing::compare::assert_array_equal;
+    use sedona_testing::create::{create_array, create_scalar};
     use sedona_testing::testers::ScalarUdfTester;
 
     use super::*;
 
     #[rstest]
     fn udf(
-        #[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY, WKB_GEOMETRY_ITEM_CRS.clone())] left_sedona_type: SedonaType,
-        #[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY, WKB_GEOMETRY_ITEM_CRS.clone())] right_sedona_type: SedonaType,
+        #[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY, WKB_GEOMETRY_ITEM_CRS.clone())]
+        left_sedona_type: SedonaType,
+        #[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY, WKB_GEOMETRY_ITEM_CRS.clone())]
+        right_sedona_type: SedonaType,
     ) {
         let udf = SedonaScalarUDF::from_impl("st_distance", st_distance_impl());
         let tester = ScalarUdfTester::new(
@@ -123,4 +141,65 @@ mod tests {
             .unwrap();
         assert!(result.is_null());
     }
+
+    /// Exercises the array/scalar fast-path routing: a scalar operand is parsed
+    /// once and reused per row, the point-to-point fast path matches the general
+    /// distance, and mixed Point/non-Point inputs fall back correctly. Covered
+    /// across storage encodings, including item-CRS (which arrives unwrapped as
+    /// plain WKB, so the parse-once scalar handling still applies).
+    #[rstest]
+    fn array_scalar_routing(
+        #[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY, WKB_GEOMETRY_ITEM_CRS.clone())] left: SedonaType,
+        #[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY, WKB_GEOMETRY_ITEM_CRS.clone())] right: SedonaType,
+    ) {
+        let udf = SedonaScalarUDF::from_impl("st_distance", st_distance_impl());
+        let tester = ScalarUdfTester::new(udf.into(), vec![left.clone(), right.clone()]);
+
+        // Array of Points vs a scalar Point: the point-to-point fast path.
+        let result = tester
+            .invoke_wkb_array_scalar(
+                vec![Some("POINT (3 4)"), None, Some("POINT (0 0)")],
+                create_scalar(Some("POINT (0 0)"), &right),
+            )
+            .unwrap();
+        let expected: ArrayRef = Arc::new(Float64Array::from(vec![Some(5.0), None, Some(0.0)]));
+        assert_array_equal(&result, &expected);
+
+        // Array of Polygons vs a scalar Point: the scalar is parsed once and
+        // reused across rows; each row falls back to the general distance. The
+        // scalar Point sits on a polygon vertex, so every distance is 0.
+        let result = tester
+            .invoke_wkb_array_scalar(
+                vec![
+                    Some("POLYGON ((10 10, 10 20, 20 20, 20 10, 10 10))"),
+                    Some("POLYGON ((10 10, 10 20, 20 20, 20 10, 10 10))"),
+                ],
+                create_scalar(Some("POINT (10 10)"), &right),
+            )
+            .unwrap();
+        let expected: ArrayRef = Arc::new(Float64Array::from(vec![Some(0.0), Some(0.0)]));
+        assert_array_equal(&result, &expected);
+
+        // Scalar Point vs array of Points (mirror orientation).
+        let result = tester
+            .invoke_scalar_array(
+                create_scalar(Some("POINT (0 0)"), &left),
+                create_array(&[Some("POINT (3 4)"), Some("POINT (0 0)")], &right),
+            )
+            .unwrap();
+        let expected: ArrayRef = Arc::new(Float64Array::from(vec![Some(5.0), Some(0.0)]));
+        assert_array_equal(&result, &expected);
+
+        // Array vs array with mixed Point / Polygon rows.
+        let arg0 = create_array(
+            &[
+                Some("POINT (0 0)"),
+                Some("POLYGON ((10 10, 10 20, 20 20, 20 10, 10 10))"),
+            ],
+            &left,
+        );
+        let arg1 = create_array(&[Some("POINT (3 4)"), Some("POINT (10 10)")], &right);
+        let expected: ArrayRef = Arc::new(Float64Array::from(vec![Some(5.0), Some(0.0)]));
+        assert_array_equal(&tester.invoke_array_array(arg0, arg1).unwrap(), &expected);
+    }
 }