paleolimbot commented on code in PR #504:
URL: https://github.com/apache/sedona-db/pull/504#discussion_r2688381442
##########
rust/sedona-geometry/src/transform.rs:
##########
@@ -58,13 +58,27 @@ 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.
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)
+ }
Review Comment:
Thank you for this!
To ensure coverage in this file I think this could use a quick test of the
non-default 3D transform case:
```rust
#[derive(Debug)]
struct Mock3DTransform {}
impl CrsTransform for Mock3DTransform {
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(())
}
}
#[test]
fn test_transform_dimensions_3d() {
// like test_transform_dimensions but when there's a non-default 3D
transform
}
```
##########
rust/sedona-functions/src/st_rotate.rs:
##########
@@ -0,0 +1,241 @@
+// 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 arrow_array::Array;
+ 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_rotate_udf: ScalarUDF = st_rotate_udf().into();
+ assert_eq!(st_rotate_udf.name(), "st_rotate");
+ assert!(st_rotate_udf.documentation().is_some());
+ }
+
+ #[rstest]
+ fn udf(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType)
{
+ let tester = ScalarUdfTester::new(
+ st_rotate_udf().into(),
+ vec![sedona_type.clone(), SedonaType::Arrow(DataType::Float64)],
+ );
+ tester.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,
+ );
+ 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
+ .invoke_arrays(prepare_args(points.clone(),
&[Some(0.0_f64.to_radians())]))
+ .unwrap();
+ assert_array_equal(&result_identity, &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 geom = create_scalar(Some("POINT (1 2)"), &sedona_type);
+ let args = vec![
+ ColumnarValue::Scalar(geom),
+ ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))),
+ ];
+
+ let result = tester.invoke(args).unwrap();
+ if let ColumnarValue::Scalar(scalar) = result {
+ tester.assert_scalar_result_equals(scalar, "POINT (1 2)");
+ } else {
+ panic!("Expected scalar result from item CRS rotate invoke");
+ }
Review Comment:
```suggestion
let result = tester.invoke_scalar_scalar("POINT (1 2)", 0.0);
tester.assert_scalar_result_equals("POINT (1 2)");
```
##########
rust/sedona-geometry/src/transform.rs:
##########
@@ -58,13 +58,27 @@ 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;
Review Comment:
I think this is OK...copying small arguments is usually not faster and there
are possibly lower hanging fruit to optimize (optimized transform of an array
of coordinates, or possibly just non-aligned bytes -> bytes).
##########
rust/sedona-functions/src/st_scale.rs:
##########
@@ -0,0 +1,359 @@
+// 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_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_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 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(3.0))),
+ ];
+
+ let result = tester.invoke(args).unwrap();
+ if let ColumnarValue::Scalar(scalar) = result {
+ tester.assert_scalar_result_equals(scalar, "POINT (2 6)");
+ } else {
+ panic!("Expected scalar result from item CRS scale invoke");
+ }
Review Comment:
```suggestion
let result = tester.invoke_scalar_scalar_scalar("POINT (1 2)", 2,
3).unwrap();
tester.assert_scalar_result_equals("POINT (2 6)");
```
##########
rust/sedona-functions/src/st_rotate.rs:
##########
@@ -0,0 +1,241 @@
+// 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 arrow_array::Array;
+ 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_rotate_udf: ScalarUDF = st_rotate_udf().into();
+ assert_eq!(st_rotate_udf.name(), "st_rotate");
+ assert!(st_rotate_udf.documentation().is_some());
+ }
+
+ #[rstest]
+ fn udf(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType)
{
+ let tester = ScalarUdfTester::new(
+ st_rotate_udf().into(),
+ vec![sedona_type.clone(), SedonaType::Arrow(DataType::Float64)],
+ );
Review Comment:
Maybe a quick check that `st_rotate_y()` and `st_rotate_x()` is also plugged
in? It can be very short (just enough to check for crossed wires).
##########
rust/sedona-functions/src/st_affine.rs:
##########
@@ -0,0 +1,450 @@
+// 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;
+use datafusion_expr::{
+ scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation,
Volatility,
+};
+use geo_traits::GeometryTrait;
+use sedona_expr::{
+ item_crs::ItemCrsKernel,
+ scalar_udf::{SedonaScalarKernel, SedonaScalarUDF},
+};
+use sedona_geometry::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};
+
+/// 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>>>>()?;
Review Comment:
My thinking was that all our existing functions (like ST_Translate()) have a
small number of arguments but 12 arguments might be enough that we'd notice the
copies. That's a very much future concern...implementing it like you did here I
think is the right thing to do.
##########
rust/sedona-functions/src/st_affine_helpers.rs:
##########
@@ -0,0 +1,531 @@
+// 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 datafusion_common::exec_err;
+use datafusion_common::internal_err;
+use datafusion_common::DataFusionError;
+use geo_traits::{
+ CoordTrait, GeometryCollectionTrait as _, GeometryTrait, LineStringTrait,
+ MultiLineStringTrait as _, MultiPointTrait as _, MultiPolygonTrait as _,
PointTrait,
+ PolygonTrait as _,
+};
+use sedona_geometry::wkb_factory::{
+ write_wkb_coord, write_wkb_empty_point,
write_wkb_geometrycollection_header,
+ write_wkb_linestring_header, write_wkb_multilinestring_header,
write_wkb_multipoint_header,
+ write_wkb_multipolygon_header, write_wkb_point_header,
write_wkb_polygon_header,
+ write_wkb_polygon_ring_header,
+};
+use std::io::Write;
+use std::sync::Arc;
+
+pub(crate) fn invoke_affine(
+ geom: &impl GeometryTrait<T = f64>,
+ writer: &mut impl Write,
+ mat: &DAffine,
+ dim: &geo_traits::Dimensions,
+) -> Result<()> {
+ let dims = geom.dim();
+ match geom.as_type() {
+ geo_traits::GeometryType::Point(pt) => {
+ if pt.coord().is_some() {
+ write_wkb_point_header(writer, dims)
+ .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+ write_transformed_coord(writer, pt.coord().unwrap(), mat,
dim)?;
+ } else {
+ write_wkb_empty_point(writer, dims)
+ .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+ }
+ }
+
+ geo_traits::GeometryType::MultiPoint(multi_point) => {
+ write_wkb_multipoint_header(writer, dims,
multi_point.points().count())
+ .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+ for pt in multi_point.points() {
+ invoke_affine(&pt, writer, mat, dim)?;
+ }
+ }
+
+ geo_traits::GeometryType::LineString(ls) => {
+ write_wkb_linestring_header(writer, dims, ls.coords().count())
+ .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+ write_transformed_coords(writer, ls.coords(), mat, dim)?;
+ }
+
+ geo_traits::GeometryType::Polygon(pgn) => {
+ let num_rings = pgn.interiors().count() + pgn.exterior().is_some()
as usize;
+ write_wkb_polygon_header(writer, dims, num_rings)
+ .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+
+ if let Some(exterior) = pgn.exterior() {
+ write_transformed_ring(writer, exterior, mat, dim)?;
+ }
+
+ for interior in pgn.interiors() {
+ write_transformed_ring(writer, interior, mat, dim)?;
+ }
+ }
+
+ geo_traits::GeometryType::MultiLineString(mls) => {
+ write_wkb_multilinestring_header(writer, dims,
mls.line_strings().count())
+ .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+ for ls in mls.line_strings() {
+ invoke_affine(&ls, writer, mat, dim)?;
+ }
+ }
+
+ geo_traits::GeometryType::MultiPolygon(mpgn) => {
+ write_wkb_multipolygon_header(writer, dims,
mpgn.polygons().count())
+ .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+ for pgn in mpgn.polygons() {
+ invoke_affine(&pgn, writer, mat, dim)?;
+ }
+ }
+
+ geo_traits::GeometryType::GeometryCollection(gcn) => {
+ write_wkb_geometrycollection_header(writer, dims,
gcn.geometries().count())
+ .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+ for geom in gcn.geometries() {
+ invoke_affine(&geom, writer, mat, dim)?;
+ }
+ }
+
+ _ => {
+ return Err(DataFusionError::Execution(
+ "Unsupported geometry type for reversal operation".to_string(),
+ ));
+ }
+ }
+ Ok(())
+}
+
+fn write_transformed_ring(
+ writer: &mut impl Write,
+ ring: impl LineStringTrait<T = f64>,
+ affine: &DAffine,
+ dim: &geo_traits::Dimensions,
+) -> Result<()> {
+ write_wkb_polygon_ring_header(writer, ring.coords().count())
+ .map_err(|e| DataFusionError::Execution(e.to_string()))?;
+ write_transformed_coords(writer, ring.coords(), affine, dim)
+}
+
+fn write_transformed_coords<I>(
+ writer: &mut impl Write,
+ coords: I,
+ affine: &DAffine,
+ dim: &geo_traits::Dimensions,
+) -> Result<()>
+where
+ I: DoubleEndedIterator,
+ I::Item: CoordTrait<T = f64>,
+{
+ coords
+ .into_iter()
+ .try_for_each(|coord| write_transformed_coord(writer, coord, affine,
dim))
+}
+
+fn write_transformed_coord<C>(
+ writer: &mut impl Write,
+ coord: C,
+ affine: &DAffine,
+ dim: &geo_traits::Dimensions,
+) -> Result<()>
+where
+ C: CoordTrait<T = f64>,
+{
+ match dim {
+ geo_traits::Dimensions::Xy => {
+ let transformed = affine.transform_point2(coord.x(), coord.y());
+ write_wkb_coord(writer, transformed)
+ .map_err(|e| DataFusionError::Execution(e.to_string()))
+ }
+ geo_traits::Dimensions::Xym => {
+ let transformed = affine.transform_point2(coord.x(), coord.y());
+ // Preserve m value
+ let m = coord.nth(2).unwrap();
+ write_wkb_coord(writer, (transformed.0, transformed.1, m))
+ .map_err(|e| DataFusionError::Execution(e.to_string()))
+ }
+ geo_traits::Dimensions::Xyz => {
+ let transformed = affine.transform_point3(coord.x(), coord.y(),
coord.nth(2).unwrap());
+ write_wkb_coord(writer, transformed)
+ .map_err(|e| DataFusionError::Execution(e.to_string()))
+ }
+ geo_traits::Dimensions::Xyzm => {
+ let transformed = affine.transform_point3(coord.x(), coord.y(),
coord.nth(2).unwrap());
+ // Preserve m value
+ let m = coord.nth(3).unwrap();
+ write_wkb_coord(writer, (transformed.0, transformed.1,
transformed.2, m))
+ .map_err(|e| DataFusionError::Execution(e.to_string()))
+ }
+ geo_traits::Dimensions::Unknown(_) => {
+ exec_err!("A geometry with unknown dimension cannot be
transformed")
+ }
+ }
+}
+
+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>,
+}
+
+impl<'a> DAffine2Iterator<'a> {
+ pub(crate) fn new(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ if array_args.len() != 6 {
+ return 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,
+ })
+ }
+}
+
+impl<'a> Iterator for DAffine2Iterator<'a> {
+ type Item = glam::DAffine2;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+ 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>,
+}
+
+impl<'a> DAffine3Iterator<'a> {
+ pub(crate) fn new(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ if array_args.len() != 12 {
+ return 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,
+ })
+ }
+}
+
+impl<'a> Iterator for DAffine3Iterator<'a> {
+ type Item = glam::DAffine3;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+ 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>,
+}
+
+impl<'a> DAffine2ScaleIterator<'a> {
+ pub(crate) fn new(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ if array_args.len() != 2 {
+ return 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,
+ })
+ }
+}
+
+impl<'a> Iterator for DAffine2ScaleIterator<'a> {
+ type Item = glam::DAffine2;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+ let scale = glam::DVec2::new(self.x_scale.value(i),
self.y_scale.value(i));
+ 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>,
+}
+
+impl<'a> DAffine3ScaleIterator<'a> {
+ pub(crate) fn new(array_args: &'a [Arc<dyn Array>]) -> Result<Self> {
+ if array_args.len() != 3 {
+ return 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,
+ })
+ }
+}
+
+impl<'a> Iterator for DAffine3ScaleIterator<'a> {
+ type Item = glam::DAffine3;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+ let scale = glam::DVec3::new(
+ self.x_scale.value(i),
+ self.y_scale.value(i),
+ self.z_scale.value(i),
+ );
+ 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,
+}
+
+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,
+ })
+ }
+}
+
+impl<'a> Iterator for DAffineRotateIterator<'a> {
+ type Item = glam::DAffine3;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+ match self.axis {
+ RotateAxis::X =>
Some(glam::DAffine3::from_rotation_x(self.angle.value(i))),
+ RotateAxis::Y =>
Some(glam::DAffine3::from_rotation_y(self.angle.value(i))),
+ RotateAxis::Z =>
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,
+ )?))
+ }
+}
+
+pub(crate) enum DAffine {
+ DAffine2(glam::DAffine2),
+ DAffine3(glam::DAffine3),
+}
+
+impl DAffine {
+ pub(crate) fn transform_point3(&self, x: f64, y: f64, z: f64) -> (f64,
f64, f64) {
+ match self {
+ DAffine::DAffine2(daffine2) => {
+ let transformed = daffine2.transform_point2(glam::DVec2 { x, y
});
+ (transformed.x, transformed.y, z)
+ }
+ DAffine::DAffine3(daffine3) => {
+ let transformed = daffine3.transform_point3(glam::DVec3 { x,
y, z });
+ (transformed.x, transformed.y, transformed.z)
+ }
+ }
+ }
+
+ pub(crate) fn transform_point2(&self, x: f64, y: f64) -> (f64, f64) {
+ match self {
+ DAffine::DAffine2(daffine2) => {
+ let transformed = daffine2.transform_point2(glam::DVec2 { x, y
});
+ (transformed.x, transformed.y)
+ }
+ DAffine::DAffine3(daffine3) => {
+ let transformed = daffine3.transform_point3(glam::DVec3 { x,
y, z: 0.0 });
+ (transformed.x, transformed.y)
+ }
+ }
+ }
+}
+
+impl<'a> Iterator for DAffineIterator<'a> {
+ type Item = DAffine;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ match self {
+ DAffineIterator::DAffine2(daffine2_iterator) => {
+ daffine2_iterator.next().map(DAffine::DAffine2)
+ }
+ DAffineIterator::DAffine3(daffine3_iterator) => {
+ daffine3_iterator.next().map(DAffine::DAffine3)
+ }
+ DAffineIterator::DAffine2Scale(daffine2_scale_iterator) => {
+ daffine2_scale_iterator.next().map(DAffine::DAffine2)
+ }
+ DAffineIterator::DAffine3Scale(daffine3_scale_iterator) => {
+ daffine3_scale_iterator.next().map(DAffine::DAffine3)
+ }
+ DAffineIterator::DAffineRotate(daffine_rotate_iterator) => {
+ daffine_rotate_iterator.next().map(DAffine::DAffine3)
+ }
+ }
+ }
+}
Review Comment:
I think we can have fewer tests now that you've implemented the transform
(thanks!), but I think we do need *something* here (or more tests in the Rust
functions). (For example, you check every null by name individually and an LLM
may come through and obliterate one of them in a refactor without a test
failing that I can see).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
