paleolimbot commented on code in PR #504: URL: https://github.com/apache/sedona-db/pull/504#discussion_r2682777488
########## 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(), Review Comment: No need to do this here because I think we need a helper for it to be compact, but in this particular case we know the output will have the exact number of bytes as the input) ########## 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"); Review Comment: ```suggestion return sedona_internal_err!("Invalid number of arguments are passed"); ``` ########## 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"); Review Comment: ```suggestion return sedona_internal_err!("Invalid number of arguments are passed"); ``` ########## python/sedonadb/tests/functions/test_functions.py: ########## @@ -189,6 +189,177 @@ 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"), + [ + ( + "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"), + [ + ( + "POINT Z (1 2 3)", + 1.0, 0.0, 0.0, + 0.0, 1.0, 0.0, + 0.0, 0.0, 1.0, + 0.0, 0.0, 0.0, + "POINT Z (1 2 3)", + ), + ( + "POINT Z (1 2 3)", + 2.0, 0.0, 0.0, + 0.0, 2.0, 0.0, + 0.0, 0.0, 2.0, + 1.0, 3.0, 5.0, + "POINT Z (3 7 11)", + ), + ], +) +def test_st_affine_3d( + eng, geom, a, b, c, d, e, f, g, h, i, xoff, yoff, zoff, expected +): + eng = eng.create_or_skip() + query = ( + "SELECT ST_Affine(" + f"{geom_or_null(geom)}, " + f"{val_or_null(a)}, {val_or_null(b)}, {val_or_null(c)}, " + f"{val_or_null(d)}, {val_or_null(e)}, {val_or_null(f)}, " + f"{val_or_null(g)}, {val_or_null(h)}, {val_or_null(i)}, " + f"{val_or_null(xoff)}, {val_or_null(yoff)}, {val_or_null(zoff)})" + ) + try: + eng.assert_query_result(query, expected) + except Exception as exc: + if isinstance(eng, PostGIS): + pytest.skip(f"PostGIS may not support 3D ST_Affine: {exc}") + raise Review Comment: Is this an issue with the version of PostGIS in our docker compose setup? ########## 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: This module would benefit from tests (e.g., for null item handling since that is non-trivial across 12 input arguments). One benefit of using the existing `transform()` would be that you don't have to write tests for it, and a benefit of consolidating iteration over `n` numeric arguments into a single struct is that you don't have to test it four times. In particular, I'm not sure there is test coverage at a higher level for null propagation or the full matrix of geometry types by XY XYM XYZ XYZM including empties, which I think is a concern of the transformed writer here. ########## 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()))?; + } + } + Review Comment: If this can be written as a `CrsTransform` we might be able re-use some existing infrastructure: https://github.com/apache/sedona-db/blob/9bf02f474d475b61a444e561ec74df449413f46b/rust/sedona-functions/src/st_translate.rs#L123-L135 https://github.com/apache/sedona-db/blob/9bf02f474d475b61a444e561ec74df449413f46b/rust/sedona-functions/src/st_translate.rs#L106-L108 https://github.com/apache/sedona-db/blob/main/rust/sedona-geometry/src/transform.rs#L259-L263 There's a chance what you have here is faster, though (no `dyn` like our current `transform()`), and we haven't implemented non-XY support for the `CrsTransform` yet ( https://github.com/apache/sedona-db/issues/47 ). No need to take on any of that here (I'll just add a note to that issue that we can update this section when we do get there). ########## 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<()> Review Comment: (Optional) It may be unfounded, but my intuition would be that `match (dim)` should be in `write_transformed_coords()` to avoid a `match()` for every coordinate. Perhaps the compiler does this anyway or the transform operation is sufficiently expensive that it doesn't matter 🙂 ########## 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: No need to do this now, but a future optimization would be to handle the case where all 9/12 input arguments are scalars. In general I haven't been able to notice a speed improvement for the case where one or two scalar numerics have to get expanded into arrays but I'm not sure that's the case with 12 input arguments 🙂 ########## 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>, +} Review Comment: Do we need multiple versions of this (e.g., can the core logic be distilled to `struct NumericArgs { args: Vec<&'a PrimitiveArray<Float64Type>> }`? Or maybe the lifetimes make this difficult?). My reading of these is that they don't handle nulls in any of the arguments. I am not sure if it's optimizing this but you could check ahead of time if the `null_count()` is >0 for any of the inputs and use the slower iterator if this is the case. ########## python/sedonadb/tests/functions/test_functions.py: ########## @@ -189,6 +189,177 @@ 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"), + [ + ( + "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"), + [ + ( + "POINT Z (1 2 3)", + 1.0, 0.0, 0.0, + 0.0, 1.0, 0.0, + 0.0, 0.0, 1.0, + 0.0, 0.0, 0.0, + "POINT Z (1 2 3)", + ), + ( + "POINT Z (1 2 3)", + 2.0, 0.0, 0.0, + 0.0, 2.0, 0.0, + 0.0, 0.0, 2.0, + 1.0, 3.0, 5.0, + "POINT Z (3 7 11)", + ), + ], +) +def test_st_affine_3d( + eng, geom, a, b, c, d, e, f, g, h, i, xoff, yoff, zoff, expected +): + eng = eng.create_or_skip() + query = ( + "SELECT ST_Affine(" + f"{geom_or_null(geom)}, " + f"{val_or_null(a)}, {val_or_null(b)}, {val_or_null(c)}, " + f"{val_or_null(d)}, {val_or_null(e)}, {val_or_null(f)}, " + f"{val_or_null(g)}, {val_or_null(h)}, {val_or_null(i)}, " + f"{val_or_null(xoff)}, {val_or_null(yoff)}, {val_or_null(zoff)})" + ) + try: + eng.assert_query_result(query, expected) + except Exception as exc: + if isinstance(eng, PostGIS): + pytest.skip(f"PostGIS may not support 3D ST_Affine: {exc}") + raise +# fmt: on + + [email protected]("eng", [SedonaDB, PostGIS]) [email protected]( + ("geom", "sx", "sy", "expected"), + [ + (None, 2.0, 3.0, None), + ("POINT (1 2)", 1.0, 1.0, "POINT (1 2)"), + ("POINT (1 2)", 2.0, 3.0, "POINT (2 6)"), + ("LINESTRING (0 0, 1 1)", 2.0, 3.0, "LINESTRING (0 0, 2 3)"), + ( + "POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0))", + 2.0, + 3.0, + "POLYGON ((0 0, 2 0, 2 3, 0 3, 0 0))", + ), + ( + "MULTIPOINT (1 2, 3 4)", + 2.0, + 3.0, + "MULTIPOINT (2 6, 6 12)", + ), + ( + "MULTILINESTRING ((0 0, 1 1), (2 2, 3 3))", + 2.0, + 3.0, + "MULTILINESTRING ((0 0, 2 3), (4 6, 6 9))", + ), + ( + "MULTIPOLYGON (((0 0, 1 0, 1 1, 0 1, 0 0)))", + 2.0, + 3.0, + "MULTIPOLYGON (((0 0, 2 0, 2 3, 0 3, 0 0)))", + ), + ( + "GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (0 0, 1 1))", + 2.0, + 3.0, + "GEOMETRYCOLLECTION (POINT (2 6), LINESTRING (0 0, 2 3))", + ), + ("POINT Z (1 2 3)", 2.0, 3.0, "POINT Z (2 6 3)"), + ], +) +def test_st_scale_2d(eng, geom, sx, sy, expected): + eng = eng.create_or_skip() + eng.assert_query_result( + f"SELECT ST_Scale({geom_or_null(geom)}, {val_or_null(sx)}, {val_or_null(sy)})", + expected, + ) + + [email protected]("eng", [SedonaDB, PostGIS]) [email protected]( + ("geom", "sx", "sy", "sz", "expected"), + [ + ("POINT Z (1 2 3)", 1.0, 1.0, 1.0, "POINT Z (1 2 3)"), + ("POINT Z (1 2 3)", 2.0, 3.0, 4.0, "POINT Z (2 6 12)"), + ], +) +def test_st_scale_3d(eng, geom, sx, sy, sz, expected): + eng = eng.create_or_skip() + query = ( + "SELECT ST_Scale(" + f"{geom_or_null(geom)}, {val_or_null(sx)}, {val_or_null(sy)}, {val_or_null(sz)})" + ) + try: + eng.assert_query_result(query, expected) + except Exception as exc: + if isinstance(eng, PostGIS): + pytest.skip(f"PostGIS may not support 3D ST_Scale: {exc}") + raise + + [email protected]("eng", [SedonaDB, PostGIS]) [email protected]( + ("geom", "angle", "expected_x", "expected_y"), + [ + (None, 0.0, None, None), + ("POINT (1 2)", 0.0, 1.0, 2.0), + ("POINT (1 2)", 1.5707963267948966, -2.0, 1.0), + ("POINT (1 2)", 3.141592653589793, -1.0, -2.0), + ], +) +def test_st_rotate(eng, geom, angle, expected_x, expected_y): + eng = eng.create_or_skip() + x_query = f"SELECT ST_X(ST_Rotate({geom_or_null(geom)}, {val_or_null(angle)}))" + y_query = f"SELECT ST_Y(ST_Rotate({geom_or_null(geom)}, {val_or_null(angle)}))" + eng.assert_query_result(x_query, expected_x, numeric_epsilon=1e-12) + eng.assert_query_result(y_query, expected_y, numeric_epsilon=1e-12) Review Comment: There is also `wkt_precision` if you just want to check the result of `ST_Rotate()`: https://github.com/apache/sedona-db/blob/9bf02f474d475b61a444e561ec74df449413f46b/python/sedonadb/python/sedonadb/testing.py#L186 ########## 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"); Review Comment: ```suggestion return sedona_internal_err!("Invalid number of arguments are passed"); ``` -- 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]
