feat(rust/sedona-functions): Implement ST_Affine(), ST_Rotate(), and ST_Scale()

Cargo.toml

@@ -92,6 +92,7 @@ dirs = "6.0.0"
 env_logger = "0.11"
 fastrand = "2.0"
 futures = "0.3"
+glam = "0.30.10"
 object_store = { version = "0.12.4", default-features = false }
 float_next_after = "1"
 num-traits = { version = "0.2", default-features = false, features = ["libm"] }

python/sedonadb/tests/functions/test_functions.py

@@ -189,6 +189,177 @@ def test_st_azimuth(eng, geom1, geom2, expected):
     )
 
 
+# fmt: off
+@pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
+@pytest.mark.parametrize(
+    ("geom", "a", "b", "d", "e", "xoff", "yoff", "expected"),
+    [
+        (
+            "POINT (1 2)",
+            1.0, 0.0,
+            0.0, 1.0,
+            0.0, 0.0,
+            "POINT (1 2)"),
+        (
+            "POINT (1 2)",
+            2.0, 0.0,
+            0.0, 2.0,
+            1.0, 3.0,
+            "POINT (3 7)"),
+        (
+            "LINESTRING (0 0, 1 1)",
+            1.0, 0.0,
+            0.0, 1.0,
+            1.0, 2.0,
+            "LINESTRING (1 2, 2 3)"
+        ),
+    ],
+)
+def test_st_affine_2d(eng, geom, a, b, d, e, xoff, yoff, expected):
+    eng = eng.create_or_skip()
+    eng.assert_query_result(
+        "SELECT ST_Affine("
+        f"{geom_or_null(geom)}, "
+        f"{val_or_null(a)}, {val_or_null(b)}, {val_or_null(d)}, {val_or_null(e)}, "
+        f"{val_or_null(xoff)}, {val_or_null(yoff)})",
+        expected,
+    )
+# fmt: on
+
+
+# fmt: off
+@pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
+@pytest.mark.parametrize(
+    ("geom", "a", "b", "c", "d", "e", "f", "g", "h", "i", "xoff", "yoff", "zoff", "expected"),
+    [
+        (
+            "POINT Z (1 2 3)",
+            1.0, 0.0, 0.0,
+            0.0, 1.0, 0.0,
+            0.0, 0.0, 1.0,
+            0.0, 0.0, 0.0,
+            "POINT Z (1 2 3)",
+        ),
+        (
+            "POINT Z (1 2 3)",
+            2.0, 0.0, 0.0,
+            0.0, 2.0, 0.0,
+            0.0, 0.0, 2.0,
+            1.0, 3.0, 5.0,
+            "POINT Z (3 7 11)",
+        ),
+    ],
+)
+def test_st_affine_3d(
+    eng, geom, a, b, c, d, e, f, g, h, i, xoff, yoff, zoff, expected
+):
+    eng = eng.create_or_skip()
+    query = (
+        "SELECT ST_Affine("
+        f"{geom_or_null(geom)}, "
+        f"{val_or_null(a)}, {val_or_null(b)}, {val_or_null(c)}, "
+        f"{val_or_null(d)}, {val_or_null(e)}, {val_or_null(f)}, "
+        f"{val_or_null(g)}, {val_or_null(h)}, {val_or_null(i)}, "
+        f"{val_or_null(xoff)}, {val_or_null(yoff)}, {val_or_null(zoff)})"
+    )
+    try:
+        eng.assert_query_result(query, expected)
+    except Exception as exc:
+        if isinstance(eng, PostGIS):
+            pytest.skip(f"PostGIS may not support 3D ST_Affine: {exc}")
+        raise
+# fmt: on
+
+
+@pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
+@pytest.mark.parametrize(
+    ("geom", "sx", "sy", "expected"),
+    [
+        (None, 2.0, 3.0, None),
+        ("POINT (1 2)", 1.0, 1.0, "POINT (1 2)"),
+        ("POINT (1 2)", 2.0, 3.0, "POINT (2 6)"),
+        ("LINESTRING (0 0, 1 1)", 2.0, 3.0, "LINESTRING (0 0, 2 3)"),
+        (
+            "POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0))",
+            2.0,
+            3.0,
+            "POLYGON ((0 0, 2 0, 2 3, 0 3, 0 0))",
+        ),
+        (
+            "MULTIPOINT (1 2, 3 4)",
+            2.0,
+            3.0,
+            "MULTIPOINT (2 6, 6 12)",
+        ),
+        (
+            "MULTILINESTRING ((0 0, 1 1), (2 2, 3 3))",
+            2.0,
+            3.0,
+            "MULTILINESTRING ((0 0, 2 3), (4 6, 6 9))",
+        ),
+        (
+            "MULTIPOLYGON (((0 0, 1 0, 1 1, 0 1, 0 0)))",
+            2.0,
+            3.0,
+            "MULTIPOLYGON (((0 0, 2 0, 2 3, 0 3, 0 0)))",
+        ),
+        (
+            "GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (0 0, 1 1))",
+            2.0,
+            3.0,
+            "GEOMETRYCOLLECTION (POINT (2 6), LINESTRING (0 0, 2 3))",
+        ),
+        ("POINT Z (1 2 3)", 2.0, 3.0, "POINT Z (2 6 3)"),
+    ],
+)
+def test_st_scale_2d(eng, geom, sx, sy, expected):
+    eng = eng.create_or_skip()
+    eng.assert_query_result(
+        f"SELECT ST_Scale({geom_or_null(geom)}, {val_or_null(sx)}, {val_or_null(sy)})",
+        expected,
+    )
+
+
+@pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
+@pytest.mark.parametrize(
+    ("geom", "sx", "sy", "sz", "expected"),
+    [
+        ("POINT Z (1 2 3)", 1.0, 1.0, 1.0, "POINT Z (1 2 3)"),
+        ("POINT Z (1 2 3)", 2.0, 3.0, 4.0, "POINT Z (2 6 12)"),
+    ],
+)
+def test_st_scale_3d(eng, geom, sx, sy, sz, expected):
+    eng = eng.create_or_skip()
+    query = (
+        "SELECT ST_Scale("
+        f"{geom_or_null(geom)}, {val_or_null(sx)}, {val_or_null(sy)}, {val_or_null(sz)})"
+    )
+    try:
+        eng.assert_query_result(query, expected)
+    except Exception as exc:
+        if isinstance(eng, PostGIS):
+            pytest.skip(f"PostGIS may not support 3D ST_Scale: {exc}")
+        raise
+
+
+@pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
+@pytest.mark.parametrize(
+    ("geom", "angle", "expected_x", "expected_y"),
+    [
+        (None, 0.0, None, None),
+        ("POINT (1 2)", 0.0, 1.0, 2.0),
+        ("POINT (1 2)", 1.5707963267948966, -2.0, 1.0),
+        ("POINT (1 2)", 3.141592653589793, -1.0, -2.0),
+    ],
+)
+def test_st_rotate(eng, geom, angle, expected_x, expected_y):
+    eng = eng.create_or_skip()
+    x_query = f"SELECT ST_X(ST_Rotate({geom_or_null(geom)}, {val_or_null(angle)}))"
+    y_query = f"SELECT ST_Y(ST_Rotate({geom_or_null(geom)}, {val_or_null(angle)}))"
+    eng.assert_query_result(x_query, expected_x, numeric_epsilon=1e-12)
+    eng.assert_query_result(y_query, expected_y, numeric_epsilon=1e-12)
+
+
 @pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
 @pytest.mark.parametrize(
     ("geom", "expected_boundary"),

rust/sedona-functions/Cargo.toml

@@ -52,6 +52,7 @@ sedona-schema = { workspace = true }
 wkb = { workspace = true }
 wkt = { workspace = true }
 serde_json = { workspace = true }
+glam = { workspace = true }
 
 [[bench]]
 harness = false

rust/sedona-functions/src/lib.rs

@@ -23,6 +23,8 @@ mod referencing;
 pub mod register;
 mod sd_format;
 pub mod sd_order;
+mod st_affine;
+mod st_affine_helpers;
 pub mod st_analyze_agg;
 mod st_area;
 mod st_asbinary;
@@ -60,6 +62,8 @@ mod st_points;
 mod st_pointzm;
 mod st_polygonize_agg;
 mod st_reverse;
+mod st_rotate;
+mod st_scale;
 mod st_setsrid;
 mod st_srid;
 mod st_start_point;

rust/sedona-functions/src/register.rs

@@ -63,6 +63,7 @@ pub fn default_function_set() -> FunctionSet {
         crate::referencing::st_line_locate_point_udf,
         crate::sd_format::sd_format_udf,
         crate::sd_order::sd_order_udf,
+        crate::st_affine::st_affine_udf,
         crate::st_area::st_area_udf,
         crate::st_asbinary::st_asbinary_udf,
         crate::st_asgeojson::st_asgeojson_udf,
@@ -102,6 +103,10 @@ pub fn default_function_set() -> FunctionSet {
         crate::st_pointzm::st_pointz_udf,
         crate::st_pointzm::st_pointzm_udf,
         crate::st_reverse::st_reverse_udf,
+        crate::st_rotate::st_rotate_udf,
+        crate::st_rotate::st_rotate_x_udf,
+        crate::st_rotate::st_rotate_y_udf,
+        crate::st_scale::st_scale_udf,
         crate::st_setsrid::st_set_crs_udf,
         crate::st_setsrid::st_set_srid_udf,
         crate::st_srid::st_crs_udf,