From 46724277a8245f517b471db0011f4e709e45df5f Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 00:19:15 +0900 Subject: [PATCH 01/15] Let CrsTransform decide the output dimension --- rust/sedona-geometry/src/transform.rs | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/rust/sedona-geometry/src/transform.rs b/rust/sedona-geometry/src/transform.rs index 9d115a666..af6c9bffb 100644 --- a/rust/sedona-geometry/src/transform.rs +++ b/rust/sedona-geometry/src/transform.rs @@ -68,6 +68,11 @@ pub trait CrsTransform: std::fmt::Debug { coord.1 = coord_2d.1; Ok(()) } + + // The dimension of the output. If `None`, the input dimension is preserved. + fn output_dim(&self) -> Option { + None + } } /// A boxed trait object for dynamic dispatch of CRS transformations. @@ -275,30 +280,32 @@ pub fn transform( trans: &dyn CrsTransform, out: &mut impl Write, ) -> Result<(), SedonaGeometryError> { - let dims = geom.dim(); + // If the CrsTransform specifies the dimension, use it. + // Otherwise, the input dimension is preserved. + let dims = trans.output_dim().unwrap_or_else(|| geom.dim()); match geom.as_type() { GeometryType::Point(pt) => { if pt.coord().is_some() { write_wkb_point_header(out, dims)?; - transform_and_write_coords(out, trans, pt.coord().into_iter())?; + transform_and_write_coords(out, trans, pt.coord().into_iter(), dims)?; } else { write_wkb_empty_point(out, dims)?; } } GeometryType::LineString(ls) => { - write_wkb_linestring_header(out, ls.dim(), ls.coords().count())?; - transform_and_write_coords(out, trans, ls.coords())?; + write_wkb_linestring_header(out, dims, ls.coords().count())?; + transform_and_write_coords(out, trans, ls.coords(), dims)?; } GeometryType::Polygon(pl) => { let num_rings = pl.interiors().count() + pl.exterior().is_some() as usize; - write_wkb_polygon_header(out, pl.dim(), num_rings)?; + write_wkb_polygon_header(out, dims, num_rings)?; if let Some(exterior) = pl.exterior() { - transform_and_write_ring(out, trans, exterior)?; + transform_and_write_ring(out, trans, exterior, dims)?; } for interior in pl.interiors() { - transform_and_write_ring(out, trans, interior)?; + transform_and_write_ring(out, trans, interior, dims)?; } } GeometryType::MultiPoint(multi_pt) => { @@ -339,13 +346,14 @@ fn transform_and_write_ring<'a, L>( buf: &mut impl Write, trans: &dyn CrsTransform, ring: L, + dims: Dimensions, ) -> Result<(), SedonaGeometryError> where L: LineStringTrait + 'a, { let num_points = ring.coords().count(); write_wkb_polygon_ring_header(buf, num_points)?; - transform_and_write_coords(buf, trans, ring.coords())?; + transform_and_write_coords(buf, trans, ring.coords(), dims)?; Ok(()) } @@ -353,13 +361,14 @@ fn transform_and_write_coords<'a, C, I>( buf: &mut impl Write, trans: &dyn CrsTransform, coords: I, + dims: Dimensions, ) -> Result<(), SedonaGeometryError> where C: CoordTrait + 'a, I: Iterator, { for coord in coords { - match coord.dim() { + match dims { Dimensions::Xy => { let mut xy: (f64, f64) = (coord.x(), coord.y()); trans.transform_coord(&mut xy)?; From b89f593e23bc572d48436502883f0a0616a0203c Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 00:19:37 +0900 Subject: [PATCH 02/15] Implement ST_Force2D --- rust/sedona-functions/src/lib.rs | 1 + rust/sedona-functions/src/register.rs | 1 + rust/sedona-functions/src/st_force_dim.rs | 128 ++++++++++++++++++++++ 3 files changed, 130 insertions(+) create mode 100644 rust/sedona-functions/src/st_force_dim.rs diff --git a/rust/sedona-functions/src/lib.rs b/rust/sedona-functions/src/lib.rs index e8037a06b..1640e7a90 100644 --- a/rust/sedona-functions/src/lib.rs +++ b/rust/sedona-functions/src/lib.rs @@ -42,6 +42,7 @@ mod st_dwithin; pub mod st_envelope; pub mod st_envelope_agg; pub mod st_flipcoordinates; +mod st_force_dim; mod st_geometryn; mod st_geometrytype; mod st_geomfromewkb; diff --git a/rust/sedona-functions/src/register.rs b/rust/sedona-functions/src/register.rs index 883f5a5ae..3feafd5fa 100644 --- a/rust/sedona-functions/src/register.rs +++ b/rust/sedona-functions/src/register.rs @@ -118,6 +118,7 @@ pub fn default_function_set() -> FunctionSet { crate::st_start_point::st_start_point_udf, crate::st_transform::st_transform_udf, crate::st_translate::st_translate_udf, + crate::st_force_dim::st_force2d_udf, crate::st_xyzm_minmax::st_mmax_udf, crate::st_xyzm_minmax::st_mmin_udf, crate::st_xyzm_minmax::st_xmax_udf, diff --git a/rust/sedona-functions/src/st_force_dim.rs b/rust/sedona-functions/src/st_force_dim.rs new file mode 100644 index 000000000..606796292 --- /dev/null +++ b/rust/sedona-functions/src/st_force_dim.rs @@ -0,0 +1,128 @@ +// 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 std::sync::Arc; + +use arrow_array::builder::BinaryBuilder; +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::{ + error::SedonaGeometryError, + transform::{transform, CrsTransform}, + wkb_factory::WKB_MIN_PROBABLE_BYTES, +}; +use sedona_schema::{ + datatypes::{SedonaType, WKB_GEOGRAPHY, WKB_GEOMETRY}, + matchers::ArgMatcher, +}; + +use crate::executor::WkbExecutor; + +/// ST_Force2D() scalar UDF +pub fn st_force2d_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "st_force2d", + ItemCrsKernel::wrap_impl(vec![ + Arc::new(STForce2D { + is_geography: false, + }), + Arc::new(STForce2D { is_geography: true }), + ]), + Volatility::Immutable, + Some(st_force2d_doc()), + ) +} + +fn st_force2d_doc() -> Documentation { + Documentation::builder( + DOC_SECTION_OTHER, + "Update coordinates of geom by a fixed offset", + "ST_Force2D (geom: Geometry)", + ) + .with_argument("geom", "geometry: Input geometry") + .with_sql_example("SELECT ST_Force2D(ST_GeomFromWKT('POINT Z (1 2 3)'))") + .build() +} + +#[derive(Debug)] +struct STForce2D { + is_geography: bool, +} + +impl SedonaScalarKernel for STForce2D { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let matcher = if self.is_geography { + ArgMatcher::new(vec![ArgMatcher::is_geography()], WKB_GEOGRAPHY) + } else { + ArgMatcher::new(vec![ArgMatcher::is_geometry()], WKB_GEOMETRY) + }; + + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + let executor = WkbExecutor::new(arg_types, args); + let mut builder = BinaryBuilder::with_capacity( + executor.num_iterations(), + WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), + ); + + let trans = Force2DTransform {}; + executor.execute_wkb_void(|maybe_wkb| { + match maybe_wkb { + Some(wkb) => { + transform(wkb, &trans, &mut builder) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + builder.append_value([]); + } + _ => { + builder.append_null(); + } + } + + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + } +} + +#[derive(Debug)] +struct Force2DTransform {} + +impl CrsTransform for Force2DTransform { + fn output_dim(&self) -> Option { + Some(geo_traits::Dimensions::Xy) + } + + fn transform_coord( + &self, + _coord: &mut (f64, f64), + ) -> std::result::Result<(), SedonaGeometryError> { + Ok(()) + } +} From 0450104702c7fad3611e963110af19a52f463a33 Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 01:27:55 +0900 Subject: [PATCH 03/15] Implement ST_Force3D --- rust/sedona-functions/src/register.rs | 1 + rust/sedona-functions/src/st_force_dim.rs | 144 +++++++++++++++++++++- rust/sedona-geometry/src/transform.rs | 88 ++++++++----- 3 files changed, 200 insertions(+), 33 deletions(-) diff --git a/rust/sedona-functions/src/register.rs b/rust/sedona-functions/src/register.rs index 3feafd5fa..a44c9dd2e 100644 --- a/rust/sedona-functions/src/register.rs +++ b/rust/sedona-functions/src/register.rs @@ -119,6 +119,7 @@ pub fn default_function_set() -> FunctionSet { crate::st_transform::st_transform_udf, crate::st_translate::st_translate_udf, crate::st_force_dim::st_force2d_udf, + crate::st_force_dim::st_force3d_udf, crate::st_xyzm_minmax::st_mmax_udf, crate::st_xyzm_minmax::st_mmin_udf, crate::st_xyzm_minmax::st_xmax_udf, diff --git a/rust/sedona-functions/src/st_force_dim.rs b/rust/sedona-functions/src/st_force_dim.rs index 606796292..abb3c8b43 100644 --- a/rust/sedona-functions/src/st_force_dim.rs +++ b/rust/sedona-functions/src/st_force_dim.rs @@ -17,11 +17,13 @@ use std::sync::Arc; -use arrow_array::builder::BinaryBuilder; -use datafusion_common::{error::Result, DataFusionError}; +use arrow_array::{builder::BinaryBuilder, Array, Float64Array}; +use arrow_schema::DataType; +use datafusion_common::{cast::as_float64_array, error::Result, DataFusionError}; use datafusion_expr::{ scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility, }; +use geo_traits::Dimensions; use sedona_expr::{ item_crs::ItemCrsKernel, scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, @@ -38,6 +40,8 @@ use sedona_schema::{ use crate::executor::WkbExecutor; +// *** 2D ************************* + /// ST_Force2D() scalar UDF pub fn st_force2d_udf() -> SedonaScalarUDF { SedonaScalarUDF::new( @@ -56,7 +60,7 @@ pub fn st_force2d_udf() -> SedonaScalarUDF { fn st_force2d_doc() -> Documentation { Documentation::builder( DOC_SECTION_OTHER, - "Update coordinates of geom by a fixed offset", + "Forces the geometry into a 2-dimensional model", "ST_Force2D (geom: Geometry)", ) .with_argument("geom", "geometry: Input geometry") @@ -122,7 +126,141 @@ impl CrsTransform for Force2DTransform { fn transform_coord( &self, _coord: &mut (f64, f64), + _input_dims: Dimensions, + ) -> std::result::Result<(), SedonaGeometryError> { + Ok(()) + } +} + +// *** 3D ************************* + +/// ST_Force3D() scalar UDF +pub fn st_force3d_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "st_force3d", + ItemCrsKernel::wrap_impl(vec![ + Arc::new(STForce3D { + is_geography: false, + }), + Arc::new(STForce3D { is_geography: true }), + ]), + Volatility::Immutable, + Some(st_force3d_doc()), + ) +} + +fn st_force3d_doc() -> Documentation { + Documentation::builder( + DOC_SECTION_OTHER, + "Forces the geometry into a 3-dimensional model.", + "ST_Force3D (geom: Geometry)", + ) + .with_argument("geom", "geometry: Input geometry") + .with_argument("z", "numeric: default Z value") + .with_sql_example("SELECT ST_Force3D(ST_GeomFromWKT('POINT (1 2)'))") + .build() +} + +#[derive(Debug)] +struct STForce3D { + is_geography: bool, +} + +impl SedonaScalarKernel for STForce3D { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let matcher = if self.is_geography { + ArgMatcher::new( + vec![ + ArgMatcher::is_geography(), + ArgMatcher::optional(ArgMatcher::is_numeric()), + ], + WKB_GEOGRAPHY, + ) + } else { + ArgMatcher::new( + vec![ + ArgMatcher::is_geometry(), + ArgMatcher::optional(ArgMatcher::is_numeric()), + ], + WKB_GEOMETRY, + ) + }; + + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + let executor = WkbExecutor::new(arg_types, args); + let mut builder = BinaryBuilder::with_capacity( + executor.num_iterations(), + WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), + ); + + let z_array = match args.get(1) { + Some(arg) => arg + .cast_to(&DataType::Float64, None)? + .to_array(executor.num_iterations())?, + None => Arc::new(Float64Array::from(vec![0.0; executor.num_iterations()])), + }; + let z_array = as_float64_array(&z_array)?; + + let mut i = 0usize; + executor.execute_wkb_void(|maybe_wkb| { + match (maybe_wkb, z_array.is_null(i)) { + (Some(wkb), false) => { + let trans = Force3DTransform { + z: z_array.value(i), + }; + transform(wkb, &trans, &mut builder) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + builder.append_value([]); + } + _ => { + builder.append_null(); + } + } + i += 1; + + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + } +} + +#[derive(Debug)] +struct Force3DTransform { + z: f64, +} + +impl CrsTransform for Force3DTransform { + fn output_dim(&self) -> Option { + Some(geo_traits::Dimensions::Xyz) + } + + fn transform_coord( + &self, + _coord: &mut (f64, f64), + _input_dims: Dimensions, ) -> std::result::Result<(), SedonaGeometryError> { + unreachable!() + } + fn transform_coord_3d( + &self, + coord: &mut (f64, f64, f64), + input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { + // If the input doesn't have Z coordinate, fill with the default value + if matches!( + input_dims, + Dimensions::Xy | Dimensions::Xym | Dimensions::Unknown(_) + ) { + coord.2 = self.z + } Ok(()) } } diff --git a/rust/sedona-geometry/src/transform.rs b/rust/sedona-geometry/src/transform.rs index af6c9bffb..56a66bda0 100644 --- a/rust/sedona-geometry/src/transform.rs +++ b/rust/sedona-geometry/src/transform.rs @@ -57,13 +57,21 @@ 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>; + fn transform_coord( + &self, + coord: &mut (f64, f64), + input_dims: Dimensions, + ) -> 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> { + fn transform_coord_3d( + &self, + coord: &mut (f64, f64, f64), + input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { let mut coord_2d = (coord.0, coord.1); - self.transform_coord(&mut coord_2d)?; + self.transform_coord(&mut coord_2d, input_dims)?; coord.0 = coord_2d.0; coord.1 = coord_2d.1; Ok(()) @@ -77,12 +85,20 @@ pub trait CrsTransform: std::fmt::Debug { /// A boxed trait object for dynamic dispatch of CRS transformations. impl CrsTransform for Box { - fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { - self.as_ref().transform_coord(coord) + fn transform_coord( + &self, + coord: &mut (f64, f64), + input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { + self.as_ref().transform_coord(coord, input_dims) } - fn transform_coord_3d(&self, coord: &mut (f64, f64, f64)) -> Result<(), SedonaGeometryError> { - self.as_ref().transform_coord_3d(coord) + fn transform_coord_3d( + &self, + coord: &mut (f64, f64, f64), + input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { + self.as_ref().transform_coord_3d(coord, input_dims) } } @@ -282,52 +298,52 @@ pub fn transform( ) -> Result<(), SedonaGeometryError> { // If the CrsTransform specifies the dimension, use it. // Otherwise, the input dimension is preserved. - let dims = trans.output_dim().unwrap_or_else(|| geom.dim()); + let output_dims = trans.output_dim().unwrap_or_else(|| geom.dim()); match geom.as_type() { GeometryType::Point(pt) => { if pt.coord().is_some() { - write_wkb_point_header(out, dims)?; - transform_and_write_coords(out, trans, pt.coord().into_iter(), dims)?; + write_wkb_point_header(out, output_dims)?; + transform_and_write_coords(out, trans, pt.coord().into_iter(), output_dims)?; } else { - write_wkb_empty_point(out, dims)?; + write_wkb_empty_point(out, output_dims)?; } } GeometryType::LineString(ls) => { - write_wkb_linestring_header(out, dims, ls.coords().count())?; - transform_and_write_coords(out, trans, ls.coords(), dims)?; + write_wkb_linestring_header(out, output_dims, ls.coords().count())?; + transform_and_write_coords(out, trans, ls.coords(), output_dims)?; } GeometryType::Polygon(pl) => { let num_rings = pl.interiors().count() + pl.exterior().is_some() as usize; - write_wkb_polygon_header(out, dims, num_rings)?; + write_wkb_polygon_header(out, output_dims, num_rings)?; if let Some(exterior) = pl.exterior() { - transform_and_write_ring(out, trans, exterior, dims)?; + transform_and_write_ring(out, trans, exterior, output_dims)?; } for interior in pl.interiors() { - transform_and_write_ring(out, trans, interior, dims)?; + transform_and_write_ring(out, trans, interior, output_dims)?; } } GeometryType::MultiPoint(multi_pt) => { - write_wkb_multipoint_header(out, dims, multi_pt.points().count())?; + write_wkb_multipoint_header(out, output_dims, multi_pt.points().count())?; for pt in multi_pt.points() { transform(pt, trans, out)?; } } GeometryType::MultiLineString(multi_ls) => { - write_wkb_multilinestring_header(out, dims, multi_ls.line_strings().count())?; + write_wkb_multilinestring_header(out, output_dims, multi_ls.line_strings().count())?; for ls in multi_ls.line_strings() { transform(ls, trans, out)?; } } GeometryType::MultiPolygon(multi_pl) => { - write_wkb_multipolygon_header(out, dims, multi_pl.polygons().count())?; + write_wkb_multipolygon_header(out, output_dims, multi_pl.polygons().count())?; for pl in multi_pl.polygons() { transform(pl, trans, out)?; } } GeometryType::GeometryCollection(collection) => { - write_wkb_geometrycollection_header(out, dims, collection.geometries().count())?; + write_wkb_geometrycollection_header(out, output_dims, collection.geometries().count())?; for geom in collection.geometries() { transform(geom, trans, out)?; } @@ -346,14 +362,14 @@ fn transform_and_write_ring<'a, L>( buf: &mut impl Write, trans: &dyn CrsTransform, ring: L, - dims: Dimensions, + output_dims: Dimensions, ) -> Result<(), SedonaGeometryError> where L: LineStringTrait + 'a, { let num_points = ring.coords().count(); write_wkb_polygon_ring_header(buf, num_points)?; - transform_and_write_coords(buf, trans, ring.coords(), dims)?; + transform_and_write_coords(buf, trans, ring.coords(), output_dims)?; Ok(()) } @@ -361,32 +377,33 @@ fn transform_and_write_coords<'a, C, I>( buf: &mut impl Write, trans: &dyn CrsTransform, coords: I, - dims: Dimensions, + output_dims: Dimensions, ) -> Result<(), SedonaGeometryError> where C: CoordTrait + 'a, I: Iterator, { for coord in coords { - match dims { + let input_dims = coord.dim(); + match output_dims { Dimensions::Xy => { let mut xy: (f64, f64) = (coord.x(), coord.y()); - trans.transform_coord(&mut xy)?; + trans.transform_coord(&mut xy, input_dims)?; write_wkb_coord(buf, (xy.0, xy.1))?; } Dimensions::Xyz => { - let mut xyz: (f64, f64, f64) = (coord.x(), coord.y(), coord.nth_or_panic(2)); - trans.transform_coord_3d(&mut xyz)?; + let mut xyz = fill_or_extract_coord(&coord, input_dims); + trans.transform_coord_3d(&mut xyz, input_dims)?; write_wkb_coord(buf, (xyz.0, xyz.1, xyz.2))?; } Dimensions::Xym => { let mut xy: (f64, f64) = (coord.x(), coord.y()); - trans.transform_coord(&mut xy)?; + trans.transform_coord(&mut xy, input_dims)?; write_wkb_coord(buf, (xy.0, xy.1, coord.nth_or_panic(2)))?; } Dimensions::Xyzm => { - let mut xyz: (f64, f64, f64) = (coord.x(), coord.y(), coord.nth_or_panic(2)); - trans.transform_coord_3d(&mut xyz)?; + let mut xyz = fill_or_extract_coord(&coord, input_dims); + trans.transform_coord_3d(&mut xyz, input_dims)?; write_wkb_coord(buf, (xyz.0, xyz.1, xyz.2, coord.nth_or_panic(3)))?; } _ => { @@ -399,6 +416,17 @@ where Ok(()) } +fn fill_or_extract_coord(coord: &C, input_dims: Dimensions) -> (f64, f64, f64) +where + C: CoordTrait, +{ + match input_dims { + // If the input doesn't have Z coordinate, fill with 0. + Dimensions::Xy | Dimensions::Xym | Dimensions::Unknown(_) => (coord.x(), coord.y(), 0.0), + Dimensions::Xyz | Dimensions::Xyzm => (coord.x(), coord.y(), coord.nth_or_panic(2)), + } +} + #[cfg(test)] mod test { use super::*; From e31d8bb4930d4d975774fdf4d62fcbf8c8f3b296 Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 01:39:43 +0900 Subject: [PATCH 04/15] Fix signatures --- c/sedona-proj/src/sd_order_lnglat.rs | 3 ++- c/sedona-proj/src/transform.rs | 23 +++++++++++++++---- .../sedona-functions/src/st_affine_helpers.rs | 12 ++++++++-- .../src/st_flipcoordinates.rs | 2 ++ rust/sedona-functions/src/st_translate.rs | 8 ++++++- rust/sedona-geometry/src/transform.rs | 13 +++++++++-- 6 files changed, 51 insertions(+), 10 deletions(-) diff --git a/c/sedona-proj/src/sd_order_lnglat.rs b/c/sedona-proj/src/sd_order_lnglat.rs index 9d95de363..488faa559 100644 --- a/c/sedona-proj/src/sd_order_lnglat.rs +++ b/c/sedona-proj/src/sd_order_lnglat.rs @@ -21,6 +21,7 @@ use arrow_array::builder::UInt64Builder; use arrow_schema::DataType; use datafusion_common::{DataFusionError, Result}; use datafusion_expr::ColumnarValue; +use geo_traits::Dimensions; use sedona_expr::scalar_udf::SedonaScalarKernel; use sedona_functions::executor::WkbBytesExecutor; use sedona_geometry::{transform::CrsEngine, wkb_header::WkbHeader}; @@ -99,7 +100,7 @@ impl u64 + Send + Sync> SedonaScalarKernel for OrderLngLat< .map_err(|e| DataFusionError::Execution(format!("{e}")))?; let mut first_xy = header.first_xy(); to_lnglat - .transform_coord(&mut first_xy) + .transform_coord(&mut first_xy, Dimensions::Xy) .map_err(|e| DataFusionError::Execution(format!("{e}")))?; let order = (self.order_fn)(first_xy); builder.append_value(order); diff --git a/c/sedona-proj/src/transform.rs b/c/sedona-proj/src/transform.rs index 1801b6539..428b00c52 100644 --- a/c/sedona-proj/src/transform.rs +++ b/c/sedona-proj/src/transform.rs @@ -16,6 +16,7 @@ // under the License. use crate::error::SedonaProjError; use crate::proj::{Proj, ProjContext}; +use geo_traits::Dimensions; use sedona_geometry::bounding_box::BoundingBox; use sedona_geometry::error::SedonaGeometryError; use sedona_geometry::interval::IntervalTrait; @@ -218,7 +219,11 @@ impl ProjTransform { } impl CrsTransform for ProjTransform { - fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { + fn transform_coord( + &self, + coord: &mut (f64, f64), + _input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { let res = self.proj.borrow_mut().transform_xy(*coord).map_err(|e| { SedonaGeometryError::Invalid(format!( "PROJ coordinate transformation failed with error: {e}" @@ -229,7 +234,11 @@ impl CrsTransform for ProjTransform { Ok(()) } - fn transform_coord_3d(&self, coord: &mut (f64, f64, f64)) -> Result<(), SedonaGeometryError> { + fn transform_coord_3d( + &self, + coord: &mut (f64, f64, f64), + _input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { let res = self.proj.borrow_mut().transform_xyz(*coord).map_err(|e| { SedonaGeometryError::Invalid(format!( "PROJ coordinate transformation failed with error: {e}" @@ -306,12 +315,18 @@ mod test { .unwrap(); let mut coord = (4_760_096.4, 3_744_293.5); - trans.as_ref().transform_coord(&mut coord).unwrap(); + trans + .as_ref() + .transform_coord(&mut coord, Dimensions::Xy) + .unwrap(); assert_relative_eq!(coord.x(), 1_450_880.284_378, epsilon = 1e-6); assert_relative_eq!(coord.y(), 1_141_262.941_224, epsilon = 1e-6); coord = (f64::NAN, f64::NAN); - trans.as_ref().transform_coord(&mut coord).unwrap(); + trans + .as_ref() + .transform_coord(&mut coord, Dimensions::Xy) + .unwrap(); assert!( coord.x().is_nan() && coord.y().is_nan(), "Expected NaN coordinates" diff --git a/rust/sedona-functions/src/st_affine_helpers.rs b/rust/sedona-functions/src/st_affine_helpers.rs index 7de1ba945..bf0166aea 100644 --- a/rust/sedona-functions/src/st_affine_helpers.rs +++ b/rust/sedona-functions/src/st_affine_helpers.rs @@ -19,6 +19,7 @@ use arrow_array::Array; use arrow_array::PrimitiveArray; use datafusion_common::cast::as_float64_array; use datafusion_common::error::Result; +use geo_traits::Dimensions; use sedona_common::sedona_internal_err; use sedona_geometry::transform::CrsTransform; use std::sync::Arc; @@ -438,6 +439,7 @@ impl CrsTransform for DAffine { fn transform_coord_3d( &self, coord: &mut (f64, f64, f64), + _input_dims: Dimensions, ) -> std::result::Result<(), sedona_geometry::error::SedonaGeometryError> { match self { DAffine::DAffine2(daffine2) => { @@ -466,6 +468,7 @@ impl CrsTransform for DAffine { fn transform_coord( &self, coord: &mut (f64, f64), + _input_dims: Dimensions, ) -> std::result::Result<(), sedona_geometry::error::SedonaGeometryError> { match self { DAffine::DAffine2(daffine2) => { @@ -536,6 +539,7 @@ mod tests { use super::*; use arrow_array::Array; use arrow_array::Float64Array; + use geo_traits::Dimensions; use std::sync::Arc; fn float_array(values: Vec>) -> Arc { @@ -622,13 +626,17 @@ mod tests { fn daffine_crs_transform_changes_coords() { let mut coord_2d = (1.0, 2.0); let affine_2d = DAffine::DAffine2(glam::DAffine2::from_scale(glam::DVec2::new(2.0, 3.0))); - affine_2d.transform_coord(&mut coord_2d).unwrap(); + affine_2d + .transform_coord(&mut coord_2d, Dimensions::Xy) + .unwrap(); assert_eq!(coord_2d, (2.0, 6.0)); let mut coord_3d = (1.0, 2.0, 3.0); let affine_3d = DAffine::DAffine3(glam::DAffine3::from_scale(glam::DVec3::new(2.0, 3.0, 4.0))); - affine_3d.transform_coord_3d(&mut coord_3d).unwrap(); + affine_3d + .transform_coord_3d(&mut coord_3d, Dimensions::Xyz) + .unwrap(); assert_eq!(coord_3d, (2.0, 6.0, 12.0)); } } diff --git a/rust/sedona-functions/src/st_flipcoordinates.rs b/rust/sedona-functions/src/st_flipcoordinates.rs index 42e5bcdc1..b478befe0 100644 --- a/rust/sedona-functions/src/st_flipcoordinates.rs +++ b/rust/sedona-functions/src/st_flipcoordinates.rs @@ -22,6 +22,7 @@ use datafusion_common::error::{DataFusionError, Result}; use datafusion_expr::{ scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility, }; +use geo_traits::Dimensions; use sedona_expr::{ item_crs::ItemCrsKernel, scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, @@ -122,6 +123,7 @@ impl CrsTransform for SwapXy { fn transform_coord( &self, coord: &mut (f64, f64), + _input_dims: Dimensions, ) -> std::result::Result<(), SedonaGeometryError> { let (x, y) = *coord; *coord = (y, x); diff --git a/rust/sedona-functions/src/st_translate.rs b/rust/sedona-functions/src/st_translate.rs index 647e29cf4..1bb60c532 100644 --- a/rust/sedona-functions/src/st_translate.rs +++ b/rust/sedona-functions/src/st_translate.rs @@ -20,6 +20,7 @@ use datafusion_common::{cast::as_float64_array, error::Result, DataFusionError}; use datafusion_expr::{ scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility, }; +use geo_traits::Dimensions; use sedona_common::sedona_internal_err; use sedona_expr::{ @@ -225,7 +226,11 @@ struct Translate { } impl CrsTransform for Translate { - fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { + fn transform_coord( + &self, + coord: &mut (f64, f64), + _input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { coord.0 += self.deltax; coord.1 += self.deltay; Ok(()) @@ -234,6 +239,7 @@ impl CrsTransform for Translate { fn transform_coord_3d( &self, coord: &mut (f64, f64, f64), + _input_dims: Dimensions, ) -> std::result::Result<(), SedonaGeometryError> { coord.0 += self.deltax; coord.1 += self.deltay; diff --git a/rust/sedona-geometry/src/transform.rs b/rust/sedona-geometry/src/transform.rs index 56a66bda0..b004db7c4 100644 --- a/rust/sedona-geometry/src/transform.rs +++ b/rust/sedona-geometry/src/transform.rs @@ -437,7 +437,11 @@ mod test { #[derive(Debug)] struct MockTransform {} impl CrsTransform for MockTransform { - fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { + fn transform_coord( + &self, + coord: &mut (f64, f64), + _input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { coord.0 += 10.0; coord.1 += 20.0; Ok(()) @@ -448,7 +452,11 @@ mod test { struct Mock3DTransform {} impl CrsTransform for Mock3DTransform { // This transforms 2D and 3D differently for testing purposes - fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { + fn transform_coord( + &self, + coord: &mut (f64, f64), + _input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { coord.0 += 100.0; coord.1 += 200.0; Ok(()) @@ -457,6 +465,7 @@ mod test { fn transform_coord_3d( &self, coord: &mut (f64, f64, f64), + _input_dims: Dimensions, ) -> Result<(), SedonaGeometryError> { coord.0 += 10.0; coord.1 += 20.0; From 8536401c1ebc42534a39b708b4102d201f81f9d4 Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 01:47:35 +0900 Subject: [PATCH 05/15] Tweak --- c/sedona-proj/src/sd_order_lnglat.rs | 2 +- c/sedona-proj/src/transform.rs | 10 ++---- .../sedona-functions/src/st_affine_helpers.rs | 5 +-- .../src/st_flipcoordinates.rs | 2 -- rust/sedona-functions/src/st_force_dim.rs | 2 -- rust/sedona-functions/src/st_translate.rs | 6 +--- rust/sedona-geometry/src/transform.rs | 34 +++++-------------- 7 files changed, 14 insertions(+), 47 deletions(-) diff --git a/c/sedona-proj/src/sd_order_lnglat.rs b/c/sedona-proj/src/sd_order_lnglat.rs index 488faa559..0a68f6f30 100644 --- a/c/sedona-proj/src/sd_order_lnglat.rs +++ b/c/sedona-proj/src/sd_order_lnglat.rs @@ -100,7 +100,7 @@ impl u64 + Send + Sync> SedonaScalarKernel for OrderLngLat< .map_err(|e| DataFusionError::Execution(format!("{e}")))?; let mut first_xy = header.first_xy(); to_lnglat - .transform_coord(&mut first_xy, Dimensions::Xy) + .transform_coord(&mut first_xy) .map_err(|e| DataFusionError::Execution(format!("{e}")))?; let order = (self.order_fn)(first_xy); builder.append_value(order); diff --git a/c/sedona-proj/src/transform.rs b/c/sedona-proj/src/transform.rs index 428b00c52..487d16a25 100644 --- a/c/sedona-proj/src/transform.rs +++ b/c/sedona-proj/src/transform.rs @@ -315,18 +315,12 @@ mod test { .unwrap(); let mut coord = (4_760_096.4, 3_744_293.5); - trans - .as_ref() - .transform_coord(&mut coord, Dimensions::Xy) - .unwrap(); + trans.as_ref().transform_coord(&mut coord).unwrap(); assert_relative_eq!(coord.x(), 1_450_880.284_378, epsilon = 1e-6); assert_relative_eq!(coord.y(), 1_141_262.941_224, epsilon = 1e-6); coord = (f64::NAN, f64::NAN); - trans - .as_ref() - .transform_coord(&mut coord, Dimensions::Xy) - .unwrap(); + trans.as_ref().transform_coord(&mut coord).unwrap(); assert!( coord.x().is_nan() && coord.y().is_nan(), "Expected NaN coordinates" diff --git a/rust/sedona-functions/src/st_affine_helpers.rs b/rust/sedona-functions/src/st_affine_helpers.rs index bf0166aea..40ce0b9d8 100644 --- a/rust/sedona-functions/src/st_affine_helpers.rs +++ b/rust/sedona-functions/src/st_affine_helpers.rs @@ -468,7 +468,6 @@ impl CrsTransform for DAffine { fn transform_coord( &self, coord: &mut (f64, f64), - _input_dims: Dimensions, ) -> std::result::Result<(), sedona_geometry::error::SedonaGeometryError> { match self { DAffine::DAffine2(daffine2) => { @@ -626,9 +625,7 @@ mod tests { fn daffine_crs_transform_changes_coords() { let mut coord_2d = (1.0, 2.0); let affine_2d = DAffine::DAffine2(glam::DAffine2::from_scale(glam::DVec2::new(2.0, 3.0))); - affine_2d - .transform_coord(&mut coord_2d, Dimensions::Xy) - .unwrap(); + affine_2d.transform_coord(&mut coord_2d).unwrap(); assert_eq!(coord_2d, (2.0, 6.0)); let mut coord_3d = (1.0, 2.0, 3.0); diff --git a/rust/sedona-functions/src/st_flipcoordinates.rs b/rust/sedona-functions/src/st_flipcoordinates.rs index b478befe0..42e5bcdc1 100644 --- a/rust/sedona-functions/src/st_flipcoordinates.rs +++ b/rust/sedona-functions/src/st_flipcoordinates.rs @@ -22,7 +22,6 @@ use datafusion_common::error::{DataFusionError, Result}; use datafusion_expr::{ scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility, }; -use geo_traits::Dimensions; use sedona_expr::{ item_crs::ItemCrsKernel, scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, @@ -123,7 +122,6 @@ impl CrsTransform for SwapXy { fn transform_coord( &self, coord: &mut (f64, f64), - _input_dims: Dimensions, ) -> std::result::Result<(), SedonaGeometryError> { let (x, y) = *coord; *coord = (y, x); diff --git a/rust/sedona-functions/src/st_force_dim.rs b/rust/sedona-functions/src/st_force_dim.rs index abb3c8b43..b7af99cd4 100644 --- a/rust/sedona-functions/src/st_force_dim.rs +++ b/rust/sedona-functions/src/st_force_dim.rs @@ -126,7 +126,6 @@ impl CrsTransform for Force2DTransform { fn transform_coord( &self, _coord: &mut (f64, f64), - _input_dims: Dimensions, ) -> std::result::Result<(), SedonaGeometryError> { Ok(()) } @@ -245,7 +244,6 @@ impl CrsTransform for Force3DTransform { fn transform_coord( &self, _coord: &mut (f64, f64), - _input_dims: Dimensions, ) -> std::result::Result<(), SedonaGeometryError> { unreachable!() } diff --git a/rust/sedona-functions/src/st_translate.rs b/rust/sedona-functions/src/st_translate.rs index 1bb60c532..c34dab1ed 100644 --- a/rust/sedona-functions/src/st_translate.rs +++ b/rust/sedona-functions/src/st_translate.rs @@ -226,11 +226,7 @@ struct Translate { } impl CrsTransform for Translate { - fn transform_coord( - &self, - coord: &mut (f64, f64), - _input_dims: Dimensions, - ) -> Result<(), SedonaGeometryError> { + fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { coord.0 += self.deltax; coord.1 += self.deltay; Ok(()) diff --git a/rust/sedona-geometry/src/transform.rs b/rust/sedona-geometry/src/transform.rs index b004db7c4..e1abcec21 100644 --- a/rust/sedona-geometry/src/transform.rs +++ b/rust/sedona-geometry/src/transform.rs @@ -57,21 +57,17 @@ 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), - input_dims: Dimensions, - ) -> Result<(), SedonaGeometryError>; + 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), - input_dims: Dimensions, + _input_dims: Dimensions, ) -> Result<(), SedonaGeometryError> { let mut coord_2d = (coord.0, coord.1); - self.transform_coord(&mut coord_2d, input_dims)?; + self.transform_coord(&mut coord_2d)?; coord.0 = coord_2d.0; coord.1 = coord_2d.1; Ok(()) @@ -85,12 +81,8 @@ pub trait CrsTransform: std::fmt::Debug { /// A boxed trait object for dynamic dispatch of CRS transformations. impl CrsTransform for Box { - fn transform_coord( - &self, - coord: &mut (f64, f64), - input_dims: Dimensions, - ) -> Result<(), SedonaGeometryError> { - self.as_ref().transform_coord(coord, input_dims) + fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { + self.as_ref().transform_coord(coord) } fn transform_coord_3d( @@ -388,7 +380,7 @@ where match output_dims { Dimensions::Xy => { let mut xy: (f64, f64) = (coord.x(), coord.y()); - trans.transform_coord(&mut xy, input_dims)?; + trans.transform_coord(&mut xy)?; write_wkb_coord(buf, (xy.0, xy.1))?; } Dimensions::Xyz => { @@ -398,7 +390,7 @@ where } Dimensions::Xym => { let mut xy: (f64, f64) = (coord.x(), coord.y()); - trans.transform_coord(&mut xy, input_dims)?; + trans.transform_coord(&mut xy)?; write_wkb_coord(buf, (xy.0, xy.1, coord.nth_or_panic(2)))?; } Dimensions::Xyzm => { @@ -437,11 +429,7 @@ mod test { #[derive(Debug)] struct MockTransform {} impl CrsTransform for MockTransform { - fn transform_coord( - &self, - coord: &mut (f64, f64), - _input_dims: Dimensions, - ) -> Result<(), SedonaGeometryError> { + fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { coord.0 += 10.0; coord.1 += 20.0; Ok(()) @@ -452,11 +440,7 @@ mod test { struct Mock3DTransform {} impl CrsTransform for Mock3DTransform { // This transforms 2D and 3D differently for testing purposes - fn transform_coord( - &self, - coord: &mut (f64, f64), - _input_dims: Dimensions, - ) -> Result<(), SedonaGeometryError> { + fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { coord.0 += 100.0; coord.1 += 200.0; Ok(()) From e8fda8af8b662ed88b57a768a46bba66b4a954ae Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 06:53:23 +0900 Subject: [PATCH 06/15] Fix errors --- c/sedona-proj/src/sd_order_lnglat.rs | 1 - c/sedona-proj/src/transform.rs | 6 +----- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/c/sedona-proj/src/sd_order_lnglat.rs b/c/sedona-proj/src/sd_order_lnglat.rs index 0a68f6f30..9d95de363 100644 --- a/c/sedona-proj/src/sd_order_lnglat.rs +++ b/c/sedona-proj/src/sd_order_lnglat.rs @@ -21,7 +21,6 @@ use arrow_array::builder::UInt64Builder; use arrow_schema::DataType; use datafusion_common::{DataFusionError, Result}; use datafusion_expr::ColumnarValue; -use geo_traits::Dimensions; use sedona_expr::scalar_udf::SedonaScalarKernel; use sedona_functions::executor::WkbBytesExecutor; use sedona_geometry::{transform::CrsEngine, wkb_header::WkbHeader}; diff --git a/c/sedona-proj/src/transform.rs b/c/sedona-proj/src/transform.rs index 487d16a25..a7bd9f077 100644 --- a/c/sedona-proj/src/transform.rs +++ b/c/sedona-proj/src/transform.rs @@ -219,11 +219,7 @@ impl ProjTransform { } impl CrsTransform for ProjTransform { - fn transform_coord( - &self, - coord: &mut (f64, f64), - _input_dims: Dimensions, - ) -> Result<(), SedonaGeometryError> { + fn transform_coord(&self, coord: &mut (f64, f64)) -> Result<(), SedonaGeometryError> { let res = self.proj.borrow_mut().transform_xy(*coord).map_err(|e| { SedonaGeometryError::Invalid(format!( "PROJ coordinate transformation failed with error: {e}" From 8137c4d41f4b005452cb9f5f21c8e2b9457ade23 Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 07:12:23 +0900 Subject: [PATCH 07/15] Add tests --- rust/sedona-functions/src/st_force_dim.rs | 135 ++++++++++++++++++---- 1 file changed, 114 insertions(+), 21 deletions(-) diff --git a/rust/sedona-functions/src/st_force_dim.rs b/rust/sedona-functions/src/st_force_dim.rs index b7af99cd4..2ae0f3833 100644 --- a/rust/sedona-functions/src/st_force_dim.rs +++ b/rust/sedona-functions/src/st_force_dim.rs @@ -118,18 +118,18 @@ impl SedonaScalarKernel for STForce2D { #[derive(Debug)] struct Force2DTransform {} -impl CrsTransform for Force2DTransform { - fn output_dim(&self) -> Option { - Some(geo_traits::Dimensions::Xy) - } - - fn transform_coord( - &self, - _coord: &mut (f64, f64), - ) -> std::result::Result<(), SedonaGeometryError> { - Ok(()) - } -} +impl CrsTransform for Force2DTransform { + fn output_dim(&self) -> Option { + Some(geo_traits::Dimensions::Xy) + } + + fn transform_coord( + &self, + _coord: &mut (f64, f64), + ) -> std::result::Result<(), SedonaGeometryError> { + Ok(()) + } +} // *** 3D ************************* @@ -241,12 +241,12 @@ impl CrsTransform for Force3DTransform { Some(geo_traits::Dimensions::Xyz) } - fn transform_coord( - &self, - _coord: &mut (f64, f64), - ) -> std::result::Result<(), SedonaGeometryError> { - unreachable!() - } + fn transform_coord( + &self, + _coord: &mut (f64, f64), + ) -> std::result::Result<(), SedonaGeometryError> { + unreachable!() + } fn transform_coord_3d( &self, coord: &mut (f64, f64, f64), @@ -259,6 +259,99 @@ impl CrsTransform for Force3DTransform { ) { coord.2 = self.z } - Ok(()) - } -} + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::create_array; + use datafusion_expr::ScalarUDF; + use rstest::rstest; + use sedona_schema::datatypes::WKB_VIEW_GEOMETRY; + use sedona_testing::{ + compare::assert_array_equal, create::create_array, testers::ScalarUdfTester, + }; + + use super::*; + + #[test] + fn udf_metadata() { + let st_force3d: ScalarUDF = st_force3d_udf().into(); + assert_eq!(st_force3d.name(), "st_force3d"); + assert!(st_force3d.documentation().is_some()); + } + + #[rstest] + fn udf_3d_without_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new(st_force3d_udf().into(), vec![sedona_type.clone()]); + tester.assert_return_type(WKB_GEOMETRY); + + let points = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT ZM (8 9 10 11)"), + ], + &sedona_type, + ); + + let expected = create_array( + &[ + None, + Some("POINT Z EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT Z (1 2 0)"), + Some("POINT Z (3 4 5)"), + Some("POINT Z (8 9 10)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points]).unwrap(); + assert_array_equal(&result, &expected); + } + + #[rstest] + fn udf_3d_with_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new( + st_force3d_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 (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT (6 7)"), + Some("POINT ZM (8 9 10 11)"), + ], + &sedona_type, + ); + let z = create_array!( + Float64, + [Some(9.0), Some(9.0), Some(9.0), Some(9.0), None, Some(9.0)] + ); + + let expected = create_array( + &[ + None, + Some("POINT Z EMPTY"), + Some("POINT Z (1 2 9)"), + Some("POINT Z (3 4 5)"), + None, + Some("POINT Z (8 9 10)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points, z]).unwrap(); + assert_array_equal(&result, &expected); + } +} From 6912bab7dbcb05dbef6005f853be8c39f784ce8d Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 07:41:55 +0900 Subject: [PATCH 08/15] Add Python tests --- .../tests/functions/test_functions.py | 42 ++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/python/sedonadb/tests/functions/test_functions.py b/python/sedonadb/tests/functions/test_functions.py index 34159360d..4621f3780 100644 --- a/python/sedonadb/tests/functions/test_functions.py +++ b/python/sedonadb/tests/functions/test_functions.py @@ -17,8 +17,8 @@ import math import pytest -import shapely import sedonadb +import shapely from sedonadb.testing import PostGIS, SedonaDB, geom_or_null, val_or_null @@ -1479,6 +1479,46 @@ def test_st_flipcoordinates(eng, geom, expected): ) +@pytest.mark.parametrize("eng", [SedonaDB, PostGIS]) +@pytest.mark.parametrize( + ("geom", "expected_2d", "expected_3d"), + [ + (None, None), + ("POINT EMPTY", "POINT (nan nan)", "POINT Z (nan nan nan)"), + ("POLYGON EMPTY", "POLYGON EMPTY", "POLYGON Z EMPTY"), + ("LINESTRING EMPTY", "LINESTRING EMPTY", "LINESTRING Z EMPTY"), + ("MULTIPOINT EMPTY", "MULTIPOINT EMPTY", "MULTIPOINT Z EMPTY"), + ("MULTILINESTRING EMPTY", "MULTILINESTRING EMPTY", "MULTILINESTRING Z EMPTY"), + ("MULTIPOLYGON EMPTY", "MULTIPOLYGON EMPTY", "MULTIPOLYGON Z EMPTY"), + ( + "GEOMETRYCOLLECTION EMPTY", + "GEOMETRYCOLLECTION EMPTY", + "GEOMETRYCOLLECTION EMPTY", + ), + ("POINT (0 1)", "POINT (1 0)", "POINT Z (1 0 5)"), + ( + "LINESTRING (0 1, 2 3)", + "LINESTRING (1 0, 3 2)", + "LINESTRING Z (1 0 5, 3 2 5)", + ), + ( + "MULTIPOINT (0 1, 2 3)", + "MULTIPOINT (1 0, 3 2)", + "MULTIPOINT Z (1 0 5, 3 2 5)", + ), + ( + "GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (3 4, 5 6), POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0)))", + "GEOMETRYCOLLECTION (POINT (2 1), LINESTRING (4 3, 6 5), POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0)))", + "GEOMETRYCOLLECTION (POINT Z (2 1 5), LINESTRING Z (4 3 5, 6 5 5), POLYGON Z ((0 0 5, 1 0 5, 1 1 5, 0 1 5, 0 0 5)))", + ), + ], +) +def test_st_force_dims(eng, geom, expected_2d, expected_3d): + eng = eng.create_or_skip() + eng.assert_query_result(f"SELECT ST_Force2D({geom_or_null(geom)})", expected_2d) + eng.assert_query_result(f"SELECT ST_Force3D({geom_or_null(geom)}, 5)", expected_3d) + + @pytest.mark.parametrize("eng", [SedonaDB, PostGIS]) @pytest.mark.parametrize( ("geom", "expected"), From e222cb493c8f6a2035eacc0bbee5a575709cd704 Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 08:25:24 +0900 Subject: [PATCH 09/15] Add more tests --- rust/sedona-functions/src/st_force_dim.rs | 265 ++++++++++++---------- 1 file changed, 151 insertions(+), 114 deletions(-) diff --git a/rust/sedona-functions/src/st_force_dim.rs b/rust/sedona-functions/src/st_force_dim.rs index 2ae0f3833..8b2609595 100644 --- a/rust/sedona-functions/src/st_force_dim.rs +++ b/rust/sedona-functions/src/st_force_dim.rs @@ -118,18 +118,18 @@ impl SedonaScalarKernel for STForce2D { #[derive(Debug)] struct Force2DTransform {} -impl CrsTransform for Force2DTransform { - fn output_dim(&self) -> Option { - Some(geo_traits::Dimensions::Xy) - } - - fn transform_coord( - &self, - _coord: &mut (f64, f64), - ) -> std::result::Result<(), SedonaGeometryError> { - Ok(()) - } -} +impl CrsTransform for Force2DTransform { + fn output_dim(&self) -> Option { + Some(geo_traits::Dimensions::Xy) + } + + fn transform_coord( + &self, + _coord: &mut (f64, f64), + ) -> std::result::Result<(), SedonaGeometryError> { + Ok(()) + } +} // *** 3D ************************* @@ -241,12 +241,12 @@ impl CrsTransform for Force3DTransform { Some(geo_traits::Dimensions::Xyz) } - fn transform_coord( - &self, - _coord: &mut (f64, f64), - ) -> std::result::Result<(), SedonaGeometryError> { - unreachable!() - } + fn transform_coord( + &self, + _coord: &mut (f64, f64), + ) -> std::result::Result<(), SedonaGeometryError> { + unreachable!() + } fn transform_coord_3d( &self, coord: &mut (f64, f64, f64), @@ -259,99 +259,136 @@ impl CrsTransform for Force3DTransform { ) { coord.2 = self.z } - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use arrow_array::create_array; - use datafusion_expr::ScalarUDF; - use rstest::rstest; - use sedona_schema::datatypes::WKB_VIEW_GEOMETRY; - use sedona_testing::{ - compare::assert_array_equal, create::create_array, testers::ScalarUdfTester, - }; - - use super::*; - - #[test] - fn udf_metadata() { - let st_force3d: ScalarUDF = st_force3d_udf().into(); - assert_eq!(st_force3d.name(), "st_force3d"); - assert!(st_force3d.documentation().is_some()); - } - - #[rstest] - fn udf_3d_without_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { - let tester = ScalarUdfTester::new(st_force3d_udf().into(), vec![sedona_type.clone()]); - tester.assert_return_type(WKB_GEOMETRY); - - let points = create_array( - &[ - None, - Some("POINT EMPTY"), - Some("POINT Z EMPTY"), - Some("POINT (1 2)"), - Some("POINT Z (3 4 5)"), - Some("POINT ZM (8 9 10 11)"), - ], - &sedona_type, - ); - - let expected = create_array( - &[ - None, - Some("POINT Z EMPTY"), - Some("POINT Z EMPTY"), - Some("POINT Z (1 2 0)"), - Some("POINT Z (3 4 5)"), - Some("POINT Z (8 9 10)"), - ], - &WKB_GEOMETRY, - ); - - let result = tester.invoke_arrays(vec![points]).unwrap(); - assert_array_equal(&result, &expected); - } - - #[rstest] - fn udf_3d_with_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { - let tester = ScalarUdfTester::new( - st_force3d_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 (1 2)"), - Some("POINT Z (3 4 5)"), - Some("POINT (6 7)"), - Some("POINT ZM (8 9 10 11)"), - ], - &sedona_type, - ); - let z = create_array!( - Float64, - [Some(9.0), Some(9.0), Some(9.0), Some(9.0), None, Some(9.0)] - ); - - let expected = create_array( - &[ - None, - Some("POINT Z EMPTY"), - Some("POINT Z (1 2 9)"), - Some("POINT Z (3 4 5)"), - None, - Some("POINT Z (8 9 10)"), - ], - &WKB_GEOMETRY, - ); - - let result = tester.invoke_arrays(vec![points, z]).unwrap(); - assert_array_equal(&result, &expected); - } -} + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::create_array; + use datafusion_expr::ScalarUDF; + use rstest::rstest; + use sedona_schema::datatypes::WKB_VIEW_GEOMETRY; + use sedona_testing::{ + compare::assert_array_equal, create::create_array, testers::ScalarUdfTester, + }; + + use super::*; + + #[test] + fn udf_metadata() { + let st_force2d: ScalarUDF = st_force2d_udf().into(); + assert_eq!(st_force2d.name(), "st_force2d"); + assert!(st_force2d.documentation().is_some()); + + let st_force3d: ScalarUDF = st_force3d_udf().into(); + assert_eq!(st_force3d.name(), "st_force3d"); + assert!(st_force3d.documentation().is_some()); + } + + #[rstest] + fn udf_2d(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new(st_force2d_udf().into(), vec![sedona_type.clone()]); + tester.assert_return_type(WKB_GEOMETRY); + + let points = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT ZM (8 9 10 11)"), + ], + &sedona_type, + ); + + let expected = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT EMPTY"), + Some("POINT (1 2)"), + Some("POINT (3 4)"), + Some("POINT (8 9)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points]).unwrap(); + assert_array_equal(&result, &expected); + } + + #[rstest] + fn udf_3d_without_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new(st_force3d_udf().into(), vec![sedona_type.clone()]); + tester.assert_return_type(WKB_GEOMETRY); + + let points = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT ZM (8 9 10 11)"), + ], + &sedona_type, + ); + + let expected = create_array( + &[ + None, + Some("POINT Z EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT Z (1 2 0)"), + Some("POINT Z (3 4 5)"), + Some("POINT Z (8 9 10)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points]).unwrap(); + assert_array_equal(&result, &expected); + } + + #[rstest] + fn udf_3d_with_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new( + st_force3d_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 (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT (6 7)"), + Some("POINT ZM (8 9 10 11)"), + ], + &sedona_type, + ); + let z = create_array!( + Float64, + [Some(9.0), Some(9.0), Some(9.0), Some(9.0), None, Some(9.0)] + ); + + let expected = create_array( + &[ + None, + Some("POINT Z EMPTY"), + Some("POINT Z (1 2 9)"), + Some("POINT Z (3 4 5)"), + None, + Some("POINT Z (8 9 10)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points, z]).unwrap(); + assert_array_equal(&result, &expected); + } +} From 96cfe6a1c3d2a357ca6b3e5f60097661ba10ea0a Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 08:28:23 +0900 Subject: [PATCH 10/15] Tweak --- python/sedonadb/tests/functions/test_functions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sedonadb/tests/functions/test_functions.py b/python/sedonadb/tests/functions/test_functions.py index 4621f3780..13e236546 100644 --- a/python/sedonadb/tests/functions/test_functions.py +++ b/python/sedonadb/tests/functions/test_functions.py @@ -1483,7 +1483,7 @@ def test_st_flipcoordinates(eng, geom, expected): @pytest.mark.parametrize( ("geom", "expected_2d", "expected_3d"), [ - (None, None), + (None, None, None), ("POINT EMPTY", "POINT (nan nan)", "POINT Z (nan nan nan)"), ("POLYGON EMPTY", "POLYGON EMPTY", "POLYGON Z EMPTY"), ("LINESTRING EMPTY", "LINESTRING EMPTY", "LINESTRING Z EMPTY"), @@ -1513,7 +1513,7 @@ def test_st_flipcoordinates(eng, geom, expected): ), ], ) -def test_st_force_dims(eng, geom, expected_2d, expected_3d): +def test_st_force_dim(eng, geom, expected_2d, expected_3d): eng = eng.create_or_skip() eng.assert_query_result(f"SELECT ST_Force2D({geom_or_null(geom)})", expected_2d) eng.assert_query_result(f"SELECT ST_Force3D({geom_or_null(geom)}, 5)", expected_3d) From b43e5b19320e3744b6edcfa16d85240bb98e40e3 Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 08:36:18 +0900 Subject: [PATCH 11/15] Fix test expectations --- .../sedonadb/tests/functions/test_functions.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/python/sedonadb/tests/functions/test_functions.py b/python/sedonadb/tests/functions/test_functions.py index 13e236546..779dc7623 100644 --- a/python/sedonadb/tests/functions/test_functions.py +++ b/python/sedonadb/tests/functions/test_functions.py @@ -1493,23 +1493,23 @@ def test_st_flipcoordinates(eng, geom, expected): ( "GEOMETRYCOLLECTION EMPTY", "GEOMETRYCOLLECTION EMPTY", - "GEOMETRYCOLLECTION EMPTY", + "GEOMETRYCOLLECTION Z EMPTY", ), - ("POINT (0 1)", "POINT (1 0)", "POINT Z (1 0 5)"), + ("POINT (0 1)", "POINT (0 1)", "POINT Z (0 1 5)"), ( "LINESTRING (0 1, 2 3)", - "LINESTRING (1 0, 3 2)", - "LINESTRING Z (1 0 5, 3 2 5)", + "LINESTRING (0 1, 2 3)", + "LINESTRING Z (0 1 5, 2 3 5)", ), ( "MULTIPOINT (0 1, 2 3)", - "MULTIPOINT (1 0, 3 2)", - "MULTIPOINT Z (1 0 5, 3 2 5)", + "MULTIPOINT (0 1, 2 3)", + "MULTIPOINT Z (0 1 5, 2 3 5)", ), ( "GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (3 4, 5 6), POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0)))", - "GEOMETRYCOLLECTION (POINT (2 1), LINESTRING (4 3, 6 5), POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0)))", - "GEOMETRYCOLLECTION (POINT Z (2 1 5), LINESTRING Z (4 3 5, 6 5 5), POLYGON Z ((0 0 5, 1 0 5, 1 1 5, 0 1 5, 0 0 5)))", + "GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (3 4, 5 6), POLYGON ((0 0, 0 1, 1 1, 1 0, 0 0)))", + "GEOMETRYCOLLECTION Z (POINT Z (1 2 5), LINESTRING Z (3 4 5, 5 6 5), POLYGON Z ((0 0 5, 0 1 5, 1 1 5, 1 0 5, 0 0 5)))", ), ], ) From a61a2997452b118d79524eaf2a25d082684a57ba Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 20:40:48 +0900 Subject: [PATCH 12/15] Apply suggestions from code review Co-authored-by: Dewey Dunnington --- rust/sedona-functions/src/st_force_dim.rs | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/rust/sedona-functions/src/st_force_dim.rs b/rust/sedona-functions/src/st_force_dim.rs index 8b2609595..3a98848a8 100644 --- a/rust/sedona-functions/src/st_force_dim.rs +++ b/rust/sedona-functions/src/st_force_dim.rs @@ -207,13 +207,11 @@ impl SedonaScalarKernel for STForce3D { }; let z_array = as_float64_array(&z_array)?; - let mut i = 0usize; + let mut z_iter = z_array.iter(); executor.execute_wkb_void(|maybe_wkb| { - match (maybe_wkb, z_array.is_null(i)) { - (Some(wkb), false) => { - let trans = Force3DTransform { - z: z_array.value(i), - }; + match (maybe_wkb, z_array.next().unwrap()) { + (Some(wkb), z) => { + let trans = Force3DTransform { z }; transform(wkb, &trans, &mut builder) .map_err(|e| DataFusionError::External(Box::new(e)))?; builder.append_value([]); @@ -222,7 +220,6 @@ impl SedonaScalarKernel for STForce3D { builder.append_null(); } } - i += 1; Ok(()) })?; @@ -245,7 +242,7 @@ impl CrsTransform for Force3DTransform { &self, _coord: &mut (f64, f64), ) -> std::result::Result<(), SedonaGeometryError> { - unreachable!() + sedona_internal_err!("Unexpected call to transform_coord()") } fn transform_coord_3d( &self, From 26d31f8b3b9e4a1a899f03e4609b1221dccd11d5 Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 20:48:22 +0900 Subject: [PATCH 13/15] Use z_iter --- rust/sedona-functions/src/st_force_dim.rs | 782 +++++++++++----------- 1 file changed, 391 insertions(+), 391 deletions(-) diff --git a/rust/sedona-functions/src/st_force_dim.rs b/rust/sedona-functions/src/st_force_dim.rs index 3a98848a8..cd813d63e 100644 --- a/rust/sedona-functions/src/st_force_dim.rs +++ b/rust/sedona-functions/src/st_force_dim.rs @@ -1,391 +1,391 @@ -// 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 std::sync::Arc; - -use arrow_array::{builder::BinaryBuilder, Array, Float64Array}; -use arrow_schema::DataType; -use datafusion_common::{cast::as_float64_array, error::Result, DataFusionError}; -use datafusion_expr::{ - scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility, -}; -use geo_traits::Dimensions; -use sedona_expr::{ - item_crs::ItemCrsKernel, - scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, -}; -use sedona_geometry::{ - error::SedonaGeometryError, - transform::{transform, CrsTransform}, - wkb_factory::WKB_MIN_PROBABLE_BYTES, -}; -use sedona_schema::{ - datatypes::{SedonaType, WKB_GEOGRAPHY, WKB_GEOMETRY}, - matchers::ArgMatcher, -}; - -use crate::executor::WkbExecutor; - -// *** 2D ************************* - -/// ST_Force2D() scalar UDF -pub fn st_force2d_udf() -> SedonaScalarUDF { - SedonaScalarUDF::new( - "st_force2d", - ItemCrsKernel::wrap_impl(vec![ - Arc::new(STForce2D { - is_geography: false, - }), - Arc::new(STForce2D { is_geography: true }), - ]), - Volatility::Immutable, - Some(st_force2d_doc()), - ) -} - -fn st_force2d_doc() -> Documentation { - Documentation::builder( - DOC_SECTION_OTHER, - "Forces the geometry into a 2-dimensional model", - "ST_Force2D (geom: Geometry)", - ) - .with_argument("geom", "geometry: Input geometry") - .with_sql_example("SELECT ST_Force2D(ST_GeomFromWKT('POINT Z (1 2 3)'))") - .build() -} - -#[derive(Debug)] -struct STForce2D { - is_geography: bool, -} - -impl SedonaScalarKernel for STForce2D { - fn return_type(&self, args: &[SedonaType]) -> Result> { - let matcher = if self.is_geography { - ArgMatcher::new(vec![ArgMatcher::is_geography()], WKB_GEOGRAPHY) - } else { - ArgMatcher::new(vec![ArgMatcher::is_geometry()], WKB_GEOMETRY) - }; - - matcher.match_args(args) - } - - fn invoke_batch( - &self, - arg_types: &[SedonaType], - args: &[ColumnarValue], - ) -> Result { - let executor = WkbExecutor::new(arg_types, args); - let mut builder = BinaryBuilder::with_capacity( - executor.num_iterations(), - WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), - ); - - let trans = Force2DTransform {}; - executor.execute_wkb_void(|maybe_wkb| { - match maybe_wkb { - Some(wkb) => { - transform(wkb, &trans, &mut builder) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - builder.append_value([]); - } - _ => { - builder.append_null(); - } - } - - Ok(()) - })?; - - executor.finish(Arc::new(builder.finish())) - } -} - -#[derive(Debug)] -struct Force2DTransform {} - -impl CrsTransform for Force2DTransform { - fn output_dim(&self) -> Option { - Some(geo_traits::Dimensions::Xy) - } - - fn transform_coord( - &self, - _coord: &mut (f64, f64), - ) -> std::result::Result<(), SedonaGeometryError> { - Ok(()) - } -} - -// *** 3D ************************* - -/// ST_Force3D() scalar UDF -pub fn st_force3d_udf() -> SedonaScalarUDF { - SedonaScalarUDF::new( - "st_force3d", - ItemCrsKernel::wrap_impl(vec![ - Arc::new(STForce3D { - is_geography: false, - }), - Arc::new(STForce3D { is_geography: true }), - ]), - Volatility::Immutable, - Some(st_force3d_doc()), - ) -} - -fn st_force3d_doc() -> Documentation { - Documentation::builder( - DOC_SECTION_OTHER, - "Forces the geometry into a 3-dimensional model.", - "ST_Force3D (geom: Geometry)", - ) - .with_argument("geom", "geometry: Input geometry") - .with_argument("z", "numeric: default Z value") - .with_sql_example("SELECT ST_Force3D(ST_GeomFromWKT('POINT (1 2)'))") - .build() -} - -#[derive(Debug)] -struct STForce3D { - is_geography: bool, -} - -impl SedonaScalarKernel for STForce3D { - fn return_type(&self, args: &[SedonaType]) -> Result> { - let matcher = if self.is_geography { - ArgMatcher::new( - vec![ - ArgMatcher::is_geography(), - ArgMatcher::optional(ArgMatcher::is_numeric()), - ], - WKB_GEOGRAPHY, - ) - } else { - ArgMatcher::new( - vec![ - ArgMatcher::is_geometry(), - ArgMatcher::optional(ArgMatcher::is_numeric()), - ], - WKB_GEOMETRY, - ) - }; - - matcher.match_args(args) - } - - fn invoke_batch( - &self, - arg_types: &[SedonaType], - args: &[ColumnarValue], - ) -> Result { - let executor = WkbExecutor::new(arg_types, args); - let mut builder = BinaryBuilder::with_capacity( - executor.num_iterations(), - WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), - ); - - let z_array = match args.get(1) { - Some(arg) => arg - .cast_to(&DataType::Float64, None)? - .to_array(executor.num_iterations())?, - None => Arc::new(Float64Array::from(vec![0.0; executor.num_iterations()])), - }; - let z_array = as_float64_array(&z_array)?; - - let mut z_iter = z_array.iter(); - executor.execute_wkb_void(|maybe_wkb| { - match (maybe_wkb, z_array.next().unwrap()) { - (Some(wkb), z) => { - let trans = Force3DTransform { z }; - transform(wkb, &trans, &mut builder) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - builder.append_value([]); - } - _ => { - builder.append_null(); - } - } - - Ok(()) - })?; - - executor.finish(Arc::new(builder.finish())) - } -} - -#[derive(Debug)] -struct Force3DTransform { - z: f64, -} - -impl CrsTransform for Force3DTransform { - fn output_dim(&self) -> Option { - Some(geo_traits::Dimensions::Xyz) - } - - fn transform_coord( - &self, - _coord: &mut (f64, f64), - ) -> std::result::Result<(), SedonaGeometryError> { - sedona_internal_err!("Unexpected call to transform_coord()") - } - fn transform_coord_3d( - &self, - coord: &mut (f64, f64, f64), - input_dims: Dimensions, - ) -> Result<(), SedonaGeometryError> { - // If the input doesn't have Z coordinate, fill with the default value - if matches!( - input_dims, - Dimensions::Xy | Dimensions::Xym | Dimensions::Unknown(_) - ) { - coord.2 = self.z - } - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use arrow_array::create_array; - use datafusion_expr::ScalarUDF; - use rstest::rstest; - use sedona_schema::datatypes::WKB_VIEW_GEOMETRY; - use sedona_testing::{ - compare::assert_array_equal, create::create_array, testers::ScalarUdfTester, - }; - - use super::*; - - #[test] - fn udf_metadata() { - let st_force2d: ScalarUDF = st_force2d_udf().into(); - assert_eq!(st_force2d.name(), "st_force2d"); - assert!(st_force2d.documentation().is_some()); - - let st_force3d: ScalarUDF = st_force3d_udf().into(); - assert_eq!(st_force3d.name(), "st_force3d"); - assert!(st_force3d.documentation().is_some()); - } - - #[rstest] - fn udf_2d(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { - let tester = ScalarUdfTester::new(st_force2d_udf().into(), vec![sedona_type.clone()]); - tester.assert_return_type(WKB_GEOMETRY); - - let points = create_array( - &[ - None, - Some("POINT EMPTY"), - Some("POINT Z EMPTY"), - Some("POINT (1 2)"), - Some("POINT Z (3 4 5)"), - Some("POINT ZM (8 9 10 11)"), - ], - &sedona_type, - ); - - let expected = create_array( - &[ - None, - Some("POINT EMPTY"), - Some("POINT EMPTY"), - Some("POINT (1 2)"), - Some("POINT (3 4)"), - Some("POINT (8 9)"), - ], - &WKB_GEOMETRY, - ); - - let result = tester.invoke_arrays(vec![points]).unwrap(); - assert_array_equal(&result, &expected); - } - - #[rstest] - fn udf_3d_without_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { - let tester = ScalarUdfTester::new(st_force3d_udf().into(), vec![sedona_type.clone()]); - tester.assert_return_type(WKB_GEOMETRY); - - let points = create_array( - &[ - None, - Some("POINT EMPTY"), - Some("POINT Z EMPTY"), - Some("POINT (1 2)"), - Some("POINT Z (3 4 5)"), - Some("POINT ZM (8 9 10 11)"), - ], - &sedona_type, - ); - - let expected = create_array( - &[ - None, - Some("POINT Z EMPTY"), - Some("POINT Z EMPTY"), - Some("POINT Z (1 2 0)"), - Some("POINT Z (3 4 5)"), - Some("POINT Z (8 9 10)"), - ], - &WKB_GEOMETRY, - ); - - let result = tester.invoke_arrays(vec![points]).unwrap(); - assert_array_equal(&result, &expected); - } - - #[rstest] - fn udf_3d_with_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { - let tester = ScalarUdfTester::new( - st_force3d_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 (1 2)"), - Some("POINT Z (3 4 5)"), - Some("POINT (6 7)"), - Some("POINT ZM (8 9 10 11)"), - ], - &sedona_type, - ); - let z = create_array!( - Float64, - [Some(9.0), Some(9.0), Some(9.0), Some(9.0), None, Some(9.0)] - ); - - let expected = create_array( - &[ - None, - Some("POINT Z EMPTY"), - Some("POINT Z (1 2 9)"), - Some("POINT Z (3 4 5)"), - None, - Some("POINT Z (8 9 10)"), - ], - &WKB_GEOMETRY, - ); - - let result = tester.invoke_arrays(vec![points, z]).unwrap(); - assert_array_equal(&result, &expected); - } -} +// 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 std::sync::Arc; + +use arrow_array::{builder::BinaryBuilder, Array, Float64Array}; +use arrow_schema::DataType; +use datafusion_common::{cast::as_float64_array, error::Result, DataFusionError}; +use datafusion_expr::{ + scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility, +}; +use geo_traits::Dimensions; +use sedona_expr::{ + item_crs::ItemCrsKernel, + scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, +}; +use sedona_geometry::{ + error::SedonaGeometryError, + transform::{transform, CrsTransform}, + wkb_factory::WKB_MIN_PROBABLE_BYTES, +}; +use sedona_schema::{ + datatypes::{SedonaType, WKB_GEOGRAPHY, WKB_GEOMETRY}, + matchers::ArgMatcher, +}; + +use crate::executor::WkbExecutor; + +// *** 2D ************************* + +/// ST_Force2D() scalar UDF +pub fn st_force2d_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "st_force2d", + ItemCrsKernel::wrap_impl(vec![ + Arc::new(STForce2D { + is_geography: false, + }), + Arc::new(STForce2D { is_geography: true }), + ]), + Volatility::Immutable, + Some(st_force2d_doc()), + ) +} + +fn st_force2d_doc() -> Documentation { + Documentation::builder( + DOC_SECTION_OTHER, + "Forces the geometry into a 2-dimensional model", + "ST_Force2D (geom: Geometry)", + ) + .with_argument("geom", "geometry: Input geometry") + .with_sql_example("SELECT ST_Force2D(ST_GeomFromWKT('POINT Z (1 2 3)'))") + .build() +} + +#[derive(Debug)] +struct STForce2D { + is_geography: bool, +} + +impl SedonaScalarKernel for STForce2D { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let matcher = if self.is_geography { + ArgMatcher::new(vec![ArgMatcher::is_geography()], WKB_GEOGRAPHY) + } else { + ArgMatcher::new(vec![ArgMatcher::is_geometry()], WKB_GEOMETRY) + }; + + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + let executor = WkbExecutor::new(arg_types, args); + let mut builder = BinaryBuilder::with_capacity( + executor.num_iterations(), + WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), + ); + + let trans = Force2DTransform {}; + executor.execute_wkb_void(|maybe_wkb| { + match maybe_wkb { + Some(wkb) => { + transform(wkb, &trans, &mut builder) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + builder.append_value([]); + } + _ => { + builder.append_null(); + } + } + + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + } +} + +#[derive(Debug)] +struct Force2DTransform {} + +impl CrsTransform for Force2DTransform { + fn output_dim(&self) -> Option { + Some(geo_traits::Dimensions::Xy) + } + + fn transform_coord( + &self, + _coord: &mut (f64, f64), + ) -> std::result::Result<(), SedonaGeometryError> { + Ok(()) + } +} + +// *** 3D ************************* + +/// ST_Force3D() scalar UDF +pub fn st_force3d_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "st_force3d", + ItemCrsKernel::wrap_impl(vec![ + Arc::new(STForce3D { + is_geography: false, + }), + Arc::new(STForce3D { is_geography: true }), + ]), + Volatility::Immutable, + Some(st_force3d_doc()), + ) +} + +fn st_force3d_doc() -> Documentation { + Documentation::builder( + DOC_SECTION_OTHER, + "Forces the geometry into a 3-dimensional model.", + "ST_Force3D (geom: Geometry)", + ) + .with_argument("geom", "geometry: Input geometry") + .with_argument("z", "numeric: default Z value") + .with_sql_example("SELECT ST_Force3D(ST_GeomFromWKT('POINT (1 2)'))") + .build() +} + +#[derive(Debug)] +struct STForce3D { + is_geography: bool, +} + +impl SedonaScalarKernel for STForce3D { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let matcher = if self.is_geography { + ArgMatcher::new( + vec![ + ArgMatcher::is_geography(), + ArgMatcher::optional(ArgMatcher::is_numeric()), + ], + WKB_GEOGRAPHY, + ) + } else { + ArgMatcher::new( + vec![ + ArgMatcher::is_geometry(), + ArgMatcher::optional(ArgMatcher::is_numeric()), + ], + WKB_GEOMETRY, + ) + }; + + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + let executor = WkbExecutor::new(arg_types, args); + let mut builder = BinaryBuilder::with_capacity( + executor.num_iterations(), + WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), + ); + + let z_array = match args.get(1) { + Some(arg) => arg + .cast_to(&DataType::Float64, None)? + .to_array(executor.num_iterations())?, + None => Arc::new(Float64Array::from(vec![0.0; executor.num_iterations()])), + }; + let z_array = as_float64_array(&z_array)?; + + let mut z_iter = z_array.iter(); + executor.execute_wkb_void(|maybe_wkb| { + match (maybe_wkb, z_iter.next().unwrap()) { + (Some(wkb), Some(z)) => { + let trans = Force3DTransform { z }; + transform(wkb, &trans, &mut builder) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + builder.append_value([]); + } + _ => { + builder.append_null(); + } + } + + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + } +} + +#[derive(Debug)] +struct Force3DTransform { + z: f64, +} + +impl CrsTransform for Force3DTransform { + fn output_dim(&self) -> Option { + Some(geo_traits::Dimensions::Xyz) + } + + fn transform_coord( + &self, + _coord: &mut (f64, f64), + ) -> std::result::Result<(), SedonaGeometryError> { + sedona_internal_err!("Unexpected call to transform_coord()") + } + fn transform_coord_3d( + &self, + coord: &mut (f64, f64, f64), + input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { + // If the input doesn't have Z coordinate, fill with the default value + if matches!( + input_dims, + Dimensions::Xy | Dimensions::Xym | Dimensions::Unknown(_) + ) { + coord.2 = self.z + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::create_array; + use datafusion_expr::ScalarUDF; + use rstest::rstest; + use sedona_schema::datatypes::WKB_VIEW_GEOMETRY; + use sedona_testing::{ + compare::assert_array_equal, create::create_array, testers::ScalarUdfTester, + }; + + use super::*; + + #[test] + fn udf_metadata() { + let st_force2d: ScalarUDF = st_force2d_udf().into(); + assert_eq!(st_force2d.name(), "st_force2d"); + assert!(st_force2d.documentation().is_some()); + + let st_force3d: ScalarUDF = st_force3d_udf().into(); + assert_eq!(st_force3d.name(), "st_force3d"); + assert!(st_force3d.documentation().is_some()); + } + + #[rstest] + fn udf_2d(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new(st_force2d_udf().into(), vec![sedona_type.clone()]); + tester.assert_return_type(WKB_GEOMETRY); + + let points = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT ZM (8 9 10 11)"), + ], + &sedona_type, + ); + + let expected = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT EMPTY"), + Some("POINT (1 2)"), + Some("POINT (3 4)"), + Some("POINT (8 9)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points]).unwrap(); + assert_array_equal(&result, &expected); + } + + #[rstest] + fn udf_3d_without_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new(st_force3d_udf().into(), vec![sedona_type.clone()]); + tester.assert_return_type(WKB_GEOMETRY); + + let points = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT ZM (8 9 10 11)"), + ], + &sedona_type, + ); + + let expected = create_array( + &[ + None, + Some("POINT Z EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT Z (1 2 0)"), + Some("POINT Z (3 4 5)"), + Some("POINT Z (8 9 10)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points]).unwrap(); + assert_array_equal(&result, &expected); + } + + #[rstest] + fn udf_3d_with_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new( + st_force3d_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 (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT (6 7)"), + Some("POINT ZM (8 9 10 11)"), + ], + &sedona_type, + ); + let z = create_array!( + Float64, + [Some(9.0), Some(9.0), Some(9.0), Some(9.0), None, Some(9.0)] + ); + + let expected = create_array( + &[ + None, + Some("POINT Z EMPTY"), + Some("POINT Z (1 2 9)"), + Some("POINT Z (3 4 5)"), + None, + Some("POINT Z (8 9 10)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points, z]).unwrap(); + assert_array_equal(&result, &expected); + } +} From e790168f23cb116a712fa74a0068dda76135b785 Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 20:48:49 +0900 Subject: [PATCH 14/15] Tweak error --- rust/sedona-functions/src/st_force_dim.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rust/sedona-functions/src/st_force_dim.rs b/rust/sedona-functions/src/st_force_dim.rs index cd813d63e..efddbd883 100644 --- a/rust/sedona-functions/src/st_force_dim.rs +++ b/rust/sedona-functions/src/st_force_dim.rs @@ -17,7 +17,7 @@ use std::sync::Arc; -use arrow_array::{builder::BinaryBuilder, Array, Float64Array}; +use arrow_array::{builder::BinaryBuilder, Float64Array}; use arrow_schema::DataType; use datafusion_common::{cast::as_float64_array, error::Result, DataFusionError}; use datafusion_expr::{ @@ -242,7 +242,9 @@ impl CrsTransform for Force3DTransform { &self, _coord: &mut (f64, f64), ) -> std::result::Result<(), SedonaGeometryError> { - sedona_internal_err!("Unexpected call to transform_coord()") + Err(SedonaGeometryError::Invalid( + "Unexpected call to transform_coord()".to_string(), + )) } fn transform_coord_3d( &self, From 8b6488359c674681f9b47594e748cb2829d9a172 Mon Sep 17 00:00:00 2001 From: Hiroaki Yutani Date: Fri, 13 Feb 2026 20:57:08 +0900 Subject: [PATCH 15/15] Fix tests --- rust/sedona-functions/src/st_force_dim.rs | 798 +++++++++++----------- 1 file changed, 405 insertions(+), 393 deletions(-) diff --git a/rust/sedona-functions/src/st_force_dim.rs b/rust/sedona-functions/src/st_force_dim.rs index efddbd883..14d0835ac 100644 --- a/rust/sedona-functions/src/st_force_dim.rs +++ b/rust/sedona-functions/src/st_force_dim.rs @@ -1,393 +1,405 @@ -// 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 std::sync::Arc; - -use arrow_array::{builder::BinaryBuilder, Float64Array}; -use arrow_schema::DataType; -use datafusion_common::{cast::as_float64_array, error::Result, DataFusionError}; -use datafusion_expr::{ - scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility, -}; -use geo_traits::Dimensions; -use sedona_expr::{ - item_crs::ItemCrsKernel, - scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, -}; -use sedona_geometry::{ - error::SedonaGeometryError, - transform::{transform, CrsTransform}, - wkb_factory::WKB_MIN_PROBABLE_BYTES, -}; -use sedona_schema::{ - datatypes::{SedonaType, WKB_GEOGRAPHY, WKB_GEOMETRY}, - matchers::ArgMatcher, -}; - -use crate::executor::WkbExecutor; - -// *** 2D ************************* - -/// ST_Force2D() scalar UDF -pub fn st_force2d_udf() -> SedonaScalarUDF { - SedonaScalarUDF::new( - "st_force2d", - ItemCrsKernel::wrap_impl(vec![ - Arc::new(STForce2D { - is_geography: false, - }), - Arc::new(STForce2D { is_geography: true }), - ]), - Volatility::Immutable, - Some(st_force2d_doc()), - ) -} - -fn st_force2d_doc() -> Documentation { - Documentation::builder( - DOC_SECTION_OTHER, - "Forces the geometry into a 2-dimensional model", - "ST_Force2D (geom: Geometry)", - ) - .with_argument("geom", "geometry: Input geometry") - .with_sql_example("SELECT ST_Force2D(ST_GeomFromWKT('POINT Z (1 2 3)'))") - .build() -} - -#[derive(Debug)] -struct STForce2D { - is_geography: bool, -} - -impl SedonaScalarKernel for STForce2D { - fn return_type(&self, args: &[SedonaType]) -> Result> { - let matcher = if self.is_geography { - ArgMatcher::new(vec![ArgMatcher::is_geography()], WKB_GEOGRAPHY) - } else { - ArgMatcher::new(vec![ArgMatcher::is_geometry()], WKB_GEOMETRY) - }; - - matcher.match_args(args) - } - - fn invoke_batch( - &self, - arg_types: &[SedonaType], - args: &[ColumnarValue], - ) -> Result { - let executor = WkbExecutor::new(arg_types, args); - let mut builder = BinaryBuilder::with_capacity( - executor.num_iterations(), - WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), - ); - - let trans = Force2DTransform {}; - executor.execute_wkb_void(|maybe_wkb| { - match maybe_wkb { - Some(wkb) => { - transform(wkb, &trans, &mut builder) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - builder.append_value([]); - } - _ => { - builder.append_null(); - } - } - - Ok(()) - })?; - - executor.finish(Arc::new(builder.finish())) - } -} - -#[derive(Debug)] -struct Force2DTransform {} - -impl CrsTransform for Force2DTransform { - fn output_dim(&self) -> Option { - Some(geo_traits::Dimensions::Xy) - } - - fn transform_coord( - &self, - _coord: &mut (f64, f64), - ) -> std::result::Result<(), SedonaGeometryError> { - Ok(()) - } -} - -// *** 3D ************************* - -/// ST_Force3D() scalar UDF -pub fn st_force3d_udf() -> SedonaScalarUDF { - SedonaScalarUDF::new( - "st_force3d", - ItemCrsKernel::wrap_impl(vec![ - Arc::new(STForce3D { - is_geography: false, - }), - Arc::new(STForce3D { is_geography: true }), - ]), - Volatility::Immutable, - Some(st_force3d_doc()), - ) -} - -fn st_force3d_doc() -> Documentation { - Documentation::builder( - DOC_SECTION_OTHER, - "Forces the geometry into a 3-dimensional model.", - "ST_Force3D (geom: Geometry)", - ) - .with_argument("geom", "geometry: Input geometry") - .with_argument("z", "numeric: default Z value") - .with_sql_example("SELECT ST_Force3D(ST_GeomFromWKT('POINT (1 2)'))") - .build() -} - -#[derive(Debug)] -struct STForce3D { - is_geography: bool, -} - -impl SedonaScalarKernel for STForce3D { - fn return_type(&self, args: &[SedonaType]) -> Result> { - let matcher = if self.is_geography { - ArgMatcher::new( - vec![ - ArgMatcher::is_geography(), - ArgMatcher::optional(ArgMatcher::is_numeric()), - ], - WKB_GEOGRAPHY, - ) - } else { - ArgMatcher::new( - vec![ - ArgMatcher::is_geometry(), - ArgMatcher::optional(ArgMatcher::is_numeric()), - ], - WKB_GEOMETRY, - ) - }; - - matcher.match_args(args) - } - - fn invoke_batch( - &self, - arg_types: &[SedonaType], - args: &[ColumnarValue], - ) -> Result { - let executor = WkbExecutor::new(arg_types, args); - let mut builder = BinaryBuilder::with_capacity( - executor.num_iterations(), - WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), - ); - - let z_array = match args.get(1) { - Some(arg) => arg - .cast_to(&DataType::Float64, None)? - .to_array(executor.num_iterations())?, - None => Arc::new(Float64Array::from(vec![0.0; executor.num_iterations()])), - }; - let z_array = as_float64_array(&z_array)?; - - let mut z_iter = z_array.iter(); - executor.execute_wkb_void(|maybe_wkb| { - match (maybe_wkb, z_iter.next().unwrap()) { - (Some(wkb), Some(z)) => { - let trans = Force3DTransform { z }; - transform(wkb, &trans, &mut builder) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - builder.append_value([]); - } - _ => { - builder.append_null(); - } - } - - Ok(()) - })?; - - executor.finish(Arc::new(builder.finish())) - } -} - -#[derive(Debug)] -struct Force3DTransform { - z: f64, -} - -impl CrsTransform for Force3DTransform { - fn output_dim(&self) -> Option { - Some(geo_traits::Dimensions::Xyz) - } - - fn transform_coord( - &self, - _coord: &mut (f64, f64), - ) -> std::result::Result<(), SedonaGeometryError> { - Err(SedonaGeometryError::Invalid( - "Unexpected call to transform_coord()".to_string(), - )) - } - fn transform_coord_3d( - &self, - coord: &mut (f64, f64, f64), - input_dims: Dimensions, - ) -> Result<(), SedonaGeometryError> { - // If the input doesn't have Z coordinate, fill with the default value - if matches!( - input_dims, - Dimensions::Xy | Dimensions::Xym | Dimensions::Unknown(_) - ) { - coord.2 = self.z - } - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use arrow_array::create_array; - use datafusion_expr::ScalarUDF; - use rstest::rstest; - use sedona_schema::datatypes::WKB_VIEW_GEOMETRY; - use sedona_testing::{ - compare::assert_array_equal, create::create_array, testers::ScalarUdfTester, - }; - - use super::*; - - #[test] - fn udf_metadata() { - let st_force2d: ScalarUDF = st_force2d_udf().into(); - assert_eq!(st_force2d.name(), "st_force2d"); - assert!(st_force2d.documentation().is_some()); - - let st_force3d: ScalarUDF = st_force3d_udf().into(); - assert_eq!(st_force3d.name(), "st_force3d"); - assert!(st_force3d.documentation().is_some()); - } - - #[rstest] - fn udf_2d(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { - let tester = ScalarUdfTester::new(st_force2d_udf().into(), vec![sedona_type.clone()]); - tester.assert_return_type(WKB_GEOMETRY); - - let points = create_array( - &[ - None, - Some("POINT EMPTY"), - Some("POINT Z EMPTY"), - Some("POINT (1 2)"), - Some("POINT Z (3 4 5)"), - Some("POINT ZM (8 9 10 11)"), - ], - &sedona_type, - ); - - let expected = create_array( - &[ - None, - Some("POINT EMPTY"), - Some("POINT EMPTY"), - Some("POINT (1 2)"), - Some("POINT (3 4)"), - Some("POINT (8 9)"), - ], - &WKB_GEOMETRY, - ); - - let result = tester.invoke_arrays(vec![points]).unwrap(); - assert_array_equal(&result, &expected); - } - - #[rstest] - fn udf_3d_without_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { - let tester = ScalarUdfTester::new(st_force3d_udf().into(), vec![sedona_type.clone()]); - tester.assert_return_type(WKB_GEOMETRY); - - let points = create_array( - &[ - None, - Some("POINT EMPTY"), - Some("POINT Z EMPTY"), - Some("POINT (1 2)"), - Some("POINT Z (3 4 5)"), - Some("POINT ZM (8 9 10 11)"), - ], - &sedona_type, - ); - - let expected = create_array( - &[ - None, - Some("POINT Z EMPTY"), - Some("POINT Z EMPTY"), - Some("POINT Z (1 2 0)"), - Some("POINT Z (3 4 5)"), - Some("POINT Z (8 9 10)"), - ], - &WKB_GEOMETRY, - ); - - let result = tester.invoke_arrays(vec![points]).unwrap(); - assert_array_equal(&result, &expected); - } - - #[rstest] - fn udf_3d_with_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { - let tester = ScalarUdfTester::new( - st_force3d_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 (1 2)"), - Some("POINT Z (3 4 5)"), - Some("POINT (6 7)"), - Some("POINT ZM (8 9 10 11)"), - ], - &sedona_type, - ); - let z = create_array!( - Float64, - [Some(9.0), Some(9.0), Some(9.0), Some(9.0), None, Some(9.0)] - ); - - let expected = create_array( - &[ - None, - Some("POINT Z EMPTY"), - Some("POINT Z (1 2 9)"), - Some("POINT Z (3 4 5)"), - None, - Some("POINT Z (8 9 10)"), - ], - &WKB_GEOMETRY, - ); - - let result = tester.invoke_arrays(vec![points, z]).unwrap(); - assert_array_equal(&result, &expected); - } -} +// 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 std::sync::Arc; + +use arrow_array::{builder::BinaryBuilder, Float64Array}; +use arrow_schema::DataType; +use datafusion_common::{cast::as_float64_array, error::Result, DataFusionError}; +use datafusion_expr::{ + scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation, Volatility, +}; +use geo_traits::Dimensions; +use sedona_expr::{ + item_crs::ItemCrsKernel, + scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}, +}; +use sedona_geometry::{ + error::SedonaGeometryError, + transform::{transform, CrsTransform}, + wkb_factory::WKB_MIN_PROBABLE_BYTES, +}; +use sedona_schema::{ + datatypes::{SedonaType, WKB_GEOGRAPHY, WKB_GEOMETRY}, + matchers::ArgMatcher, +}; + +use crate::executor::WkbExecutor; + +// *** 2D ************************* + +/// ST_Force2D() scalar UDF +pub fn st_force2d_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "st_force2d", + ItemCrsKernel::wrap_impl(vec![ + Arc::new(STForce2D { + is_geography: false, + }), + Arc::new(STForce2D { is_geography: true }), + ]), + Volatility::Immutable, + Some(st_force2d_doc()), + ) +} + +fn st_force2d_doc() -> Documentation { + Documentation::builder( + DOC_SECTION_OTHER, + "Forces the geometry into a 2-dimensional model", + "ST_Force2D (geom: Geometry)", + ) + .with_argument("geom", "geometry: Input geometry") + .with_sql_example("SELECT ST_Force2D(ST_GeomFromWKT('POINT Z (1 2 3)'))") + .build() +} + +#[derive(Debug)] +struct STForce2D { + is_geography: bool, +} + +impl SedonaScalarKernel for STForce2D { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let matcher = if self.is_geography { + ArgMatcher::new(vec![ArgMatcher::is_geography()], WKB_GEOGRAPHY) + } else { + ArgMatcher::new(vec![ArgMatcher::is_geometry()], WKB_GEOMETRY) + }; + + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + let executor = WkbExecutor::new(arg_types, args); + let mut builder = BinaryBuilder::with_capacity( + executor.num_iterations(), + WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), + ); + + let trans = Force2DTransform {}; + executor.execute_wkb_void(|maybe_wkb| { + match maybe_wkb { + Some(wkb) => { + transform(wkb, &trans, &mut builder) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + builder.append_value([]); + } + _ => { + builder.append_null(); + } + } + + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + } +} + +#[derive(Debug)] +struct Force2DTransform {} + +impl CrsTransform for Force2DTransform { + fn output_dim(&self) -> Option { + Some(geo_traits::Dimensions::Xy) + } + + fn transform_coord( + &self, + _coord: &mut (f64, f64), + ) -> std::result::Result<(), SedonaGeometryError> { + Ok(()) + } +} + +// *** 3D ************************* + +/// ST_Force3D() scalar UDF +pub fn st_force3d_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "st_force3d", + ItemCrsKernel::wrap_impl(vec![ + Arc::new(STForce3D { + is_geography: false, + }), + Arc::new(STForce3D { is_geography: true }), + ]), + Volatility::Immutable, + Some(st_force3d_doc()), + ) +} + +fn st_force3d_doc() -> Documentation { + Documentation::builder( + DOC_SECTION_OTHER, + "Forces the geometry into a 3-dimensional model.", + "ST_Force3D (geom: Geometry)", + ) + .with_argument("geom", "geometry: Input geometry") + .with_argument("z", "numeric: default Z value") + .with_sql_example("SELECT ST_Force3D(ST_GeomFromWKT('POINT (1 2)'))") + .build() +} + +#[derive(Debug)] +struct STForce3D { + is_geography: bool, +} + +impl SedonaScalarKernel for STForce3D { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let matcher = if self.is_geography { + ArgMatcher::new( + vec![ + ArgMatcher::is_geography(), + ArgMatcher::optional(ArgMatcher::is_numeric()), + ], + WKB_GEOGRAPHY, + ) + } else { + ArgMatcher::new( + vec![ + ArgMatcher::is_geometry(), + ArgMatcher::optional(ArgMatcher::is_numeric()), + ], + WKB_GEOMETRY, + ) + }; + + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + let executor = WkbExecutor::new(arg_types, args); + let mut builder = BinaryBuilder::with_capacity( + executor.num_iterations(), + WKB_MIN_PROBABLE_BYTES * executor.num_iterations(), + ); + + let z_array = match args.get(1) { + Some(arg) => arg + .cast_to(&DataType::Float64, None)? + .to_array(executor.num_iterations())?, + None => Arc::new(Float64Array::from(vec![0.0; executor.num_iterations()])), + }; + let z_array = as_float64_array(&z_array)?; + + let mut z_iter = z_array.iter(); + executor.execute_wkb_void(|maybe_wkb| { + match (maybe_wkb, z_iter.next().unwrap()) { + (Some(wkb), Some(z)) => { + let trans = Force3DTransform { z }; + transform(wkb, &trans, &mut builder) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + builder.append_value([]); + } + _ => { + builder.append_null(); + } + } + + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + } +} + +#[derive(Debug)] +struct Force3DTransform { + z: f64, +} + +impl CrsTransform for Force3DTransform { + fn output_dim(&self) -> Option { + Some(geo_traits::Dimensions::Xyz) + } + + fn transform_coord( + &self, + _coord: &mut (f64, f64), + ) -> std::result::Result<(), SedonaGeometryError> { + Err(SedonaGeometryError::Invalid( + "Unexpected call to transform_coord()".to_string(), + )) + } + fn transform_coord_3d( + &self, + coord: &mut (f64, f64, f64), + input_dims: Dimensions, + ) -> Result<(), SedonaGeometryError> { + // If the input doesn't have Z coordinate, fill with the default value + if matches!( + input_dims, + Dimensions::Xy | Dimensions::Xym | Dimensions::Unknown(_) + ) { + coord.2 = self.z + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::create_array; + use datafusion_expr::ScalarUDF; + use rstest::rstest; + use sedona_schema::datatypes::WKB_VIEW_GEOMETRY; + use sedona_testing::{ + compare::assert_array_equal, create::create_array, testers::ScalarUdfTester, + }; + + use super::*; + + #[test] + fn udf_metadata() { + let st_force2d: ScalarUDF = st_force2d_udf().into(); + assert_eq!(st_force2d.name(), "st_force2d"); + assert!(st_force2d.documentation().is_some()); + + let st_force3d: ScalarUDF = st_force3d_udf().into(); + assert_eq!(st_force3d.name(), "st_force3d"); + assert!(st_force3d.documentation().is_some()); + } + + #[rstest] + fn udf_2d(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new(st_force2d_udf().into(), vec![sedona_type.clone()]); + tester.assert_return_type(WKB_GEOMETRY); + + let points = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT ZM (8 9 10 11)"), + ], + &sedona_type, + ); + + let expected = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT EMPTY"), + Some("POINT (1 2)"), + Some("POINT (3 4)"), + Some("POINT (8 9)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points]).unwrap(); + assert_array_equal(&result, &expected); + } + + #[rstest] + fn udf_3d_without_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new(st_force3d_udf().into(), vec![sedona_type.clone()]); + tester.assert_return_type(WKB_GEOMETRY); + + let points = create_array( + &[ + None, + Some("POINT EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT M (6 7 8)"), + Some("POINT ZM (9 10 11 12)"), + ], + &sedona_type, + ); + + let expected = create_array( + &[ + None, + Some("POINT Z EMPTY"), + Some("POINT Z EMPTY"), + Some("POINT Z (1 2 0)"), + Some("POINT Z (3 4 5)"), + Some("POINT Z (6 7 0)"), + Some("POINT Z (9 10 11)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points]).unwrap(); + assert_array_equal(&result, &expected); + } + + #[rstest] + fn udf_3d_with_z(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType) { + let tester = ScalarUdfTester::new( + st_force3d_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 (1 2)"), + Some("POINT (1 2)"), + Some("POINT Z (3 4 5)"), + Some("POINT M (6 7 8)"), + Some("POINT ZM (9 10 11 12)"), + ], + &sedona_type, + ); + let z = create_array!( + Float64, + [ + Some(9.0), + Some(9.0), + Some(9.0), + None, + Some(9.0), + Some(9.0), + Some(9.0) + ] + ); + + let expected = create_array( + &[ + None, + Some("POINT Z EMPTY"), + Some("POINT Z (1 2 9)"), + None, + Some("POINT Z (3 4 5)"), + Some("POINT Z (6 7 9)"), + Some("POINT Z (9 10 11)"), + ], + &WKB_GEOMETRY, + ); + + let result = tester.invoke_arrays(vec![points, z]).unwrap(); + assert_array_equal(&result, &expected); + } +}