This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new 3f4acd0f feat(rust/sedona-functions): Implement ST_Affine(),
ST_Rotate(), and ST_Scale() (#504)
3f4acd0f is described below
commit 3f4acd0f074e78a0bd7e2e365e8912df0710a718
Author: Hiroaki Yutani <[email protected]>
AuthorDate: Wed Jan 14 11:20:07 2026 +0900
feat(rust/sedona-functions): Implement ST_Affine(), ST_Rotate(), and
ST_Scale() (#504)
Co-authored-by: Dewey Dunnington <[email protected]>
---
Cargo.lock | 7 +
Cargo.toml | 1 +
python/sedonadb/tests/functions/test_functions.py | 276 ++++++++++
rust/sedona-functions/Cargo.toml | 1 +
rust/sedona-functions/src/lib.rs | 4 +
rust/sedona-functions/src/register.rs | 5 +
rust/sedona-functions/src/st_affine.rs | 453 ++++++++++++++++
rust/sedona-functions/src/st_affine_helpers.rs | 634 ++++++++++++++++++++++
rust/sedona-functions/src/st_rotate.rs | 261 +++++++++
rust/sedona-functions/src/st_scale.rs | 348 ++++++++++++
rust/sedona-geometry/src/transform.rs | 108 +++-
11 files changed, 2082 insertions(+), 16 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index aa9c05ec..ca89f4e9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2853,6 +2853,12 @@ dependencies = [
"wasm-bindgen",
]
+[[package]]
+name = "glam"
+version = "0.30.10"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9"
+
[[package]]
name = "glob"
version = "0.3.3"
@@ -5089,6 +5095,7 @@ dependencies = [
"datafusion-common",
"datafusion-expr",
"geo-traits",
+ "glam",
"rstest",
"sedona-common",
"sedona-expr",
diff --git a/Cargo.toml b/Cargo.toml
index 9105526f..1e69ef66 100644
--- a/Cargo.toml
+++ b/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"] }
diff --git a/python/sedonadb/tests/functions/test_functions.py
b/python/sedonadb/tests/functions/test_functions.py
index ff566ccc..28f2e901 100644
--- a/python/sedonadb/tests/functions/test_functions.py
+++ b/python/sedonadb/tests/functions/test_functions.py
@@ -17,6 +17,7 @@
import pytest
import shapely
from sedonadb.testing import PostGIS, SedonaDB, geom_or_null, val_or_null
+import math
@pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
@@ -189,6 +190,281 @@ def test_st_azimuth(eng, geom1, geom2, expected):
)
+# fmt: off
[email protected]("eng", [SedonaDB, PostGIS])
[email protected](
+ ("geom", "a", "b", "d", "e", "xoff", "yoff", "expected"),
+ [
+ (
+ None,
+ 1.0, 0.0,
+ 0.0, 2.0,
+ 1.0, 3.0,
+ None
+ ),
+ (
+ "POINT (1 2)",
+ None, 0.0,
+ 0.0, 2.0,
+ 1.0, 3.0,
+ None
+ ),
+ (
+ "POINT (1 2)",
+ 1.0, 0.0,
+ 0.0, 2.0,
+ 1.0, None,
+ None
+ ),
+ (
+ "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
[email protected]("eng", [SedonaDB, PostGIS])
[email protected](
+ ("geom", "a", "b", "c", "d", "e", "f", "g", "h", "i", "xoff", "yoff",
"zoff", "expected"),
+ [
+ (
+ None,
+ 1.0, 0.0, 0.0,
+ 0.0, 2.0, 0.0,
+ 0.0, 0.0, 2.0,
+ 1.0, 3.0, 5.0,
+ None
+ ),
+ (
+ "POINT Z (1 2 3)",
+ None, 0.0, 0.0,
+ 0.0, 2.0, 0.0,
+ 0.0, 0.0, 2.0,
+ 1.0, 3.0, 5.0,
+ None
+ ),
+ (
+ "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, None,
+ None
+ ),
+ (
+ "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)})"
+ )
+ eng.assert_query_result(query, expected)
+# fmt: on
+
+
[email protected]("eng", [SedonaDB, PostGIS])
[email protected](
+ ("geom", "sx", "sy", "expected"),
+ [
+ (None, 1.0, 1.0, None),
+ ("POINT (1 2)", None, 1.0, None),
+ ("POINT (1 2)", 1.0, None, None),
+ ("POINT EMPTY", 1.0, 1.0, "POINT (nan nan)"),
+ ("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)"),
+ ("POINT M (1 2 3)", 2.0, 3.0, "POINT M (2 6 3)"),
+ ("POINT ZM (1 2 3 4)", 2.0, 3.0, "POINT ZM (2 6 3 4)"),
+ ],
+)
+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,
+ )
+
+
[email protected]("eng", [SedonaDB, PostGIS])
[email protected](
+ ("geom", "sx", "sy", "sz", "expected"),
+ [
+ (None, 1.0, 1.0, 1.0, None),
+ ("POINT Z (1 2 3)", None, 1.0, 1.0, None),
+ ("POINT Z (1 2 3)", 1.0, 1.0, None, None),
+ ("POINT EMPTY", 1.0, 1.0, 1.0, "POINT (nan nan)"),
+ ("POINT Z EMPTY", 1.0, 1.0, 1.0, "POINT Z (nan nan nan)"),
+ ("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)"),
+ ("POINT ZM (1 2 3 4)", 2.0, 3.0, 4.0, "POINT ZM (2 6 12 4)"),
+ ("LINESTRING Z (0 0 0, 1 1 1)", 2.0, 3.0, 4.0, "LINESTRING Z (0 0 0, 2
3 4)"),
+ (
+ "POLYGON Z ((0 0 0, 1 0 2, 1 1 4, 0 1 2, 0 0 0))",
+ 2.0,
+ 3.0,
+ 4.0,
+ "POLYGON Z ((0 0 0, 2 0 8, 2 3 16, 0 3 8, 0 0 0))",
+ ),
+ ("POINT (1 2)", 2.0, 3.0, 4.0, "POINT (2 6)"),
+ ("POINT M (1 2 3)", 2.0, 3.0, 4.0, "POINT M (2 6 3)"),
+ ],
+)
+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)})"
+ )
+ eng.assert_query_result(query, expected)
+
+
[email protected]("eng", [SedonaDB, PostGIS])
[email protected](
+ ("geom", "angle", "expected"),
+ [
+ (None, 0, None),
+ ("POINT (1 2)", None, None),
+ ("POINT EMPTY", 0, "POINT (nan nan)"),
+ ("POINT Z EMPTY", 0, "POINT Z (nan nan nan)"),
+ ("POINT (1 2)", 0, "POINT (1 2)"),
+ ("POINT (1 2)", math.pi / 2, "POINT (-2 1)"),
+ ("POINT (1 2)", math.pi, "POINT (-1 -2)"),
+ ("POINT Z (1 2 3)", math.pi, "POINT Z (-1 -2 3)"),
+ ("POINT M (1 2 3)", math.pi, "POINT M (-1 -2 3)"),
+ ("POINT ZM (1 2 3 4)", math.pi, "POINT ZM (-1 -2 3 4)"),
+ ("LINESTRING (0 0, 1 2)", math.pi, "LINESTRING (0 0, -1 -2)"),
+ ("LINESTRING Z (0 0 0, 1 2 3)", math.pi, "LINESTRING Z (0 0 0, -1 -2
3)"),
+ (
+ "POLYGON ((0 0, 1 2, 2 3, 2 1, 0 0))",
+ math.pi,
+ "POLYGON ((0 0, -1 -2, -2 -3, -2 -1, 0 0))",
+ ),
+ (
+ "POLYGON Z ((0 0 0, 1 2 4, 2 3 4, 2 1 4, 0 0 0))",
+ math.pi,
+ "POLYGON Z ((0 0 0, -1 -2 4, -2 -3 4, -2 -1 4, 0 0 0))",
+ ),
+ ],
+)
+def test_st_rotate(eng, geom, angle, expected):
+ eng = eng.create_or_skip()
+ query = f"SELECT ST_Rotate({geom_or_null(geom)}, {val_or_null(angle)})"
+ eng.assert_query_result(query, expected, wkt_precision=1e-12)
+
+
[email protected]("eng", [SedonaDB, PostGIS])
[email protected](
+ ("geom", "angle", "expected"),
+ [
+ (None, 0, None),
+ ("POINT (1 2)", None, None),
+ ("POINT Z (1 2 3)", math.pi, "POINT Z (1 -2 -3)"),
+ ],
+)
+def test_st_rotate_x(eng, geom, angle, expected):
+ eng = eng.create_or_skip()
+ query = f"SELECT ST_RotateX({geom_or_null(geom)}, {val_or_null(angle)})"
+ eng.assert_query_result(query, expected, wkt_precision=1e-12)
+
+
[email protected]("eng", [SedonaDB, PostGIS])
[email protected](
+ ("geom", "angle", "expected"),
+ [
+ (None, 0, None),
+ ("POINT (1 2)", None, None),
+ ("POINT Z (1 2 3)", math.pi, "POINT Z (-1 2 -3)"),
+ ],
+)
+def test_st_rotate_y(eng, geom, angle, expected):
+ eng = eng.create_or_skip()
+ query = f"SELECT ST_RotateY({geom_or_null(geom)}, {val_or_null(angle)})"
+ eng.assert_query_result(query, expected, wkt_precision=1e-12)
+
+
@pytest.mark.parametrize("eng", [SedonaDB, PostGIS])
@pytest.mark.parametrize(
("geom", "expected_boundary"),
diff --git a/rust/sedona-functions/Cargo.toml b/rust/sedona-functions/Cargo.toml
index 4d0cedbe..41b26e24 100644
--- a/rust/sedona-functions/Cargo.toml
+++ b/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
diff --git a/rust/sedona-functions/src/lib.rs b/rust/sedona-functions/src/lib.rs
index bef42fff..6e38e426 100644
--- a/rust/sedona-functions/src/lib.rs
+++ b/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;
diff --git a/rust/sedona-functions/src/register.rs
b/rust/sedona-functions/src/register.rs
index 32549b67..80dbbd82 100644
--- a/rust/sedona-functions/src/register.rs
+++ b/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,
diff --git a/rust/sedona-functions/src/st_affine.rs
b/rust/sedona-functions/src/st_affine.rs
new file mode 100644
index 00000000..c831b710
--- /dev/null
+++ b/rust/sedona-functions/src/st_affine.rs
@@ -0,0 +1,453 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+use arrow_array::{builder::BinaryBuilder, Array};
+use arrow_schema::DataType;
+use datafusion_common::{error::Result, DataFusionError};
+use datafusion_expr::{
+ scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation,
Volatility,
+};
+use sedona_expr::{
+ item_crs::ItemCrsKernel,
+ scalar_udf::{SedonaScalarKernel, SedonaScalarUDF},
+};
+use sedona_geometry::{transform::transform,
wkb_factory::WKB_MIN_PROBABLE_BYTES};
+use sedona_schema::{
+ datatypes::{SedonaType, WKB_GEOMETRY},
+ matchers::ArgMatcher,
+};
+use std::sync::Arc;
+
+use crate::{
+ executor::WkbExecutor,
+ st_affine_helpers::{self},
+};
+
+/// ST_Affine() scalar UDF
+///
+/// Native implementation for affine transformation
+pub fn st_affine_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "st_affine",
+ ItemCrsKernel::wrap_impl(vec![
+ Arc::new(STAffine { is_3d: true }),
+ Arc::new(STAffine { is_3d: false }),
+ ]),
+ Volatility::Immutable,
+ Some(st_affine_doc()),
+ )
+}
+
+fn st_affine_doc() -> Documentation {
+ Documentation::builder(
+ DOC_SECTION_OTHER,
+ "Apply an affine transformation to the given geometry.",
+ "ST_Affine (geom: Geometry, a: Double, b: Double, c: Double, d:
Double, e: Double, f: Double, g: Double, h: Double, i: Double, xOff: Double,
yOff: Double, zOff: Double)",
+ )
+ .with_argument("geom", "geometry: Input geometry")
+ .with_argument("a", "a component of the affine matrix")
+ .with_argument("b", "a component of the affine matrix")
+ .with_argument("c", "a component of the affine matrix")
+ .with_argument("d", "a component of the affine matrix")
+ .with_argument("e", "a component of the affine matrix")
+ .with_argument("f", "a component of the affine matrix")
+ .with_argument("g", "a component of the affine matrix")
+ .with_argument("h", "a component of the affine matrix")
+ .with_argument("i", "a component of the affine matrix")
+ .with_argument("xOff", "X offset")
+ .with_argument("yOff", "Y offset")
+ .with_argument("zOff", "Z offset")
+ .with_sql_example("SELECT ST_Affine(ST_GeomFromText('POLYGON Z ((1 0 1, 1
1 1, 2 2 2, 1 0 1))'), 1, 2, 4, 1, 1, 2, 3, 2, 5, 4, 8, 3)")
+ .build()
+}
+
+#[derive(Debug)]
+struct STAffine {
+ is_3d: bool,
+}
+
+impl SedonaScalarKernel for STAffine {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let arg_matchers = if self.is_3d {
+ vec![
+ ArgMatcher::is_geometry(),
+ ArgMatcher::is_numeric(), // a
+ ArgMatcher::is_numeric(), // b
+ ArgMatcher::is_numeric(), // c
+ ArgMatcher::is_numeric(), // d
+ ArgMatcher::is_numeric(), // e
+ ArgMatcher::is_numeric(), // f
+ ArgMatcher::is_numeric(), // g
+ ArgMatcher::is_numeric(), // h
+ ArgMatcher::is_numeric(), // i
+ ArgMatcher::is_numeric(), // xOff
+ ArgMatcher::is_numeric(), // yOff
+ ArgMatcher::is_numeric(), // zOff
+ ]
+ } else {
+ vec![
+ ArgMatcher::is_geometry(),
+ ArgMatcher::is_numeric(), // a
+ ArgMatcher::is_numeric(), // b
+ ArgMatcher::is_numeric(), // d
+ ArgMatcher::is_numeric(), // e
+ ArgMatcher::is_numeric(), // xOff
+ ArgMatcher::is_numeric(), // yOff
+ ]
+ };
+
+ let matcher = ArgMatcher::new(arg_matchers, WKB_GEOMETRY);
+
+ matcher.match_args(args)
+ }
+
+ fn invoke_batch(
+ &self,
+ arg_types: &[SedonaType],
+ args: &[ColumnarValue],
+ ) -> Result<ColumnarValue> {
+ let executor = WkbExecutor::new(arg_types, args);
+ let mut builder = BinaryBuilder::with_capacity(
+ executor.num_iterations(),
+ WKB_MIN_PROBABLE_BYTES * executor.num_iterations(),
+ );
+
+ let array_args = args[1..]
+ .iter()
+ .map(|arg| {
+ arg.cast_to(&DataType::Float64, None)?
+ .to_array(executor.num_iterations())
+ })
+ .collect::<Result<Vec<Arc<dyn Array>>>>()?;
+
+ let mut affine_iter = if self.is_3d {
+ st_affine_helpers::DAffineIterator::new_3d(&array_args)?
+ } else {
+ st_affine_helpers::DAffineIterator::new_2d(&array_args)?
+ };
+
+ executor.execute_wkb_void(|maybe_wkb| {
+ let maybe_mat = affine_iter.next().unwrap();
+ match (maybe_wkb, maybe_mat) {
+ (Some(wkb), Some(mat)) => {
+ transform(&wkb, &mat, &mut builder)
+ .map_err(|e|
DataFusionError::Execution(e.to_string()))?;
+ builder.append_value([]);
+ }
+ _ => builder.append_null(),
+ }
+
+ Ok(())
+ })?;
+
+ executor.finish(Arc::new(builder.finish()))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDF};
+ use rstest::rstest;
+ use sedona_schema::datatypes::{WKB_GEOMETRY_ITEM_CRS, WKB_VIEW_GEOMETRY};
+ use sedona_testing::{
+ compare::assert_array_equal, create::create_array,
create::create_scalar,
+ testers::ScalarUdfTester,
+ };
+
+ use super::*;
+
+ #[test]
+ fn udf_metadata() {
+ let st_affine_udf: ScalarUDF = st_affine_udf().into();
+ assert_eq!(st_affine_udf.name(), "st_affine");
+ assert!(st_affine_udf.documentation().is_some());
+ }
+
+ #[rstest]
+ fn udf_2d(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type:
SedonaType) {
+ let tester_2d = ScalarUdfTester::new(
+ st_affine_udf().into(),
+ vec![
+ sedona_type.clone(),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ ],
+ );
+ tester_2d.assert_return_type(WKB_GEOMETRY);
+
+ let points = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (1 2)"),
+ Some("POINT M (1 2 3)"),
+ ],
+ &sedona_type,
+ );
+
+ // identity transformation
+
+ #[rustfmt::skip]
+ let m_identity = &[
+ Some(1.0), Some(0.0),
+ Some(0.0), Some(1.0),
+ Some(0.0), Some(0.0),
+ ];
+
+ let expected_identity = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (1 2)"),
+ Some("POINT M (1 2 3)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_identity = tester_2d
+ .invoke_arrays(prepare_args(points.clone(), m_identity))
+ .unwrap();
+ assert_array_equal(&result_identity, &expected_identity);
+
+ // scale transformation
+
+ #[rustfmt::skip]
+ let m_scale = &[
+ Some(10.0), Some(0.0),
+ Some(0.0), Some(10.0),
+ Some(0.0), Some(0.0),
+ ];
+
+ let expected_scale = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (10 20)"),
+ Some("POINT M (10 20 3)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_scale = tester_2d
+ .invoke_arrays(prepare_args(points.clone(), m_scale))
+ .unwrap();
+ assert_array_equal(&result_scale, &expected_scale);
+
+ // 2D matrix with 3D input (z/m preserved)
+ let points_3d = create_array(
+ &[
+ None,
+ Some("POINT Z EMPTY"),
+ Some("POINT ZM EMPTY"),
+ Some("POINT Z (1 2 3)"),
+ Some("POINT ZM (1 2 3 4)"),
+ ],
+ &sedona_type,
+ );
+
+ let expected_scale_3d = create_array(
+ &[
+ None,
+ Some("POINT Z EMPTY"),
+ Some("POINT ZM EMPTY"),
+ Some("POINT Z (10 20 3)"),
+ Some("POINT ZM (10 20 3 4)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_scale_3d = tester_2d
+ .invoke_arrays(prepare_args(points_3d, m_scale))
+ .unwrap();
+ assert_array_equal(&result_scale_3d, &expected_scale_3d);
+ }
+
+ #[rstest]
+ fn udf_3d(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type:
SedonaType) {
+ let tester_3d = ScalarUdfTester::new(
+ st_affine_udf().into(),
+ vec![
+ sedona_type.clone(),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ ],
+ );
+ tester_3d.assert_return_type(WKB_GEOMETRY);
+
+ let points = create_array(
+ &[
+ None,
+ Some("POINT Z EMPTY"),
+ Some("POINT ZM EMPTY"),
+ Some("POINT Z (1 2 3)"),
+ Some("POINT ZM (1 2 3 4)"),
+ ],
+ &sedona_type,
+ );
+
+ // identity matrix
+ #[rustfmt::skip]
+ let m_identity = &[
+ Some(1.0), Some(0.0), Some(0.0),
+ Some(0.0), Some(1.0), Some(0.0),
+ Some(0.0), Some(0.0), Some(1.0),
+ Some(0.0), Some(0.0), Some(0.0),
+ ];
+
+ let expected_identity = create_array(
+ &[
+ None,
+ Some("POINT Z EMPTY"),
+ Some("POINT ZM EMPTY"),
+ Some("POINT Z (1 2 3)"),
+ Some("POINT ZM (1 2 3 4)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_identity = tester_3d
+ .invoke_arrays(prepare_args(points.clone(), m_identity))
+ .unwrap();
+ assert_array_equal(&result_identity, &expected_identity);
+
+ // scale transformation
+ #[rustfmt::skip]
+ let m_scale = &[
+ Some(10.0), Some(0.0), Some(0.0),
+ Some(0.0), Some(10.0), Some(0.0),
+ Some(0.0), Some(0.0), Some(10.0),
+ Some(0.0), Some(0.0), Some(0.0),
+ ];
+
+ let expected_scale = create_array(
+ &[
+ None,
+ Some("POINT Z EMPTY"),
+ Some("POINT ZM EMPTY"),
+ Some("POINT Z (10 20 30)"),
+ Some("POINT ZM (10 20 30 4)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_scale = tester_3d
+ .invoke_arrays(prepare_args(points, m_scale))
+ .unwrap();
+ assert_array_equal(&result_scale, &expected_scale);
+
+ // 3D matrix with 2D input (z translation ignored, m preserved)
+ let points_2d = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (1 2)"),
+ Some("POINT M (1 2 3)"),
+ ],
+ &sedona_type,
+ );
+
+ #[rustfmt::skip]
+ let m_translate = &[
+ Some(1.0), Some(0.0), Some(0.0),
+ Some(0.0), Some(1.0), Some(0.0),
+ Some(0.0), Some(0.0), Some(1.0),
+ Some(10.0), Some(20.0), Some(30.0),
+ ];
+
+ let expected_translate = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (11 22)"),
+ Some("POINT M (11 22 3)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_translate = tester_3d
+ .invoke_arrays(prepare_args(points_2d, m_translate))
+ .unwrap();
+ assert_array_equal(&result_translate, &expected_translate);
+ }
+
+ fn prepare_args(wkt: Arc<dyn Array>, mat: &[Option<f64>]) -> Vec<Arc<dyn
Array>> {
+ let n = wkt.len();
+ let mut args: Vec<Arc<dyn Array>> = mat
+ .iter()
+ .map(|a| {
+ let values = vec![*a; n];
+ Arc::new(arrow_array::Float64Array::from(values)) as Arc<dyn
Array>
+ })
+ .collect();
+ args.insert(0, wkt);
+ args
+ }
+
+ #[rstest]
+ fn udf_invoke_item_crs(#[values(WKB_GEOMETRY_ITEM_CRS.clone())]
sedona_type: SedonaType) {
+ let tester = ScalarUdfTester::new(
+ st_affine_udf().into(),
+ vec![
+ sedona_type.clone(),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ ],
+ );
+ tester.assert_return_type(sedona_type.clone());
+
+ let geom = create_scalar(Some("POINT (1 2)"), &sedona_type);
+ let args = vec![
+ ColumnarValue::Scalar(geom),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(1.0))),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(3.0))),
+ ];
+
+ let result = tester.invoke(args).unwrap();
+ if let ColumnarValue::Scalar(scalar) = result {
+ tester.assert_scalar_result_equals(scalar, "POINT (3 7)");
+ } else {
+ panic!("Expected scalar result from item CRS affine invoke");
+ }
+ }
+}
diff --git a/rust/sedona-functions/src/st_affine_helpers.rs
b/rust/sedona-functions/src/st_affine_helpers.rs
new file mode 100644
index 00000000..7de1ba94
--- /dev/null
+++ b/rust/sedona-functions/src/st_affine_helpers.rs
@@ -0,0 +1,634 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+use arrow_array::types::Float64Type;
+use arrow_array::Array;
+use arrow_array::PrimitiveArray;
+use datafusion_common::cast::as_float64_array;
+use datafusion_common::error::Result;
+use sedona_common::sedona_internal_err;
+use sedona_geometry::transform::CrsTransform;
+use std::sync::Arc;
+
+pub(crate) struct DAffine2Iterator<'a> {
+ index: usize,
+ a: &'a PrimitiveArray<Float64Type>,
+ b: &'a PrimitiveArray<Float64Type>,
+ d: &'a PrimitiveArray<Float64Type>,
+ e: &'a PrimitiveArray<Float64Type>,
+ x_offset: &'a PrimitiveArray<Float64Type>,
+ y_offset: &'a PrimitiveArray<Float64Type>,
+ no_null: bool,
+}
+
+impl<'a> DAffine2Iterator<'a> {
+ pub(crate) fn new(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ if array_args.len() != 6 {
+ return sedona_internal_err!("Invalid number of arguments are
passed");
+ }
+
+ let a = as_float64_array(&array_args[0])?;
+ let b = as_float64_array(&array_args[1])?;
+ let d = as_float64_array(&array_args[2])?;
+ let e = as_float64_array(&array_args[3])?;
+ let x_offset = as_float64_array(&array_args[4])?;
+ let y_offset = as_float64_array(&array_args[5])?;
+
+ Ok(Self {
+ index: 0,
+ a,
+ b,
+ d,
+ e,
+ x_offset,
+ y_offset,
+ no_null: a.null_count() == 0
+ && b.null_count() == 0
+ && d.null_count() == 0
+ && e.null_count() == 0
+ && x_offset.null_count() == 0
+ && y_offset.null_count() == 0,
+ })
+ }
+
+ fn is_null(&self, i: usize) -> bool {
+ if self.no_null {
+ return false;
+ }
+
+ self.a.is_null(i)
+ || self.b.is_null(i)
+ || self.d.is_null(i)
+ || self.e.is_null(i)
+ || self.x_offset.is_null(i)
+ || self.y_offset.is_null(i)
+ }
+}
+
+impl<'a> Iterator for DAffine2Iterator<'a> {
+ // As this needs to distinguish NULL, next() returns Some(Some(value))
+ type Item = Option<glam::DAffine2>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+
+ if self.is_null(i) {
+ return Some(None);
+ }
+
+ Some(Some(glam::DAffine2 {
+ matrix2: glam::DMat2 {
+ x_axis: glam::DVec2 {
+ x: self.a.value(i),
+ y: self.b.value(i),
+ },
+ y_axis: glam::DVec2 {
+ x: self.d.value(i),
+ y: self.e.value(i),
+ },
+ },
+ translation: glam::DVec2 {
+ x: self.x_offset.value(i),
+ y: self.y_offset.value(i),
+ },
+ }))
+ }
+}
+
+pub(crate) struct DAffine3Iterator<'a> {
+ index: usize,
+ a: &'a PrimitiveArray<Float64Type>,
+ b: &'a PrimitiveArray<Float64Type>,
+ c: &'a PrimitiveArray<Float64Type>,
+ d: &'a PrimitiveArray<Float64Type>,
+ e: &'a PrimitiveArray<Float64Type>,
+ f: &'a PrimitiveArray<Float64Type>,
+ g: &'a PrimitiveArray<Float64Type>,
+ h: &'a PrimitiveArray<Float64Type>,
+ i: &'a PrimitiveArray<Float64Type>,
+ x_offset: &'a PrimitiveArray<Float64Type>,
+ y_offset: &'a PrimitiveArray<Float64Type>,
+ z_offset: &'a PrimitiveArray<Float64Type>,
+ no_null: bool,
+}
+
+impl<'a> DAffine3Iterator<'a> {
+ pub(crate) fn new(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ if array_args.len() != 12 {
+ return sedona_internal_err!("Invalid number of arguments are
passed");
+ }
+
+ let a = as_float64_array(&array_args[0])?;
+ let b = as_float64_array(&array_args[1])?;
+ let c = as_float64_array(&array_args[2])?;
+ let d = as_float64_array(&array_args[3])?;
+ let e = as_float64_array(&array_args[4])?;
+ let f = as_float64_array(&array_args[5])?;
+ let g = as_float64_array(&array_args[6])?;
+ let h = as_float64_array(&array_args[7])?;
+ let i = as_float64_array(&array_args[8])?;
+ let x_offset = as_float64_array(&array_args[9])?;
+ let y_offset = as_float64_array(&array_args[10])?;
+ let z_offset = as_float64_array(&array_args[11])?;
+
+ Ok(Self {
+ index: 0,
+ a,
+ b,
+ c,
+ d,
+ e,
+ f,
+ g,
+ h,
+ i,
+ x_offset,
+ y_offset,
+ z_offset,
+ no_null: a.null_count() == 0
+ && b.null_count() == 0
+ && c.null_count() == 0
+ && d.null_count() == 0
+ && e.null_count() == 0
+ && f.null_count() == 0
+ && g.null_count() == 0
+ && h.null_count() == 0
+ && i.null_count() == 0
+ && x_offset.null_count() == 0
+ && y_offset.null_count() == 0
+ && z_offset.null_count() == 0,
+ })
+ }
+
+ fn is_null(&self, i: usize) -> bool {
+ if self.no_null {
+ return false;
+ }
+
+ self.a.is_null(i)
+ || self.b.is_null(i)
+ || self.c.is_null(i)
+ || self.d.is_null(i)
+ || self.e.is_null(i)
+ || self.f.is_null(i)
+ || self.g.is_null(i)
+ || self.h.is_null(i)
+ || self.i.is_null(i)
+ || self.x_offset.is_null(i)
+ || self.y_offset.is_null(i)
+ || self.z_offset.is_null(i)
+ }
+}
+
+impl<'a> Iterator for DAffine3Iterator<'a> {
+ // As this needs to distinguish NULL, next() returns Some(Some(value))
+ type Item = Option<glam::DAffine3>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+
+ if self.is_null(i) {
+ return Some(None);
+ }
+
+ Some(Some(glam::DAffine3 {
+ matrix3: glam::DMat3 {
+ x_axis: glam::DVec3 {
+ x: self.a.value(i),
+ y: self.b.value(i),
+ z: self.c.value(i),
+ },
+ y_axis: glam::DVec3 {
+ x: self.d.value(i),
+ y: self.e.value(i),
+ z: self.f.value(i),
+ },
+ z_axis: glam::DVec3 {
+ x: self.g.value(i),
+ y: self.h.value(i),
+ z: self.i.value(i),
+ },
+ },
+ translation: glam::DVec3 {
+ x: self.x_offset.value(i),
+ y: self.y_offset.value(i),
+ z: self.z_offset.value(i),
+ },
+ }))
+ }
+}
+
+pub(crate) struct DAffine2ScaleIterator<'a> {
+ index: usize,
+ x_scale: &'a PrimitiveArray<Float64Type>,
+ y_scale: &'a PrimitiveArray<Float64Type>,
+ no_null: bool,
+}
+
+impl<'a> DAffine2ScaleIterator<'a> {
+ pub(crate) fn new(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ if array_args.len() != 2 {
+ return sedona_internal_err!("Invalid number of arguments are
passed");
+ }
+
+ let x_scale = as_float64_array(&array_args[0])?;
+ let y_scale = as_float64_array(&array_args[1])?;
+
+ Ok(Self {
+ index: 0,
+ x_scale,
+ y_scale,
+ no_null: x_scale.null_count() == 0 && y_scale.null_count() == 0,
+ })
+ }
+
+ fn is_null(&self, i: usize) -> bool {
+ if self.no_null {
+ return false;
+ }
+
+ self.x_scale.is_null(i) || self.y_scale.is_null(i)
+ }
+}
+
+impl<'a> Iterator for DAffine2ScaleIterator<'a> {
+ // As this needs to distinguish NULL, next() returns Some(Some(value))
+ type Item = Option<glam::DAffine2>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+
+ if self.is_null(i) {
+ return Some(None);
+ }
+
+ let scale = glam::DVec2::new(self.x_scale.value(i),
self.y_scale.value(i));
+ Some(Some(glam::DAffine2::from_scale(scale)))
+ }
+}
+
+pub(crate) struct DAffine3ScaleIterator<'a> {
+ index: usize,
+ x_scale: &'a PrimitiveArray<Float64Type>,
+ y_scale: &'a PrimitiveArray<Float64Type>,
+ z_scale: &'a PrimitiveArray<Float64Type>,
+ no_null: bool,
+}
+
+impl<'a> DAffine3ScaleIterator<'a> {
+ pub(crate) fn new(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ if array_args.len() != 3 {
+ return sedona_internal_err!("Invalid number of arguments are
passed");
+ }
+
+ let x_scale = as_float64_array(&array_args[0])?;
+ let y_scale = as_float64_array(&array_args[1])?;
+ let z_scale = as_float64_array(&array_args[2])?;
+
+ Ok(Self {
+ index: 0,
+ x_scale,
+ y_scale,
+ z_scale,
+ no_null: x_scale.null_count() == 0
+ && y_scale.null_count() == 0
+ && z_scale.null_count() == 0,
+ })
+ }
+
+ fn is_null(&self, i: usize) -> bool {
+ if self.no_null {
+ return false;
+ }
+
+ self.x_scale.is_null(i) || self.y_scale.is_null(i) ||
self.z_scale.is_null(i)
+ }
+}
+
+impl<'a> Iterator for DAffine3ScaleIterator<'a> {
+ // As this needs to distinguish NULL, next() returns Some(Some(value))
+ type Item = Option<glam::DAffine3>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+
+ if self.is_null(i) {
+ return Some(None);
+ }
+
+ let scale = glam::DVec3::new(
+ self.x_scale.value(i),
+ self.y_scale.value(i),
+ self.z_scale.value(i),
+ );
+ Some(Some(glam::DAffine3::from_scale(scale)))
+ }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) enum RotateAxis {
+ X,
+ Y,
+ Z,
+}
+
+pub(crate) struct DAffineRotateIterator<'a> {
+ index: usize,
+ angle: &'a PrimitiveArray<Float64Type>,
+ axis: RotateAxis,
+ no_null: bool,
+}
+
+impl<'a> DAffineRotateIterator<'a> {
+ pub(crate) fn new(angle: &'a Arc<dyn Array>, axis: RotateAxis) ->
Result<Self> {
+ let angle = as_float64_array(angle)?;
+ Ok(Self {
+ index: 0,
+ angle,
+ axis,
+ no_null: angle.null_count() == 0,
+ })
+ }
+
+ fn is_null(&self, i: usize) -> bool {
+ if self.no_null {
+ return false;
+ }
+
+ self.angle.is_null(i)
+ }
+}
+
+impl<'a> Iterator for DAffineRotateIterator<'a> {
+ // As this needs to distinguish NULL, next() returns Some(Some(value))
+ type Item = Option<glam::DAffine3>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+
+ if self.is_null(i) {
+ return Some(None);
+ }
+
+ match self.axis {
+ RotateAxis::X =>
Some(Some(glam::DAffine3::from_rotation_x(self.angle.value(i)))),
+ RotateAxis::Y =>
Some(Some(glam::DAffine3::from_rotation_y(self.angle.value(i)))),
+ RotateAxis::Z =>
Some(Some(glam::DAffine3::from_rotation_z(self.angle.value(i)))),
+ }
+ }
+}
+
+pub(crate) enum DAffineIterator<'a> {
+ DAffine2(DAffine2Iterator<'a>),
+ DAffine3(DAffine3Iterator<'a>),
+ DAffine2Scale(DAffine2ScaleIterator<'a>),
+ DAffine3Scale(DAffine3ScaleIterator<'a>),
+ DAffineRotate(DAffineRotateIterator<'a>),
+}
+
+impl<'a> DAffineIterator<'a> {
+ pub(crate) fn new_2d(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ Ok(Self::DAffine2(DAffine2Iterator::new(array_args)?))
+ }
+
+ pub(crate) fn new_3d(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ Ok(Self::DAffine3(DAffine3Iterator::new(array_args)?))
+ }
+
+ pub(crate) fn from_scale_2d(array_args: &'a [Arc<dyn Array>]) ->
Result<Self> {
+ Ok(Self::DAffine2Scale(DAffine2ScaleIterator::new(array_args)?))
+ }
+
+ pub(crate) fn from_scale_3d(array_args: &'a [Arc<dyn Array>]) ->
Result<Self> {
+ Ok(Self::DAffine3Scale(DAffine3ScaleIterator::new(array_args)?))
+ }
+
+ pub(crate) fn from_angle(angle: &'a Arc<dyn Array>, axis: RotateAxis) ->
Result<Self> {
+ Ok(Self::DAffineRotate(DAffineRotateIterator::new(
+ angle, axis,
+ )?))
+ }
+}
+
+#[derive(Debug, PartialEq)]
+pub(crate) enum DAffine {
+ DAffine2(glam::DAffine2),
+ DAffine3(glam::DAffine3),
+}
+
+impl CrsTransform for DAffine {
+ fn transform_coord_3d(
+ &self,
+ coord: &mut (f64, f64, f64),
+ ) -> std::result::Result<(), sedona_geometry::error::SedonaGeometryError> {
+ match self {
+ DAffine::DAffine2(daffine2) => {
+ let transformed = daffine2.transform_point2(glam::DVec2 {
+ x: coord.0,
+ y: coord.1,
+ });
+ coord.0 = transformed.x;
+ coord.1 = transformed.y;
+ }
+ DAffine::DAffine3(daffine3) => {
+ let transformed = daffine3.transform_point3(glam::DVec3 {
+ x: coord.0,
+ y: coord.1,
+ z: coord.2,
+ });
+ coord.0 = transformed.x;
+ coord.1 = transformed.y;
+ coord.2 = transformed.z;
+ }
+ }
+
+ Ok(())
+ }
+
+ fn transform_coord(
+ &self,
+ coord: &mut (f64, f64),
+ ) -> std::result::Result<(), sedona_geometry::error::SedonaGeometryError> {
+ match self {
+ DAffine::DAffine2(daffine2) => {
+ let transformed = daffine2.transform_point2(glam::DVec2 {
+ x: coord.0,
+ y: coord.1,
+ });
+ coord.0 = transformed.x;
+ coord.1 = transformed.y;
+ }
+ DAffine::DAffine3(daffine3) => {
+ let transformed = daffine3.transform_point3(glam::DVec3 {
+ x: coord.0,
+ y: coord.1,
+ z: 0.0,
+ });
+ coord.0 = transformed.x;
+ coord.1 = transformed.y;
+ }
+ }
+
+ Ok(())
+ }
+}
+
+impl<'a> Iterator for DAffineIterator<'a> {
+ type Item = Option<DAffine>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ match self {
+ DAffineIterator::DAffine2(daffine2_iterator) => match
daffine2_iterator.next() {
+ Some(Some(a)) => Some(Some(DAffine::DAffine2(a))),
+ Some(None) => Some(None),
+ None => None,
+ },
+ DAffineIterator::DAffine3(daffine3_iterator) => match
daffine3_iterator.next() {
+ Some(Some(a)) => Some(Some(DAffine::DAffine3(a))),
+ Some(None) => Some(None),
+ None => None,
+ },
+ DAffineIterator::DAffine2Scale(daffine2_scale_iterator) => {
+ match daffine2_scale_iterator.next() {
+ Some(Some(a)) => Some(Some(DAffine::DAffine2(a))),
+ Some(None) => Some(None),
+ None => None,
+ }
+ }
+ DAffineIterator::DAffine3Scale(daffine3_scale_iterator) => {
+ match daffine3_scale_iterator.next() {
+ Some(Some(a)) => Some(Some(DAffine::DAffine3(a))),
+ Some(None) => Some(None),
+ None => None,
+ }
+ }
+ DAffineIterator::DAffineRotate(daffine_rotate_iterator) => {
+ match daffine_rotate_iterator.next() {
+ Some(Some(a)) => Some(Some(DAffine::DAffine3(a))),
+ Some(None) => Some(None),
+ None => None,
+ }
+ }
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::Array;
+ use arrow_array::Float64Array;
+ use std::sync::Arc;
+
+ fn float_array(values: Vec<Option<f64>>) -> Arc<dyn Array> {
+ Arc::new(Float64Array::from(values)) as Arc<dyn Array>
+ }
+
+ #[test]
+ fn daffine2_iterator_handles_nulls() {
+ let args = vec![
+ float_array(vec![Some(1.0), Some(10.0)]),
+ float_array(vec![Some(2.0), Some(20.0)]),
+ float_array(vec![Some(3.0), Some(30.0)]),
+ float_array(vec![Some(4.0), None]),
+ float_array(vec![Some(5.0), Some(50.0)]),
+ float_array(vec![Some(6.0), Some(60.0)]),
+ ];
+
+ let mut iter = DAffine2Iterator::new(&args).unwrap();
+
+ let expected_first = glam::DAffine2 {
+ matrix2: glam::DMat2 {
+ x_axis: glam::DVec2 { x: 1.0, y: 2.0 },
+ y_axis: glam::DVec2 { x: 3.0, y: 4.0 },
+ },
+ translation: glam::DVec2 { x: 5.0, y: 6.0 },
+ };
+ assert_eq!(iter.next(), Some(Some(expected_first)));
+
+ // The second case contains NULL, so the result is NULL
+ assert_eq!(iter.next(), Some(None));
+ }
+
+ #[test]
+ fn daffine3_iterator_values() {
+ let values = [
+ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
+ ];
+ let args = values
+ .iter()
+ .map(|value| float_array(vec![Some(*value)]))
+ .collect::<Vec<_>>();
+
+ let mut iter = DAffine3Iterator::new(&args).unwrap();
+ let expected = glam::DAffine3 {
+ matrix3: glam::DMat3::from_cols(
+ glam::DVec3::new(1.0, 2.0, 3.0),
+ glam::DVec3::new(4.0, 5.0, 6.0),
+ glam::DVec3::new(7.0, 8.0, 9.0),
+ ),
+ translation: glam::DVec3::new(10.0, 11.0, 12.0),
+ };
+
+ assert_eq!(iter.next(), Some(Some(expected)));
+ }
+
+ #[test]
+ fn daffine_iterator_from_scale() {
+ let scale_args = vec![
+ float_array(vec![Some(2.0), None]),
+ float_array(vec![Some(3.0), Some(4.0)]),
+ ];
+ let mut iter = DAffineIterator::from_scale_2d(&scale_args).unwrap();
+
+ let expected_scale =
+ DAffine::DAffine2(glam::DAffine2::from_scale(glam::DVec2::new(2.0,
3.0)));
+ assert_eq!(iter.next(), Some(Some(expected_scale)));
+
+ // The second case contains NULL, so the result is NULL
+ assert_eq!(iter.next(), Some(None));
+ }
+
+ #[test]
+ fn daffine_iterator_from_rotate() {
+ let angle = float_array(vec![Some(0.25), None]);
+ let mut iter = DAffineIterator::from_angle(&angle,
RotateAxis::X).unwrap();
+ let expected_rotate =
DAffine::DAffine3(glam::DAffine3::from_rotation_x(0.25));
+ assert_eq!(iter.next(), Some(Some(expected_rotate)));
+
+ // The second case contains NULL, so the result is NULL
+ assert_eq!(iter.next(), Some(None));
+ }
+
+ #[test]
+ fn daffine_crs_transform_changes_coords() {
+ let mut coord_2d = (1.0, 2.0);
+ let affine_2d =
DAffine::DAffine2(glam::DAffine2::from_scale(glam::DVec2::new(2.0, 3.0)));
+ affine_2d.transform_coord(&mut coord_2d).unwrap();
+ assert_eq!(coord_2d, (2.0, 6.0));
+
+ 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();
+ assert_eq!(coord_3d, (2.0, 6.0, 12.0));
+ }
+}
diff --git a/rust/sedona-functions/src/st_rotate.rs
b/rust/sedona-functions/src/st_rotate.rs
new file mode 100644
index 00000000..c60781fc
--- /dev/null
+++ b/rust/sedona-functions/src/st_rotate.rs
@@ -0,0 +1,261 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+use arrow_array::builder::BinaryBuilder;
+use arrow_schema::DataType;
+use datafusion_common::{error::Result, DataFusionError};
+use datafusion_expr::{
+ scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation,
Volatility,
+};
+use sedona_expr::{
+ item_crs::ItemCrsKernel,
+ scalar_udf::{SedonaScalarKernel, SedonaScalarUDF},
+};
+use sedona_geometry::{transform::transform,
wkb_factory::WKB_MIN_PROBABLE_BYTES};
+use sedona_schema::{
+ datatypes::{SedonaType, WKB_GEOMETRY},
+ matchers::ArgMatcher,
+};
+use std::sync::Arc;
+
+use crate::{
+ executor::WkbExecutor,
+ st_affine_helpers::{self, RotateAxis},
+};
+
+/// ST_Rotate() scalar UDF
+///
+/// Native implementation for rotate transformation
+pub fn st_rotate_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "st_rotate",
+ ItemCrsKernel::wrap_impl(vec![Arc::new(STRotate {
+ axis: RotateAxis::Z,
+ })]),
+ Volatility::Immutable,
+ Some(st_rotate_doc("")),
+ )
+}
+
+/// ST_RotateX() scalar UDF
+///
+/// Native implementation for rotate transformation
+pub fn st_rotate_x_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "st_rotatex",
+ ItemCrsKernel::wrap_impl(vec![Arc::new(STRotate {
+ axis: RotateAxis::X,
+ })]),
+ Volatility::Immutable,
+ Some(st_rotate_doc("X")),
+ )
+}
+
+/// ST_RotateY() scalar UDF
+///
+/// Native implementation for rotate transformation
+pub fn st_rotate_y_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "st_rotatey",
+ ItemCrsKernel::wrap_impl(vec![Arc::new(STRotate {
+ axis: RotateAxis::Y,
+ })]),
+ Volatility::Immutable,
+ Some(st_rotate_doc("Y")),
+ )
+}
+
+fn st_rotate_doc(axis: &str) -> Documentation {
+ let suffix = match axis {
+ "Z" => "",
+ _ => axis,
+ };
+ Documentation::builder(
+ DOC_SECTION_OTHER,
+ format!("Rotates a geometry by a specified angle in radians
counter-clockwise around the {axis}-axis "),
+ format!("ST_Rotate{suffix} (geom: Geometry, rot: Double)"),
+ )
+ .with_argument("geom", "geometry: Input geometry")
+ .with_argument("rot", "angle (in radians)")
+ .with_sql_example(
+ format!("SELECT ST_Rotate{suffix}(ST_GeomFromText('POLYGON Z ((1 0 1,
1 1 1, 2 2 2, 1 0 1))'), radians(45))"),
+ )
+ .build()
+}
+
+#[derive(Debug)]
+struct STRotate {
+ axis: RotateAxis,
+}
+
+impl SedonaScalarKernel for STRotate {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let matcher = ArgMatcher::new(
+ vec![ArgMatcher::is_geometry(), ArgMatcher::is_numeric()],
+ WKB_GEOMETRY,
+ );
+
+ matcher.match_args(args)
+ }
+
+ fn invoke_batch(
+ &self,
+ arg_types: &[SedonaType],
+ args: &[ColumnarValue],
+ ) -> Result<ColumnarValue> {
+ let executor = WkbExecutor::new(arg_types, args);
+ let mut builder = BinaryBuilder::with_capacity(
+ executor.num_iterations(),
+ WKB_MIN_PROBABLE_BYTES * executor.num_iterations(),
+ );
+
+ let angle = args[1]
+ .cast_to(&DataType::Float64, None)?
+ .to_array(executor.num_iterations())?;
+
+ let mut affine_iter =
st_affine_helpers::DAffineIterator::from_angle(&angle, self.axis)?;
+
+ executor.execute_wkb_void(|maybe_wkb| {
+ let maybe_mat = affine_iter.next().unwrap();
+ match (maybe_wkb, maybe_mat) {
+ (Some(wkb), Some(mat)) => {
+ transform(&wkb, &mat, &mut builder)
+ .map_err(|e|
DataFusionError::Execution(e.to_string()))?;
+ builder.append_value([]);
+ }
+ _ => builder.append_null(),
+ }
+
+ Ok(())
+ })?;
+
+ executor.finish(Arc::new(builder.finish()))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::f64;
+
+ use arrow_array::Array;
+ use datafusion_expr::ScalarUDF;
+ use rstest::rstest;
+ use sedona_schema::datatypes::{WKB_GEOMETRY_ITEM_CRS, WKB_VIEW_GEOMETRY};
+ use sedona_testing::{
+ compare::assert_array_equal, create::create_array,
testers::ScalarUdfTester,
+ };
+
+ use super::*;
+
+ #[test]
+ fn udf_metadata() {
+ let st_rotate_udf: ScalarUDF = st_rotate_udf().into();
+ assert_eq!(st_rotate_udf.name(), "st_rotate");
+ assert!(st_rotate_udf.documentation().is_some());
+
+ let st_rotate_x_udf: ScalarUDF = st_rotate_x_udf().into();
+ assert_eq!(st_rotate_x_udf.name(), "st_rotatex");
+ assert!(st_rotate_x_udf.documentation().is_some());
+
+ let st_rotate_y_udf: ScalarUDF = st_rotate_y_udf().into();
+ assert_eq!(st_rotate_y_udf.name(), "st_rotatey");
+ assert!(st_rotate_y_udf.documentation().is_some());
+ }
+
+ #[rstest]
+ fn udf(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType)
{
+ let tester_z = ScalarUdfTester::new(
+ st_rotate_udf().into(),
+ vec![sedona_type.clone(), SedonaType::Arrow(DataType::Float64)],
+ );
+ let tester_x = ScalarUdfTester::new(
+ st_rotate_x_udf().into(),
+ vec![sedona_type.clone(), SedonaType::Arrow(DataType::Float64)],
+ );
+ let tester_y = ScalarUdfTester::new(
+ st_rotate_y_udf().into(),
+ vec![sedona_type.clone(), SedonaType::Arrow(DataType::Float64)],
+ );
+ tester_z.assert_return_type(WKB_GEOMETRY);
+ tester_x.assert_return_type(WKB_GEOMETRY);
+ tester_y.assert_return_type(WKB_GEOMETRY);
+
+ let points = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (1 2)"),
+ Some("POINT M (1 2 3)"),
+ Some("POINT Z (1 2 3)"),
+ ],
+ &sedona_type,
+ );
+ let expected_identity = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (1 2)"),
+ Some("POINT M (1 2 3)"),
+ Some("POINT Z (1 2 3)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_identity_z = tester_z
+ .invoke_arrays(prepare_args(points.clone(), &[Some(0.0_f64)]))
+ .unwrap();
+ assert_array_equal(&result_identity_z, &expected_identity);
+
+ let result_identity_x = tester_x
+ .invoke_arrays(prepare_args(points.clone(), &[Some(0.0_f64)]))
+ .unwrap();
+ assert_array_equal(&result_identity_x, &expected_identity);
+
+ let result_identity_y = tester_y
+ .invoke_arrays(prepare_args(points.clone(), &[Some(0.0_f64)]))
+ .unwrap();
+ assert_array_equal(&result_identity_y, &expected_identity);
+
+ // Don't test the rotated results here since it's hard to match with
the exact number.
+ }
+
+ fn prepare_args(wkt: Arc<dyn Array>, mat: &[Option<f64>]) -> Vec<Arc<dyn
Array>> {
+ let n = wkt.len();
+ let mut args: Vec<Arc<dyn Array>> = mat
+ .iter()
+ .map(|a| {
+ let values = vec![*a; n];
+ Arc::new(arrow_array::Float64Array::from(values)) as Arc<dyn
Array>
+ })
+ .collect();
+ args.insert(0, wkt);
+ args
+ }
+
+ #[rstest]
+ fn udf_invoke_item_crs(#[values(WKB_GEOMETRY_ITEM_CRS.clone())]
sedona_type: SedonaType) {
+ let tester = ScalarUdfTester::new(
+ st_rotate_udf().into(),
+ vec![sedona_type.clone(), SedonaType::Arrow(DataType::Float64)],
+ );
+ tester.assert_return_type(sedona_type.clone());
+
+ let result = tester.invoke_scalar_scalar("POINT (1 2)", 0.0).unwrap();
+ tester.assert_scalar_result_equals(result, "POINT (1 2)");
+ }
+}
diff --git a/rust/sedona-functions/src/st_scale.rs
b/rust/sedona-functions/src/st_scale.rs
new file mode 100644
index 00000000..419e0810
--- /dev/null
+++ b/rust/sedona-functions/src/st_scale.rs
@@ -0,0 +1,348 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+use arrow_array::{builder::BinaryBuilder, Array};
+use arrow_schema::DataType;
+use datafusion_common::{error::Result, DataFusionError};
+use datafusion_expr::{
+ scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation,
Volatility,
+};
+use sedona_expr::{
+ item_crs::ItemCrsKernel,
+ scalar_udf::{SedonaScalarKernel, SedonaScalarUDF},
+};
+use sedona_geometry::{transform::transform,
wkb_factory::WKB_MIN_PROBABLE_BYTES};
+use sedona_schema::{
+ datatypes::{SedonaType, WKB_GEOMETRY},
+ matchers::ArgMatcher,
+};
+use std::sync::Arc;
+
+use crate::{
+ executor::WkbExecutor,
+ st_affine_helpers::{self},
+};
+
+/// ST_Scale() scalar UDF
+///
+/// Native implementation for scale transformation
+pub fn st_scale_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "st_scale",
+ ItemCrsKernel::wrap_impl(vec![
+ Arc::new(STScale { is_3d: true }),
+ Arc::new(STScale { is_3d: false }),
+ ]),
+ Volatility::Immutable,
+ Some(st_scale_doc()),
+ )
+}
+
+fn st_scale_doc() -> Documentation {
+ Documentation::builder(
+ DOC_SECTION_OTHER,
+ "Scales the geometry to a new size by multiplying the ordinates with
the corresponding scaling factors",
+ "ST_Scale (geom: Geometry, scaleX: Double, scaleY: Double, scaleZ:
Double)",
+ )
+ .with_argument("geom", "geometry: Input geometry")
+ .with_argument("scaleX", "scaling factor for X")
+ .with_argument("scaleY", "scaling factor for Y")
+ .with_argument("scaleZ", "scaling factor for Z")
+ .with_sql_example(
+ "SELECT ST_Scale(ST_GeomFromText('POLYGON Z ((1 0 1, 1 1 1, 2 2 2, 1 0
1))'), 1, 2, 3)",
+ )
+ .build()
+}
+
+#[derive(Debug)]
+struct STScale {
+ is_3d: bool,
+}
+
+impl SedonaScalarKernel for STScale {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let arg_matchers = if self.is_3d {
+ vec![
+ ArgMatcher::is_geometry(),
+ ArgMatcher::is_numeric(),
+ ArgMatcher::is_numeric(),
+ ArgMatcher::is_numeric(),
+ ]
+ } else {
+ vec![
+ ArgMatcher::is_geometry(),
+ ArgMatcher::is_numeric(),
+ ArgMatcher::is_numeric(),
+ ]
+ };
+
+ let matcher = ArgMatcher::new(arg_matchers, WKB_GEOMETRY);
+
+ matcher.match_args(args)
+ }
+
+ fn invoke_batch(
+ &self,
+ arg_types: &[SedonaType],
+ args: &[ColumnarValue],
+ ) -> Result<ColumnarValue> {
+ let executor = WkbExecutor::new(arg_types, args);
+ let mut builder = BinaryBuilder::with_capacity(
+ executor.num_iterations(),
+ WKB_MIN_PROBABLE_BYTES * executor.num_iterations(),
+ );
+
+ let array_args = args[1..]
+ .iter()
+ .map(|arg| {
+ arg.cast_to(&DataType::Float64, None)?
+ .to_array(executor.num_iterations())
+ })
+ .collect::<Result<Vec<Arc<dyn Array>>>>()?;
+
+ let mut affine_iter = if self.is_3d {
+ st_affine_helpers::DAffineIterator::from_scale_3d(&array_args)?
+ } else {
+ st_affine_helpers::DAffineIterator::from_scale_2d(&array_args)?
+ };
+
+ executor.execute_wkb_void(|maybe_wkb| {
+ let maybe_mat = affine_iter.next().unwrap();
+ match (maybe_wkb, maybe_mat) {
+ (Some(wkb), Some(mat)) => {
+ transform(&wkb, &mat, &mut builder)
+ .map_err(|e|
DataFusionError::Execution(e.to_string()))?;
+ builder.append_value([]);
+ }
+ _ => builder.append_null(),
+ }
+
+ Ok(())
+ })?;
+
+ executor.finish(Arc::new(builder.finish()))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use arrow_array::Array;
+ use datafusion_expr::ScalarUDF;
+ use rstest::rstest;
+ use sedona_schema::datatypes::{WKB_GEOMETRY_ITEM_CRS, WKB_VIEW_GEOMETRY};
+ use sedona_testing::{
+ compare::assert_array_equal, create::create_array,
testers::ScalarUdfTester,
+ };
+
+ use super::*;
+
+ #[test]
+ fn udf_metadata() {
+ let st_scale_udf: ScalarUDF = st_scale_udf().into();
+ assert_eq!(st_scale_udf.name(), "st_scale");
+ assert!(st_scale_udf.documentation().is_some());
+ }
+
+ #[rstest]
+ fn udf(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType)
{
+ let tester_2d = ScalarUdfTester::new(
+ st_scale_udf().into(),
+ vec![
+ sedona_type.clone(),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ ],
+ );
+ tester_2d.assert_return_type(WKB_GEOMETRY);
+
+ let points_2d = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (1 2)"),
+ Some("POINT M (1 2 3)"),
+ ],
+ &sedona_type,
+ );
+
+ let expected_identity_2d = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (1 2)"),
+ Some("POINT M (1 2 3)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_identity_2d = tester_2d
+ .invoke_arrays(prepare_args(points_2d.clone(), &[Some(1.0),
Some(1.0)]))
+ .unwrap();
+ assert_array_equal(&result_identity_2d, &expected_identity_2d);
+
+ let expected_scale_2d = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (10 40)"),
+ Some("POINT M (10 40 3)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_scale_2d = tester_2d
+ .invoke_arrays(prepare_args(points_2d, &[Some(10.0), Some(20.0)]))
+ .unwrap();
+ assert_array_equal(&result_scale_2d, &expected_scale_2d);
+
+ let points_3d = create_array(
+ &[
+ None,
+ Some("POINT Z EMPTY"),
+ Some("POINT ZM EMPTY"),
+ Some("POINT Z (1 2 3)"),
+ Some("POINT ZM (1 2 3 4)"),
+ ],
+ &sedona_type,
+ );
+
+ let expected_scale_2d_on_3d = create_array(
+ &[
+ None,
+ Some("POINT Z EMPTY"),
+ Some("POINT ZM EMPTY"),
+ Some("POINT Z (2 6 3)"),
+ Some("POINT ZM (2 6 3 4)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_scale_2d_on_3d = tester_2d
+ .invoke_arrays(prepare_args(points_3d.clone(), &[Some(2.0),
Some(3.0)]))
+ .unwrap();
+ assert_array_equal(&result_scale_2d_on_3d, &expected_scale_2d_on_3d);
+
+ let tester_3d = ScalarUdfTester::new(
+ st_scale_udf().into(),
+ vec![
+ sedona_type.clone(),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ ],
+ );
+ tester_3d.assert_return_type(WKB_GEOMETRY);
+
+ let expected_identity_3d = create_array(
+ &[
+ None,
+ Some("POINT Z EMPTY"),
+ Some("POINT ZM EMPTY"),
+ Some("POINT Z (1 2 3)"),
+ Some("POINT ZM (1 2 3 4)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_identity_3d = tester_3d
+ .invoke_arrays(prepare_args(
+ points_3d.clone(),
+ &[Some(1.0), Some(1.0), Some(1.0)],
+ ))
+ .unwrap();
+ assert_array_equal(&result_identity_3d, &expected_identity_3d);
+
+ let expected_scale_3d = create_array(
+ &[
+ None,
+ Some("POINT Z EMPTY"),
+ Some("POINT ZM EMPTY"),
+ Some("POINT Z (2 6 12)"),
+ Some("POINT ZM (2 6 12 4)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_scale_3d = tester_3d
+ .invoke_arrays(prepare_args(points_3d, &[Some(2.0), Some(3.0),
Some(4.0)]))
+ .unwrap();
+ assert_array_equal(&result_scale_3d, &expected_scale_3d);
+
+ let points_2d_for_3d = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (1 2)"),
+ Some("POINT M (1 2 3)"),
+ ],
+ &sedona_type,
+ );
+
+ let expected_scale_3d_on_2d = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT M EMPTY"),
+ Some("POINT (2 6)"),
+ Some("POINT M (2 6 3)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result_scale_3d_on_2d = tester_3d
+ .invoke_arrays(prepare_args(
+ points_2d_for_3d,
+ &[Some(2.0), Some(3.0), Some(4.0)],
+ ))
+ .unwrap();
+ assert_array_equal(&result_scale_3d_on_2d, &expected_scale_3d_on_2d);
+ }
+
+ fn prepare_args(wkt: Arc<dyn Array>, mat: &[Option<f64>]) -> Vec<Arc<dyn
Array>> {
+ let n = wkt.len();
+ let mut args: Vec<Arc<dyn Array>> = mat
+ .iter()
+ .map(|a| {
+ let values = vec![*a; n];
+ Arc::new(arrow_array::Float64Array::from(values)) as Arc<dyn
Array>
+ })
+ .collect();
+ args.insert(0, wkt);
+ args
+ }
+
+ #[rstest]
+ fn udf_invoke_item_crs(#[values(WKB_GEOMETRY_ITEM_CRS.clone())]
sedona_type: SedonaType) {
+ let tester = ScalarUdfTester::new(
+ st_scale_udf().into(),
+ vec![
+ sedona_type.clone(),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ ],
+ );
+ tester.assert_return_type(sedona_type.clone());
+
+ let result = tester
+ .invoke_scalar_scalar_scalar("POINT (1 2)", 2, 3)
+ .unwrap();
+ tester.assert_scalar_result_equals(result, "POINT (2 6)");
+ }
+}
diff --git a/rust/sedona-geometry/src/transform.rs
b/rust/sedona-geometry/src/transform.rs
index 409bc159..9d115a66 100644
--- a/rust/sedona-geometry/src/transform.rs
+++ b/rust/sedona-geometry/src/transform.rs
@@ -58,6 +58,16 @@ pub trait CrsEngine: Debug {
/// Trait for transforming coordinates in a geometry from one CRS to another.
pub trait CrsTransform: std::fmt::Debug {
fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(),
SedonaGeometryError>;
+
+ // CrsTransform can optionally handle 3D coordinates. If this method is
not implemented,
+ // the Z coordinate is simply ignored.
+ fn transform_coord_3d(&self, coord: &mut (f64, f64, f64)) -> Result<(),
SedonaGeometryError> {
+ let mut coord_2d = (coord.0, coord.1);
+ self.transform_coord(&mut coord_2d)?;
+ coord.0 = coord_2d.0;
+ coord.1 = coord_2d.1;
+ Ok(())
+ }
}
/// A boxed trait object for dynamic dispatch of CRS transformations.
@@ -65,6 +75,10 @@ impl CrsTransform for Box<dyn CrsTransform> {
fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(),
SedonaGeometryError> {
self.as_ref().transform_coord(coord)
}
+
+ fn transform_coord_3d(&self, coord: &mut (f64, f64, f64)) -> Result<(),
SedonaGeometryError> {
+ self.as_ref().transform_coord_3d(coord)
+ }
}
/// A caching wrapper around any CRS transformation engine.
@@ -345,24 +359,26 @@ where
I: Iterator<Item = C>,
{
for coord in coords {
- let mut xy: (f64, f64) = (coord.x(), coord.y());
- trans.transform_coord(&mut xy)?;
-
match coord.dim() {
Dimensions::Xy => {
+ let mut xy: (f64, f64) = (coord.x(), coord.y());
+ trans.transform_coord(&mut xy)?;
write_wkb_coord(buf, (xy.0, xy.1))?;
}
Dimensions::Xyz => {
- write_wkb_coord(buf, (xy.0, xy.1, coord.nth_or_panic(2)))?;
+ let mut xyz: (f64, f64, f64) = (coord.x(), coord.y(),
coord.nth_or_panic(2));
+ trans.transform_coord_3d(&mut xyz)?;
+ write_wkb_coord(buf, (xyz.0, xyz.1, xyz.2))?;
}
Dimensions::Xym => {
+ let mut xy: (f64, f64) = (coord.x(), coord.y());
+ trans.transform_coord(&mut xy)?;
write_wkb_coord(buf, (xy.0, xy.1, coord.nth_or_panic(2)))?;
}
Dimensions::Xyzm => {
- write_wkb_coord(
- buf,
- (xy.0, xy.1, coord.nth_or_panic(2), coord.nth_or_panic(3)),
- )?;
+ let mut xyz: (f64, f64, f64) = (coord.x(), coord.y(),
coord.nth_or_panic(2));
+ trans.transform_coord_3d(&mut xyz)?;
+ write_wkb_coord(buf, (xyz.0, xyz.1, xyz.2,
coord.nth_or_panic(3)))?;
}
_ => {
return Err(SedonaGeometryError::Invalid(
@@ -391,6 +407,49 @@ mod test {
}
}
+ #[derive(Debug)]
+ struct Mock3DTransform {}
+ impl CrsTransform for Mock3DTransform {
+ // This transforms 2D and 3D differently for testing purposes
+ fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(),
SedonaGeometryError> {
+ coord.0 += 100.0;
+ coord.1 += 200.0;
+ Ok(())
+ }
+
+ fn transform_coord_3d(
+ &self,
+ coord: &mut (f64, f64, f64),
+ ) -> Result<(), SedonaGeometryError> {
+ coord.0 += 10.0;
+ coord.1 += 20.0;
+ coord.2 += 30.0;
+ Ok(())
+ }
+ }
+
+ fn test_transform_inner(
+ geom: impl GeometryTrait<T = f64>,
+ expected: &str,
+ mock_transform: impl CrsTransform,
+ ) {
+ let mut wkb_bytes = Vec::new();
+
+ transform(geom, &mock_transform, &mut wkb_bytes).unwrap();
+ let wkb_reader = read_wkb(&wkb_bytes).unwrap();
+ let mut wkt = String::new();
+ wkt::to_wkt::write_geometry(&mut wkt, &wkb_reader).unwrap();
+ assert_eq!(wkt, expected);
+ }
+
+ fn test_transform(geom: impl GeometryTrait<T = f64>, expected: &str) {
+ test_transform_inner(geom, expected, MockTransform {})
+ }
+
+ fn test_transform_3d(geom: impl GeometryTrait<T = f64>, expected: &str) {
+ test_transform_inner(geom, expected, Mock3DTransform {})
+ }
+
#[test]
fn test_transform_point() {
let point = geo_types::Point::new(1.0, 2.0);
@@ -525,15 +584,32 @@ mod test {
test_transform(ls_xyzm, "LINESTRING ZM(11 22 3 4,15 26 7 8)");
}
- fn test_transform(geom: impl GeometryTrait<T = f64>, expected: &str) {
- let mock_transform = MockTransform {};
- let mut wkb_bytes = Vec::new();
+ #[test]
+ fn test_transform_point_3d() {
+ let point = wkt::Wkt::from_str("POINT Z(1 2 3)").unwrap();
+ test_transform_3d(point, "POINT Z(11 22 33)");
- transform(geom, &mock_transform, &mut wkb_bytes).unwrap();
- let wkb_reader = read_wkb(&wkb_bytes).unwrap();
- let mut wkt = String::new();
- wkt::to_wkt::write_geometry(&mut wkt, &wkb_reader).unwrap();
- assert_eq!(wkt, expected);
+ let nan_point = wkt::Wkt::from_str("POINT Z EMPTY").unwrap();
+ test_transform_3d(nan_point, "POINT Z EMPTY");
+ }
+
+ #[test]
+ fn test_transform_dimensions_3d() {
+ let ls_xy_wkt = "LINESTRING(1.0 2.0, 3.0 4.0)";
+ let ls_xy: Wkt = Wkt::from_str(ls_xy_wkt).unwrap();
+ test_transform_3d(ls_xy, "LINESTRING(101 202,103 204)");
+
+ let ls_xyz_wkt = "LINESTRING Z(1.0 2.0 3.0, 4.0 5.0 6.0)";
+ let ls_xyz: Wkt = Wkt::from_str(ls_xyz_wkt).unwrap();
+ test_transform_3d(ls_xyz, "LINESTRING Z(11 22 33,14 25 36)");
+
+ let ls_xym_wkt = "LINESTRING M(1.0 2.0 3.0, 4.0 5.0 6.0)";
+ let ls_xym: Wkt = Wkt::from_str(ls_xym_wkt).unwrap();
+ test_transform_3d(ls_xym, "LINESTRING M(101 202 3,104 205 6)");
+
+ let ls_xyzm_wkt = "LINESTRING ZM(1.0 2.0 3.0 4.0, 5.0 6.0 7.0 8.0)";
+ let ls_xyzm: Wkt = Wkt::from_str(ls_xyzm_wkt).unwrap();
+ test_transform_3d(ls_xyzm, "LINESTRING ZM(11 22 33 4,15 26 37 8)");
}
/// Mock CRS engine for testing caching behavior
