perf: Optimize st_has(z/m) using WKBBytesExecutor + Implement new WKBHeader

benchmarks/test_functions.py

@@ -167,6 +167,42 @@ def queries():
 
         benchmark(queries)
 
+    @pytest.mark.parametrize(
+        "eng", [SedonaDBSingleThread, PostGISSingleThread, DuckDBSingleThread]
+    )
+    @pytest.mark.parametrize(
+        "table",
+        [
+            "collections_simple",
+            "collections_complex",
+        ],
+    )
+    def test_st_hasm(self, benchmark, eng, table):
+        eng = self._get_eng(eng)
+
+        def queries():
+            eng.execute_and_collect(f"SELECT ST_HasM(geom1) from {table}")
+
+        benchmark(queries)
+
+    @pytest.mark.parametrize(
+        "eng", [SedonaDBSingleThread, PostGISSingleThread, DuckDBSingleThread]
+    )
+    @pytest.mark.parametrize(
+        "table",
+        [
+            "collections_simple",
+            "collections_complex",
+        ],
+    )
+    def test_st_hasz(self, benchmark, eng, table):
+        eng = self._get_eng(eng)
+
+        def queries():
+            eng.execute_and_collect(f"SELECT ST_HasZ(geom1) from {table}")
+
+        benchmark(queries)
+
     @pytest.mark.parametrize(
         "eng", [SedonaDBSingleThread, PostGISSingleThread, DuckDBSingleThread]
     )

python/sedonadb/tests/functions/test_functions.py

@@ -561,8 +561,13 @@ def test_st_geomfromwkb(eng, geom):
         ("LINESTRING Z (0 0 0, 1 1 1)", True),
         ("POLYGON EMPTY", False),
         ("MULTIPOINT ((0 0), (1 1))", False),
+        ("MULTIPOINT Z ((0 0 0))", True),
+        ("MULTIPOINT ZM ((0 0 0 0))", True),
         ("GEOMETRYCOLLECTION EMPTY", False),
+        # Z-dim specified only in the nested geometry
         ("GEOMETRYCOLLECTION (POINT Z (0 0 0))", True),
+        # Z-dim specified on both levels
+        ("GEOMETRYCOLLECTION Z (POINT Z (0 0 0))", True),
         ("GEOMETRYCOLLECTION (GEOMETRYCOLLECTION (POINT Z (0 0 0)))", True),
     ],
 )

rust/sedona-functions/src/st_haszm.rs

@@ -16,19 +16,17 @@
 // under the License.
 use std::sync::Arc;
 
-use crate::executor::WkbExecutor;
+use crate::executor::WkbBytesExecutor;
 use arrow_array::builder::BooleanBuilder;
 use arrow_schema::DataType;
-use datafusion_common::error::Result;
+use datafusion_common::{error::Result, DataFusionError};
 use datafusion_expr::{
     scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility,
 };
-use geo_traits::GeometryCollectionTrait;
-use geo_traits::{Dimensions, GeometryTrait};
-use sedona_common::sedona_internal_err;
+use geo_traits::Dimensions;
 use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_geometry::wkb_header::WkbHeader;
 use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher};
-use wkb::reader::Wkb;
 
 pub fn st_hasz_udf() -> SedonaScalarUDF {
     SedonaScalarUDF::new(
@@ -91,13 +89,13 @@ impl SedonaScalarKernel for STHasZm {
             _ => unreachable!(),
         };
 
-        let executor = WkbExecutor::new(arg_types, args);
+        let executor = WkbBytesExecutor::new(arg_types, args);
         let mut builder = BooleanBuilder::with_capacity(executor.num_iterations());
 
         executor.execute_wkb_void(|maybe_item| {
             match maybe_item {
                 Some(item) => {
-                    builder.append_option(invoke_scalar(&item, dim_index)?);
+                    builder.append_option(invoke_scalar(item, dim_index)?);
                 }
                 None => builder.append_null(),
             }
@@ -108,27 +106,34 @@ impl SedonaScalarKernel for STHasZm {
     }
 }
 
-fn invoke_scalar(item: &Wkb, dim_index: usize) -> Result<Option<bool>> {
-    match item.as_type() {
-        geo_traits::GeometryType::GeometryCollection(collection) => {
-            if collection.num_geometries() == 0 {
-                Ok(Some(false))
-            } else {
-                // PostGIS doesn't allow creating a GeometryCollection with geometries of different dimensions
-                // so we can just check the dimension of the first one
-                let first_geom = unsafe { collection.geometry_unchecked(0) };
-                invoke_scalar(first_geom, dim_index)
-            }
-        }
-        _ => {
-            let geom_dim = item.dim();
-            match dim_index {
-                2 => Ok(Some(matches!(geom_dim, Dimensions::Xyz | Dimensions::Xyzm))),
-                3 => Ok(Some(matches!(geom_dim, Dimensions::Xym | Dimensions::Xyzm))),
-                _ => sedona_internal_err!("unexpected dim_index"),
-            }
-        }
+fn invoke_scalar(buf: &[u8], dim_index: usize) -> Result<Option<bool>> {
+    let header = WkbHeader::try_new(buf).map_err(|e| DataFusionError::External(Box::new(e)))?;
+    let top_level_dimensions = header
+        .dimensions()
+        .map_err(|e| DataFusionError::External(Box::new(e)))?;
+
+    // Infer dimension based on first coordinate dimension for cases where it differs from top-level
+    // e.g GEOMETRYCOLLECTION (POINT Z (1 2 3))
+    let dimensions;
+    if let Some(first_geom_dimensions) = header.first_geom_dimensions() {
+        dimensions = first_geom_dimensions;
+    } else {
+        dimensions = top_level_dimensions;
+    }
+
+    if dim_index == 2 {
+        return Ok(Some(matches!(
+            dimensions,
+            Dimensions::Xyz | Dimensions::Xyzm
+        )));
+    }
+    if dim_index == 3 {
+        return Ok(Some(matches!(
+            dimensions,
+            Dimensions::Xym | Dimensions::Xyzm
+        )));
     }
+    Ok(Some(false))
 }
 
 #[cfg(test)]
@@ -137,7 +142,9 @@ mod tests {
     use datafusion_expr::ScalarUDF;
     use rstest::rstest;
     use sedona_schema::datatypes::{WKB_GEOMETRY, WKB_VIEW_GEOMETRY};
-    use sedona_testing::testers::ScalarUdfTester;
+    use sedona_testing::{
+        fixtures::MULTIPOINT_WITH_INFERRED_Z_DIMENSION_WKB, testers::ScalarUdfTester,
+    };
 
     use super::*;
 
@@ -184,11 +191,19 @@ mod tests {
         let result = m_tester.invoke_wkb_scalar(None).unwrap();
         m_tester.assert_scalar_result_equals(result, ScalarValue::Null);
 
+        // Z-dimension specified only in the nested geometry, but not the geom collection level
         let result = z_tester
             .invoke_wkb_scalar(Some("GEOMETRYCOLLECTION (POINT Z (1 2 3))"))
             .unwrap();
         z_tester.assert_scalar_result_equals(result, ScalarValue::Boolean(Some(true)));
 
+        // Z-dimension specified on both the geom collection and nested geometry level
+        // Geometry collection with Z dimension both on the geom collection and nested geometry level
+        let result = z_tester
+            .invoke_wkb_scalar(Some("GEOMETRYCOLLECTION Z (POINT Z (1 2 3))"))
+            .unwrap();
+        z_tester.assert_scalar_result_equals(result, ScalarValue::Boolean(Some(true)));
+
         let result = m_tester
             .invoke_wkb_scalar(Some("GEOMETRYCOLLECTION (POINT M (1 2 3))"))
             .unwrap();
@@ -203,5 +218,32 @@ mod tests {
             .invoke_wkb_scalar(Some("GEOMETRYCOLLECTION EMPTY"))
             .unwrap();
         m_tester.assert_scalar_result_equals(result, ScalarValue::Boolean(Some(false)));
+
+        // Empty geometry collections with Z or M dimensions
+        let result = z_tester
+            .invoke_wkb_scalar(Some("GEOMETRYCOLLECTION Z EMPTY"))
+            .unwrap();
+        z_tester.assert_scalar_result_equals(result, ScalarValue::Boolean(Some(true)));
+
+        let result = m_tester
+            .invoke_wkb_scalar(Some("GEOMETRYCOLLECTION M EMPTY"))
+            .unwrap();
+        m_tester.assert_scalar_result_equals(result, ScalarValue::Boolean(Some(true)));
+    }
+
+    #[test]
+    fn multipoint_with_inferred_z_dimension() {
+        let z_tester = ScalarUdfTester::new(st_hasz_udf().into(), vec![WKB_GEOMETRY]);
+        let m_tester = ScalarUdfTester::new(st_hasm_udf().into(), vec![WKB_GEOMETRY]);
+
+        let scalar = ScalarValue::Binary(Some(MULTIPOINT_WITH_INFERRED_Z_DIMENSION_WKB.to_vec()));
+        assert_eq!(
+            z_tester.invoke_scalar(scalar.clone()).unwrap(),
+            ScalarValue::Boolean(Some(true))
+        );
+        assert_eq!(
+            m_tester.invoke_scalar(scalar.clone()).unwrap(),
+            ScalarValue::Boolean(Some(false))
+        );
     }
 }

rust/sedona-geometry/Cargo.toml

@@ -30,6 +30,7 @@ result_large_err = "allow"
 [dev-dependencies]
 geo-types = { workspace = true }
 rstest = { workspace = true }
+sedona-testing = { path = "../sedona-testing" }
 serde_json = { workspace = true }
 wkt = { workspace = true }
 

rust/sedona-geometry/src/lib.rs

@@ -24,3 +24,4 @@ pub mod point_count;
 pub mod transform;
 pub mod types;
 pub mod wkb_factory;
+pub mod wkb_header;