perf: Faster st_geometrytype() function

rust/sedona-functions/src/executor.rs

@@ -119,6 +119,43 @@ impl<'a, 'b, Factory0: GeometryFactory, Factory1: GeometryFactory>
         }
     }
 
+    /// Execute a function by iterating over [Wkb]'s raw binary representation.
+    /// The provided `func` can assume its input bytes is a valid [Wkb] binary.
+    ///
+    /// This function assumes the first argument from the executor is either a [SedonaType::Wkb]
+    /// or [SedonaType::WkbView] type. If other type is provided, an error will be
+    /// thrown.
+    pub fn execute_wkb_bytes_void<F: FnMut(Option<&'b [u8]>) -> Result<()>>(
+        &self,
+        mut func: F,
+    ) -> Result<()> {
+        // Ensure the first argument of the executor is either Wkb, WkbView, or
+        // a Null type (to support columns of all-null values)
+        match &self.arg_types[0] {
+            SedonaType::Wkb(_, _)
+            | SedonaType::WkbView(_, _)
+            | SedonaType::Arrow(DataType::Null) => {}
+            other => {
+                return sedona_internal_err!(
+                    "Expected SedonaType::Wkb or SedonaType::WkbView or SedonaType::Arrow(DataType::Null) for the first arg, got {}",
+                    other
+                )
+            }
+        }
+
+        // We can assume the storage type for Wkb is Binary, and also BinaryView for
+        // WkbView
+        match &self.args[0] {
+            ColumnarValue::Array(array) => {
+                array.iter_as_wkb_bytes(&self.arg_types[0], self.num_iterations, func)
+            }
+            ColumnarValue::Scalar(scalar_value) => {
+                let maybe_bytes = scalar_value.scalar_as_wkb_bytes()?;
+                func(maybe_bytes)
+            }
+        }
+    }
+
     /// Execute a binary geometry function by iterating over [Wkb] scalars in the
     /// first two arguments
     ///

rust/sedona-functions/src/st_geometrytype.rs

@@ -23,11 +23,9 @@ use datafusion_common::error::Result;
 use datafusion_expr::{
     scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility,
 };
-use geo_traits::GeometryTrait;
 use sedona_common::sedona_internal_err;
 use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
 use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher};
-use wkb::reader::Wkb;
 
 pub fn st_geometry_type_udf() -> SedonaScalarUDF {
     SedonaScalarUDF::new(
@@ -71,12 +69,12 @@ impl SedonaScalarKernel for STGeometryType {
         let min_output_size = "POINT".len() * executor.num_iterations();
         let mut builder = StringBuilder::with_capacity(executor.num_iterations(), min_output_size);
 
-        // We can do quite a lot better than this with some vectorized WKB processing,
-        // but for now we just do a slow iteration
-        executor.execute_wkb_void(|maybe_item| {
-            match maybe_item {
-                Some(item) => {
-                    builder.append_option(invoke_scalar(&item)?);
+        // Iterate over raw WKB bytes for faster type inference
+        executor.execute_wkb_bytes_void(|maybe_bytes| {
+            match maybe_bytes {
+                Some(bytes) => {
+                    let name = infer_geometry_type_name(bytes)?;
+                    builder.append_value(name);
                 }
                 None => builder.append_null(),
             }
@@ -87,20 +85,33 @@ impl SedonaScalarKernel for STGeometryType {
     }
 }
 
-fn invoke_scalar(item: &Wkb) -> Result<Option<String>> {
-    match item.as_type() {
-        geo_traits::GeometryType::Point(_) => Ok(Some("ST_Point".to_string())),
-        geo_traits::GeometryType::LineString(_) => Ok(Some("ST_LineString".to_string())),
-        geo_traits::GeometryType::Polygon(_) => Ok(Some("ST_Polygon".to_string())),
-        geo_traits::GeometryType::MultiPoint(_) => Ok(Some("ST_MultiPoint".to_string())),
-        geo_traits::GeometryType::MultiLineString(_) => Ok(Some("ST_MultiLineString".to_string())),
-        geo_traits::GeometryType::MultiPolygon(_) => Ok(Some("ST_MultiPolygon".to_string())),
-        geo_traits::GeometryType::GeometryCollection(_) => {
-            Ok(Some("ST_GeometryCollection".to_string()))
-        }
-
-        // Other geometry types in geo that we should not get here: Rect, Triangle, Line
-        _ => sedona_internal_err!("unexpected geometry type"),
+/// Fast-path inference of geometry type name from raw WKB bytes
+/// An error will be thrown for invalid WKB bytes input
+///
+/// Spec: https://libgeos.org/specifications/wkb/
+#[inline]
+fn infer_geometry_type_name(buf: &[u8]) -> Result<&'static str> {
+    if buf.len() < 5 {
+        return sedona_internal_err!("Invalid WKB: buffer too small ({} bytes)", buf.len());
+    }
+
+    let byte_order = buf[0];
+    let code = match byte_order {
+        0 => u32::from_be_bytes([buf[1], buf[2], buf[3], buf[4]]),
+        1 => u32::from_le_bytes([buf[1], buf[2], buf[3], buf[4]]),
+        other => return sedona_internal_err!("Unexpected byte order: {other}"),
+    };
+
+    // Only low 3 bits is for the base type, high bits include additional info
+    match code & 0x7 {
+        1 => Ok("ST_Point"),
+        2 => Ok("ST_LineString"),
+        3 => Ok("ST_Polygon"),
+        4 => Ok("ST_MultiPoint"),
+        5 => Ok("ST_MultiLineString"),
+        6 => Ok("ST_MultiPolygon"),
+        7 => Ok("ST_GeometryCollection"),
+        _ => sedona_internal_err!("WKB type code out of range. Got: {}", code),
     }
 }