perf(rust/sedona-geo): fast-path point-to-point ST_Distance

Inspect each operand's WKB header once via the new WkbPointLayout primitive
and, when both are Points, compute the Euclidean distance directly from the
coordinate offsets, skipping a full geometry parse and the generic distance
dispatch. A per-batch check routes a non-Point scalar operand to the
parse-once general path so it is never re-parsed per row.

Author: james-willis

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

@@ -60,7 +60,11 @@ 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", ArrayArray(Point, Polygon(10)));
     benchmark::scalar(
         c,
         &f,

rust/sedona-geo/src/st_distance.rs

@@ -18,16 +18,18 @@ use std::sync::Arc;
 
 use arrow_array::builder::Float64Builder;
 use arrow_schema::DataType;
-use datafusion_common::error::Result;
+use datafusion_common::error::{DataFusionError, Result};
+use datafusion_common::scalar::ScalarValue;
 use datafusion_expr::ColumnarValue;
 use sedona_expr::{
     item_crs::ItemCrsKernel,
     scalar_udf::{ScalarKernelRef, SedonaScalarKernel},
 };
-use sedona_functions::executor::WkbExecutor;
+use sedona_functions::executor::{WkbBytesExecutor, WkbExecutor};
 use sedona_geo_generic_alg::line_measures::DistanceExt;
+use sedona_geometry::wkb_header::WkbPointLayout;
 use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher};
-use wkb::reader::Wkb;
+use wkb::reader::read_wkb;
 
 /// ST_Distance() implementation using [DistanceExt]
 pub fn st_distance_impl() -> Vec<ScalarKernelRef> {
@@ -52,12 +54,21 @@ impl SedonaScalarKernel for STDistance {
         arg_types: &[SedonaType],
         args: &[ColumnarValue],
     ) -> Result<ColumnarValue> {
-        let executor = WkbExecutor::new(arg_types, args);
+        // The point-to-point fast path can only fire when neither operand is a
+        // non-Point scalar: every row pairs against that scalar, so reading raw
+        // bytes would re-parse it on every row for no benefit. Detect that case
+        // up front (one cheap header read per scalar) and take the general path,
+        // which parses each scalar exactly once.
+        if scalar_is_non_point(&args[0])? || scalar_is_non_point(&args[1])? {
+            return invoke_general(arg_types, args);
+        }
+
+        let executor = WkbBytesExecutor::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(|maybe_bytes0, maybe_bytes1| {
+            match (maybe_bytes0, maybe_bytes1) {
+                (Some(bytes0), Some(bytes1)) => {
+                    builder.append_value(invoke_scalar_bytes(bytes0, bytes1)?);
                 }
                 _ => builder.append_null(),
             }
@@ -69,8 +80,75 @@ impl SedonaScalarKernel for STDistance {
     }
 }
 
-fn invoke_scalar(wkb_a: &Wkb, wkb_b: &Wkb) -> Result<f64> {
-    Ok(wkb_a.distance_ext(wkb_b))
+/// Distance between two operands given their raw WKB bytes.
+fn invoke_scalar_bytes(bytes_a: &[u8], bytes_b: &[u8]) -> Result<f64> {
+    // Point-to-point fast path: inspect each operand's header once. When both
+    // are Points, read their (x, y) from the offsets just discovered and compute
+    // the Euclidean distance directly, skipping a full geometry parse and the
+    // generic distance dispatch. Inspecting both headers before reading any
+    // coordinate means a non-Point operand never triggers a wasted coordinate
+    // read, and the matched Point coordinates are read without re-parsing the
+    // header.
+    if let (Ok(Some(layout_a)), Ok(Some(layout_b))) = (
+        WkbPointLayout::try_from_wkb(bytes_a),
+        WkbPointLayout::try_from_wkb(bytes_b),
+    ) {
+        if let (Ok(Some((ax, ay))), Ok(Some((bx, by)))) =
+            (layout_a.read_xy(bytes_a), layout_b.read_xy(bytes_b))
+        {
+            let dx = ax - bx;
+            let dy = ay - by;
+            return Ok((dx * dx + dy * dy).sqrt());
+        }
+    }
+
+    // Any non-Point, POINT EMPTY, or malformed input falls back to the general
+    // parser, which preserves the exact semantics (and surfaces the same error
+    // on truly malformed bytes).
+    let wkb_a = read_wkb(bytes_a).map_err(|e| DataFusionError::External(Box::new(e)))?;
+    let wkb_b = read_wkb(bytes_b).map_err(|e| DataFusionError::External(Box::new(e)))?;
+    Ok(wkb_a.distance_ext(&wkb_b))
+}
+
+/// General path: parse both operands (parsing any scalar operand exactly once)
+/// and use the generic distance implementation.
+fn invoke_general(arg_types: &[SedonaType], args: &[ColumnarValue]) -> Result<ColumnarValue> {
+    let executor = WkbExecutor::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(wkb0.distance_ext(wkb1));
+            }
+            _ => builder.append_null(),
+        }
+
+        Ok(())
+    })?;
+
+    executor.finish(Arc::new(builder.finish()))
+}
+
+/// Returns `true` only when `arg` is a scalar WKB value that is definitively not
+/// a Point. Array arguments, nulls, and Point scalars return `false` (the fast
+/// path may still apply per row). Reads only the WKB header.
+fn scalar_is_non_point(arg: &ColumnarValue) -> Result<bool> {
+    if let ColumnarValue::Scalar(scalar) = arg {
+        if let Some(bytes) = scalar_wkb_bytes(scalar) {
+            return Ok(matches!(WkbPointLayout::try_from_wkb(bytes), Ok(None)));
+        }
+    }
+    Ok(false)
+}
+
+/// Extracts the raw WKB bytes from a scalar geometry value, if present.
+fn scalar_wkb_bytes(scalar: &ScalarValue) -> Option<&[u8]> {
+    match scalar {
+        ScalarValue::Binary(maybe)
+        | ScalarValue::BinaryView(maybe)
+        | ScalarValue::LargeBinary(maybe) => maybe.as_deref(),
+        _ => None,
+    }
 }
 
 #[cfg(test)]

rust/sedona-geometry/src/wkb_header.rs

@@ -152,74 +152,114 @@ impl WkbHeader {
     }
 }
 
-/// Read the `(x, y)` of a WKB Point without parsing the geometry generally.
+/// The coordinate layout of a WKB Point, discovered by reading only its header.
 ///
-/// This is a fast path for the common case of sampling/indexing by a single
-/// point: a Point WKB has a fixed layout (`byte order`, `type code`, optional
-/// SRID, then the coordinate), so the coordinate is at a known offset and can
-/// be read with a couple of fixed-offset loads. It deliberately handles only
-/// Points — any other geometry is an error.
+/// This is the fast-path primitive for sampling/indexing by Point. A Point WKB
+/// has a fixed layout (`byte order`, `type code`, optional SRID, then the
+/// coordinate), so [`try_from_wkb`](Self::try_from_wkb) can cheaply decide
+/// whether a buffer is a Point — reading only the first few header bytes, never
+/// a coordinate — and a later [`read_xy`](Self::read_xy) reads the coordinate
+/// from the offset it already found.
 ///
-/// Returns `Ok(None)` for `POINT EMPTY` (encoded as `NaN`/`NaN`), `Ok(Some((x, y)))`
-/// for a finite point, and an error if the bytes are not a Point or are truncated.
-/// Accepts both ISO (with Z/M dimension flags) and EWKB (with Z/M/SRID flags)
-/// encodings; only the first two ordinates are read.
-pub fn read_point_xy(buf: &[u8]) -> Result<Option<(f64, f64)>, SedonaGeometryError> {
-    let little_endian = match buf.first() {
-        Some(1) => true,
-        Some(0) => false,
-        _ => {
-            return Err(SedonaGeometryError::Invalid(
-                "WKB: missing or invalid byte order".to_string(),
-            ))
-        }
-    };
+/// Splitting the header inspection from the coordinate read lets a binary
+/// function (e.g. distance) test *both* operands first and read coordinates only
+/// when both are Points, without re-parsing the header or wasting a coordinate
+/// read on rows that fall back to the general parser.
+#[derive(Debug, Clone, Copy)]
+pub struct WkbPointLayout {
+    little_endian: bool,
+    coord_offset: usize,
+}
 
-    let read_u32 = |offset: usize| -> Result<u32, SedonaGeometryError> {
-        let bytes: [u8; 4] = buf
-            .get(offset..offset + 4)
+impl WkbPointLayout {
+    /// Inspects a WKB header, returning the Point layout or `Ok(None)` if `buf`
+    /// is not a Point. Reads only the byte order and type code (recognizing the
+    /// EWKB SRID prefix) — never a coordinate. Errors only on a missing/invalid
+    /// byte order or a truncated type code. Accepts both ISO (with Z/M dimension
+    /// flags) and EWKB (with Z/M/SRID flags) encodings.
+    #[inline]
+    pub fn try_from_wkb(buf: &[u8]) -> Result<Option<Self>, SedonaGeometryError> {
+        let little_endian = match buf.first() {
+            Some(1) => true,
+            Some(0) => false,
+            _ => {
+                return Err(SedonaGeometryError::Invalid(
+                    "WKB: missing or invalid byte order".to_string(),
+                ))
+            }
+        };
+
+        let type_bytes: [u8; 4] = buf
+            .get(1..5)
             .ok_or_else(|| SedonaGeometryError::Invalid("WKB: truncated header".to_string()))?
             .try_into()
             .unwrap();
-        Ok(if little_endian {
-            u32::from_le_bytes(bytes)
+        let type_code = if little_endian {
+            u32::from_le_bytes(type_bytes)
         } else {
-            u32::from_be_bytes(bytes)
-        })
-    };
+            u32::from_be_bytes(type_bytes)
+        };
+
+        // The low three bits of the base type identify a Point (1) regardless of
+        // the ISO dimension digits or EWKB high-bit flags.
+        if type_code & 0x7 != 1 {
+            return Ok(None);
+        }
+
+        // EWKB carries a 4-byte SRID between the type code and the coordinate.
+        let has_srid = type_code & SRID_FLAG_BIT != 0;
+        let coord_offset = if has_srid { 9 } else { 5 };
 
-    let read_f64 = |offset: usize| -> Result<f64, SedonaGeometryError> {
+        Ok(Some(Self {
+            little_endian,
+            coord_offset,
+        }))
+    }
+
+    /// Reads the `(x, y)` from `buf`, which must be the same buffer passed to
+    /// [`try_from_wkb`](Self::try_from_wkb). Returns `Ok(None)` for `POINT EMPTY`
+    /// (encoded as `NaN`/`NaN`); errors if the coordinate is truncated.
+    #[inline]
+    pub fn read_xy(&self, buf: &[u8]) -> Result<Option<(f64, f64)>, SedonaGeometryError> {
+        let x = self.read_f64(buf, self.coord_offset)?;
+        let y = self.read_f64(buf, self.coord_offset + 8)?;
+        if x.is_nan() && y.is_nan() {
+            Ok(None)
+        } else {
+            Ok(Some((x, y)))
+        }
+    }
+
+    #[inline]
+    fn read_f64(&self, buf: &[u8], offset: usize) -> Result<f64, SedonaGeometryError> {
         let bytes: [u8; 8] = buf
             .get(offset..offset + 8)
             .ok_or_else(|| SedonaGeometryError::Invalid("WKB: truncated coordinate".to_string()))?
             .try_into()
             .unwrap();
-        Ok(if little_endian {
+        Ok(if self.little_endian {
             f64::from_le_bytes(bytes)
         } else {
             f64::from_be_bytes(bytes)
         })
-    };
-
-    let type_code = read_u32(1)?;
-    // The low three bits of the base type identify a Point (1) regardless of the
-    // ISO dimension digits or EWKB high-bit flags.
-    if type_code & 0x7 != 1 {
-        return Err(SedonaGeometryError::Invalid(format!(
-            "WKB: expected a Point (type code {type_code})"
-        )));
     }
+}
 
-    // EWKB carries a 4-byte SRID between the type code and the coordinate.
-    let has_srid = type_code & SRID_FLAG_BIT != 0;
-    let coord_offset = if has_srid { 9 } else { 5 };
-
-    let x = read_f64(coord_offset)?;
-    let y = read_f64(coord_offset + 8)?;
-    if x.is_nan() && y.is_nan() {
-        Ok(None)
-    } else {
-        Ok(Some((x, y)))
+/// Read the `(x, y)` of a WKB Point without parsing the geometry generally.
+///
+/// Convenience wrapper over [`WkbPointLayout`] for the single-Point case: it
+/// deliberately handles only Points — any other geometry is an error.
+///
+/// Returns `Ok(None)` for `POINT EMPTY` (encoded as `NaN`/`NaN`), `Ok(Some((x, y)))`
+/// for a finite point, and an error if the bytes are not a Point or are truncated.
+/// Accepts both ISO (with Z/M dimension flags) and EWKB (with Z/M/SRID flags)
+/// encodings; only the first two ordinates are read.
+pub fn read_point_xy(buf: &[u8]) -> Result<Option<(f64, f64)>, SedonaGeometryError> {
+    match WkbPointLayout::try_from_wkb(buf)? {
+        Some(layout) => layout.read_xy(buf),
+        None => Err(SedonaGeometryError::Invalid(
+            "WKB: expected a Point".to_string(),
+        )),
     }
 }