perf(rust/sedona-geo): PointXYGeometryFactory point fast path for ST_Distance/ST_DWithin

james-willis · james-willis · commit 7f337d11ba69 · 2026-06-30T15:54:22.000-07:00
Add PointXYGeometryFactory to sedona-functions: a GeometryFactory whose Geom is
a PointOrWkb — a finite Point decoded straight from its fixed WKB offset (no full
parse), or a parsed Wkb for anything else (incl POINT EMPTY). The GenericExecutor
parses each scalar once, so kernels get parse-once-for-scalars for free.

Adopt it in ST_Distance (Point/Point computes Euclidean directly; mixed
Point/geometry uses a geo_types::Point so the Point is never re-parsed) and
ST_DWithin (shares the same backend, distance &lt;= bound). Replaces the
hand-rolled PreparedScalar scalar-prep. Adds point/point ST_DWithin benchmarks.

ST_X/ST_Y deliberately stay on their direct read_point_xy fast path: the factory
enum adds per-element overhead that regresses the already-optimal unary case
(~+30% in benchmarks), with no parse savings to offset it.
diff --git a/rust/sedona-functions/src/executor.rs b/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
diff --git a/rust/sedona-geo/benches/geo-functions.rs b/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) {
diff --git a/rust/sedona-geo/src/st_distance.rs b/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);
+    }
 }
diff --git a/rust/sedona-geo/src/st_dwithin.rs b/rust/sedona-geo/src/st_dwithin.rs
@@ -16,6 +16,7 @@
 // under the License.
 use std::sync::Arc;
 
+use crate::st_distance::point_or_wkb_distance;
 use arrow_array::builder::BooleanBuilder;
 use arrow_schema::DataType;
 use datafusion_common::{cast::as_float64_array, error::Result};
@@ -24,10 +25,8 @@ use sedona_expr::{
     item_crs::ItemCrsKernel,
     scalar_udf::{ScalarKernelRef, SedonaScalarKernel},
 };
-use sedona_functions::executor::WkbExecutor;
-use sedona_geo_generic_alg::line_measures::DistanceExt;
+use sedona_functions::executor::PointXYExecutor;
 use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher};
-use wkb::reader::Wkb;
 
 /// ST_DWithin() implementation using [DistanceExt]
 pub fn st_dwithin_impl() -> Vec<ScalarKernelRef> {
@@ -57,15 +56,17 @@ impl SedonaScalarKernel for STDWithin {
         args: &[ColumnarValue],
     ) -> Result<ColumnarValue> {
         let arg2 = args[2].cast_to(&DataType::Float64, None)?;
-        let executor = WkbExecutor::new(arg_types, args);
+        // Same PointXY fast path as ST_Distance's backend: Point/Point pairs
+        // skip the parse; the bound comparison is `distance <= bound`.
+        let executor = PointXYExecutor::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(|maybe_wkb0, maybe_wkb1| {
-            match (maybe_wkb0, maybe_wkb1, arg2_iter.next().unwrap()) {
-                (Some(wkb0), Some(wkb1), Some(distance)) => {
-                    builder.append_value(invoke_scalar(wkb0, wkb1, distance)?);
+        executor.execute_wkb_wkb_void(|maybe0, maybe1| {
+            match (maybe0, maybe1, arg2_iter.next().unwrap()) {
+                (Some(a), Some(b), Some(bound)) => {
+                    builder.append_value(point_or_wkb_distance(a, b) <= bound);
                 }
                 _ => builder.append_null(),
             }
@@ -77,11 +78,6 @@ impl SedonaScalarKernel for STDWithin {
     }
 }
 
-fn invoke_scalar(wkb_a: &Wkb, wkb_b: &Wkb, distance_bound: f64) -> Result<bool> {
-    let actual_distance = wkb_a.distance_ext(wkb_b);
-    Ok(actual_distance <= distance_bound)
-}
-
 #[cfg(test)]
 mod tests {
     use arrow_array::{create_array as arrow_array, ArrayRef};
diff --git a/rust/sedona-geometry/src/wkb_header.rs b/rust/sedona-geometry/src/wkb_header.rs
