diff --git a/rust/sedona-functions/benches/native-functions.rs b/rust/sedona-functions/benches/native-functions.rs index 6607e7ba2f..ff90660ace 100644 --- a/rust/sedona-functions/benches/native-functions.rs +++ b/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)); diff --git a/rust/sedona-functions/src/executor.rs b/rust/sedona-functions/src/executor.rs index 4782048609..f3b6f75f44 100644 --- 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> { + // 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-functions/src/st_azimuth.rs b/rust/sedona-functions/src/st_azimuth.rs index 5b661adfc2..3ab56035b7 100644 --- a/rust/sedona-functions/src/st_azimuth.rs +++ b/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 { - 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> { - 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> { + 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"), } diff --git a/rust/sedona-geo/benches/geo-functions.rs b/rust/sedona-geo/benches/geo-functions.rs index 513cc141d4..2ba42ceddf 100644 --- 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 index 51c7c23a2e..f128d67361 100644 --- 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 { @@ -52,16 +52,17 @@ impl SedonaScalarKernel for STDistance { arg_types: &[SedonaType], args: &[ColumnarValue], ) -> Result { - 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(()) })?; @@ -69,25 +70,42 @@ impl SedonaScalarKernel for STDistance { } } -fn invoke_scalar(wkb_a: &Wkb, wkb_b: &Wkb) -> Result { - 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 index 63bb22f45c..0b3b283c8e 100644 --- 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 { @@ -57,15 +56,17 @@ impl SedonaScalarKernel for STDWithin { args: &[ColumnarValue], ) -> Result { 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 { - 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 index 2e1b1c8e07..44b4a30bd8 100644 --- a/rust/sedona-geometry/src/wkb_header.rs +++ b/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, 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, +} + +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, 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 read_u32 = |offset: usize| -> Result { - let bytes: [u8; 4] = buf - .get(offset..offset + 4) + 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) + }; - let read_f64 = |offset: usize| -> Result { + // 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 }; + + 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, 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 { 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, SedonaGeometryError> { + match WkbPointLayout::try_from_wkb(buf)? { + Some(layout) => layout.read_xy(buf), + None => Err(SedonaGeometryError::Invalid( + "WKB: expected a Point".to_string(), + )), } } @@ -1121,6 +1161,133 @@ mod tests { assert!(header.is_empty().unwrap()); } + #[test] + fn wkb_point_layout_accepts_point_encodings() { + // ISO points of every dimensionality: the layout is found and the + // coordinate read matches the leading (x, y). + for wkt_value in [ + "POINT (1 2)", + "POINT Z (1 2 3)", + "POINT M (1 2 3)", + "POINT ZM (1 2 3 4)", + ] { + let wkb = make_wkb(wkt_value); + let layout = WkbPointLayout::try_from_wkb(&wkb) + .unwrap() + .unwrap_or_else(|| panic!("{wkt_value} not recognized as a Point")); + assert_eq!( + layout.read_xy(&wkb).unwrap(), + Some((1.0, 2.0)), + "{wkt_value}" + ); + } + + // EWKB points: the SRID between the type code and the coordinate must be + // skipped for every dimensionality. + for wkb in [ + &POINT_WITH_SRID_4326_EWKB[..], + &POINT_Z_WITH_SRID_3857_EWKB[..], + &POINT_M_WITH_SRID_4326_EWKB[..], + &POINT_ZM_WITH_SRID_4326_EWKB[..], + ] { + let layout = WkbPointLayout::try_from_wkb(wkb).unwrap().unwrap(); + assert_eq!(layout.read_xy(wkb).unwrap(), Some((1.0, 2.0))); + } + + // Big-endian EWKB: SRID flag and coordinate both decoded big-endian. + let be_ewkb: [u8; 25] = [ + 0x00, // big-endian + 0x20, 0x00, 0x00, 0x01, // type code 1 (Point) with SRID flag + 0x00, 0x00, 0x10, 0xE6, // SRID 4326 + 0x3F, 0xF0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // x = 1.0 + 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // y = 2.0 + ]; + let layout = WkbPointLayout::try_from_wkb(&be_ewkb).unwrap().unwrap(); + assert_eq!(layout.read_xy(&be_ewkb).unwrap(), Some((1.0, 2.0))); + } + + #[test] + fn wkb_point_layout_non_point_is_none() { + for wkt_value in [ + "LINESTRING (1 2, 3 4)", + "POLYGON ((0 0, 0 1, 1 0, 0 0))", + "MULTIPOINT ((1 2))", + "GEOMETRYCOLLECTION (POINT (1 2))", + ] { + let wkb = make_wkb(wkt_value); + assert!( + WkbPointLayout::try_from_wkb(&wkb).unwrap().is_none(), + "{wkt_value} unexpectedly recognized as a Point" + ); + } + + // Non-point EWKB: the SRID flag must not confuse type detection. + assert!( + WkbPointLayout::try_from_wkb(&LINESTRING_WITH_SRID_4326_EWKB) + .unwrap() + .is_none() + ); + assert!( + WkbPointLayout::try_from_wkb(&MULTIPOINT_WITH_SRID_4326_EWKB) + .unwrap() + .is_none() + ); + } + + #[test] + fn wkb_point_layout_incomplete_iso_buffer() { + let wkb = make_wkb("POINT (1 2)"); + + // No byte order at all, and an invalid byte-order flag. + assert!(WkbPointLayout::try_from_wkb(&[]).is_err()); + assert!(WkbPointLayout::try_from_wkb(&[0x02, 0, 0, 0, 1]).is_err()); + + // Shorter than byte order + type code: the header itself is truncated. + for i in 1..5 { + assert!( + WkbPointLayout::try_from_wkb(&wkb[0..i]).is_err(), + "header 0..{i} unexpectedly succeeded" + ); + } + + // Header is readable but the coordinate is truncated: read_xy must + // error rather than read out of bounds. + for i in 5..wkb.len() { + let layout = WkbPointLayout::try_from_wkb(&wkb[0..i]).unwrap().unwrap(); + assert!( + layout.read_xy(&wkb[0..i]).is_err(), + "coordinate 0..{i} unexpectedly succeeded" + ); + } + } + + #[test] + fn wkb_point_layout_incomplete_ewkb_buffer() { + // 1 byte order + 4 type code + 4 SRID + 8 x + 8 y + let wkb = POINT_WITH_SRID_4326_EWKB; + assert_eq!(wkb.len(), 25); + + for i in 1..5 { + assert!( + WkbPointLayout::try_from_wkb(&wkb[0..i]).is_err(), + "header 0..{i} unexpectedly succeeded" + ); + } + + // Everything shorter than the full SRID + coordinate must error: the + // coordinate offset accounts for the SRID, so a buffer that would fit an + // ISO point is still too short here. + for i in 5..wkb.len() { + let layout = WkbPointLayout::try_from_wkb(&wkb[0..i]).unwrap().unwrap(); + assert!( + layout.read_xy(&wkb[0..i]).is_err(), + "coordinate 0..{i} unexpectedly succeeded" + ); + } + let layout = WkbPointLayout::try_from_wkb(&wkb).unwrap().unwrap(); + assert_eq!(layout.read_xy(&wkb).unwrap(), Some((1.0, 2.0))); + } + #[test] fn read_point_xy_iso_little_endian() { // 2-D / 3-D / 4-D ISO points all expose the same leading (x, y), and the