feat(rust/sedona-functions): Implement ST_Force2D() and ST_Force3D()

c/sedona-proj/src/transform.rs

@@ -16,6 +16,7 @@
 // under the License.
 use crate::error::SedonaProjError;
 use crate::proj::{Proj, ProjContext};
+use geo_traits::Dimensions;
 use sedona_geometry::bounding_box::BoundingBox;
 use sedona_geometry::error::SedonaGeometryError;
 use sedona_geometry::interval::IntervalTrait;
@@ -229,7 +230,11 @@ impl CrsTransform for ProjTransform {
         Ok(())
     }
 
-    fn transform_coord_3d(&self, coord: &mut (f64, f64, f64)) -> Result<(), SedonaGeometryError> {
+    fn transform_coord_3d(
+        &self,
+        coord: &mut (f64, f64, f64),
+        _input_dims: Dimensions,
+    ) -> Result<(), SedonaGeometryError> {
         let res = self.proj.borrow_mut().transform_xyz(*coord).map_err(|e| {
             SedonaGeometryError::Invalid(format!(
                 "PROJ coordinate transformation failed with error: {e}"

python/sedonadb/tests/functions/test_functions.py

@@ -17,8 +17,8 @@
 import math
 
 import pytest
-import shapely
 import sedonadb
+import shapely
 from sedonadb.testing import PostGIS, SedonaDB, geom_or_null, val_or_null
 
 
@@ -1479,6 +1479,46 @@ def test_st_flipcoordinates(eng, geom, expected):
     )
 
 
+@pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
+@pytest.mark.parametrize(
+    ("geom", "expected_2d", "expected_3d"),
+    [
+        (None, None, None),
+        ("POINT EMPTY", "POINT (nan nan)", "POINT Z (nan nan nan)"),
+        ("POLYGON EMPTY", "POLYGON EMPTY", "POLYGON Z EMPTY"),
+        ("LINESTRING EMPTY", "LINESTRING EMPTY", "LINESTRING Z EMPTY"),
+        ("MULTIPOINT EMPTY", "MULTIPOINT EMPTY", "MULTIPOINT Z EMPTY"),
+        ("MULTILINESTRING EMPTY", "MULTILINESTRING EMPTY", "MULTILINESTRING Z EMPTY"),
+        ("MULTIPOLYGON EMPTY", "MULTIPOLYGON EMPTY", "MULTIPOLYGON Z EMPTY"),
+        (
+            "GEOMETRYCOLLECTION EMPTY",
+            "GEOMETRYCOLLECTION EMPTY",
+            "GEOMETRYCOLLECTION Z EMPTY",
+        ),
+        ("POINT (0 1)", "POINT (0 1)", "POINT Z (0 1 5)"),
+        (
+            "LINESTRING (0 1, 2 3)",
+            "LINESTRING (0 1, 2 3)",
+            "LINESTRING Z (0 1 5, 2 3 5)",
+        ),
+        (
+            "MULTIPOINT (0 1, 2 3)",
+            "MULTIPOINT (0 1, 2 3)",
+            "MULTIPOINT Z (0 1 5, 2 3 5)",
+        ),
+        (
+            "GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (3 4, 5 6), POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0)))",
+            "GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (3 4, 5 6), POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0)))",
+            "GEOMETRYCOLLECTION Z (POINT Z (1 2 5), LINESTRING Z (3 4 5, 5 6 5), POLYGON Z ((0 0 5, 0 1 5, 1 1 5, 1 0 5, 0 0 5)))",
+        ),
+    ],
+)
+def test_st_force_dim(eng, geom, expected_2d, expected_3d):
+    eng = eng.create_or_skip()
+    eng.assert_query_result(f"SELECT ST_Force2D({geom_or_null(geom)})", expected_2d)
+    eng.assert_query_result(f"SELECT ST_Force3D({geom_or_null(geom)}, 5)", expected_3d)
+
+
 @pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
 @pytest.mark.parametrize(
     ("geom", "expected"),

rust/sedona-functions/src/lib.rs

@@ -42,6 +42,7 @@ mod st_dwithin;
 pub mod st_envelope;
 pub mod st_envelope_agg;
 pub mod st_flipcoordinates;
+mod st_force_dim;
 mod st_geometryn;
 mod st_geometrytype;
 mod st_geomfromewkb;

rust/sedona-functions/src/register.rs

@@ -118,6 +118,8 @@ pub fn default_function_set() -> FunctionSet {
         crate::st_start_point::st_start_point_udf,
         crate::st_transform::st_transform_udf,
         crate::st_translate::st_translate_udf,
+        crate::st_force_dim::st_force2d_udf,
+        crate::st_force_dim::st_force3d_udf,
         crate::st_xyzm_minmax::st_mmax_udf,
         crate::st_xyzm_minmax::st_mmin_udf,
         crate::st_xyzm_minmax::st_xmax_udf,

rust/sedona-functions/src/st_affine_helpers.rs

@@ -19,6 +19,7 @@ use arrow_array::Array;
 use arrow_array::PrimitiveArray;
 use datafusion_common::cast::as_float64_array;
 use datafusion_common::error::Result;
+use geo_traits::Dimensions;
 use sedona_common::sedona_internal_err;
 use sedona_geometry::transform::CrsTransform;
 use std::sync::Arc;
@@ -438,6 +439,7 @@ impl CrsTransform for DAffine {
     fn transform_coord_3d(
         &self,
         coord: &mut (f64, f64, f64),
+        _input_dims: Dimensions,
     ) -> std::result::Result<(), sedona_geometry::error::SedonaGeometryError> {
         match self {
             DAffine::DAffine2(daffine2) => {
@@ -536,6 +538,7 @@ mod tests {
     use super::*;
     use arrow_array::Array;
     use arrow_array::Float64Array;
+    use geo_traits::Dimensions;
     use std::sync::Arc;
 
     fn float_array(values: Vec<Option<f64>>) -> Arc<dyn Array> {
@@ -628,7 +631,9 @@ mod tests {
         let mut coord_3d = (1.0, 2.0, 3.0);
         let affine_3d =
             DAffine::DAffine3(glam::DAffine3::from_scale(glam::DVec3::new(2.0, 3.0, 4.0)));
-        affine_3d.transform_coord_3d(&mut coord_3d).unwrap();
+        affine_3d
+            .transform_coord_3d(&mut coord_3d, Dimensions::Xyz)
+            .unwrap();
         assert_eq!(coord_3d, (2.0, 6.0, 12.0));
     }
 }