This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new b38ad859 feat(rust/sedona-functions): make ST_Translate accept deltaZ
arg (#524)
b38ad859 is described below
commit b38ad85968a0ffc029f92dc766140a41ecae58a8
Author: Hiroaki Yutani <[email protected]>
AuthorDate: Tue Jan 20 00:14:49 2026 +0900
feat(rust/sedona-functions): make ST_Translate accept deltaZ arg (#524)
---
python/sedonadb/tests/functions/test_transforms.py | 68 ++++++
rust/sedona-functions/src/st_translate.rs | 237 +++++++++++++++++++--
2 files changed, 283 insertions(+), 22 deletions(-)
diff --git a/python/sedonadb/tests/functions/test_transforms.py
b/python/sedonadb/tests/functions/test_transforms.py
index 5ae0d20d..4f82c5dd 100644
--- a/python/sedonadb/tests/functions/test_transforms.py
+++ b/python/sedonadb/tests/functions/test_transforms.py
@@ -192,3 +192,71 @@ def test_st_translate(eng, geom, dx, dy, expected):
f"SELECT ST_Translate({geom_or_null(geom)}, {val_or_null(dx)},
{val_or_null(dy)})",
expected,
)
+
+
[email protected]("eng", [SedonaDB, PostGIS])
[email protected](
+ ("geom", "dx", "dy", "dz", "expected"),
+ [
+ # Nulls
+ (None, None, None, None, None),
+ (None, 1.0, 2.0, 3.0, None),
+ ("POINT Z (0 1 2)", None, 2.0, 3.0, None),
+ ("POINT Z (0 1 2)", 1.0, None, 3.0, None),
+ ("POINT Z (0 1 2)", 1.0, 2.0, None, None),
+ ("POINT Z (0 1 2)", 1.0, 2.0, 3.0, "POINT Z (1 3 5)"), # Positives
+ ("POINT Z (0 1 2)", -1.0, -2.0, -3.0, "POINT Z (-1 -1 -1)"), #
Negatives
+ ("POINT Z (0 1 2)", 0.0, 0.0, 0.0, "POINT Z (0 1 2)"), # Zeroes
+ ("POINT Z (0 1 2)", 1, 2, 3, "POINT Z (1 3 5)"), # Integers
+ ("POINT (0 1)", 1.0, 2.0, 3.0, "POINT (1 3)"), # 2D
+ ("POINT M (0 1 2)", 1.0, 2.0, 3.0, "POINT M (1 3 2)"), # M
+ ("POINT ZM (0 1 2 3)", 1.0, 2.0, 3.0, "POINT ZM (1 3 5 3)"), # ZM
+ # Not points
+ ("LINESTRING Z (0 1 2, 2 3 4)", 1.0, 2.0, 3.0, "LINESTRING Z (1 3 5, 3
5 7)"),
+ (
+ "POLYGON Z ((0 0 0, 1 0 2, 0 1 2, 0 0 0))",
+ 1.0,
+ 2.0,
+ 3.0,
+ "POLYGON Z ((1 2 3, 2 2 5, 1 3 5, 1 2 3))",
+ ),
+ ("MULTIPOINT Z (0 1 2, 2 3 4)", 1.0, 2.0, 3.0, "MULTIPOINT Z (1 3 5, 3
5 7)"),
+ (
+ "MULTILINESTRING Z ((0 1 2, 2 3 4))",
+ 1.0,
+ 2.0,
+ 3.0,
+ "MULTILINESTRING Z ((1 3 5, 3 5 7))",
+ ),
+ (
+ "MULTIPOLYGON Z (((0 0 0, 1 0 2, 0 1 2, 0 0 0)))",
+ 1.0,
+ 2.0,
+ 3.0,
+ "MULTIPOLYGON Z (((1 2 3, 2 2 5, 1 3 5, 1 2 3)))",
+ ),
+ (
+ "GEOMETRYCOLLECTION Z (POINT Z (0 1 2))",
+ 1.0,
+ 2.0,
+ 3.0,
+ "GEOMETRYCOLLECTION Z (POINT Z (1 3 5))",
+ ),
+ # WKT output of geoarrow-c is causing this (both correctly output
+ # empties)
+ ("POINT EMPTY", 1.0, 2.0, 3.0, "POINT (nan nan)"),
+ ("POINT Z EMPTY", 1.0, 2.0, 3.0, "POINT Z (nan nan nan)"),
+ ("LINESTRING EMPTY", 1.0, 2.0, 3.0, "LINESTRING EMPTY"),
+ ("POLYGON EMPTY", 1.0, 2.0, 3.0, "POLYGON EMPTY"),
+ ("MULTIPOINT EMPTY", 1.0, 2.0, 3.0, "MULTIPOINT EMPTY"),
+ ("MULTILINESTRING EMPTY", 1.0, 2.0, 3.0, "MULTILINESTRING EMPTY"),
+ ("MULTIPOLYGON EMPTY", 1.0, 2.0, 3.0, "MULTIPOLYGON EMPTY"),
+ ("GEOMETRYCOLLECTION EMPTY", 1.0, 2.0, 3.0, "GEOMETRYCOLLECTION
EMPTY"),
+ ],
+)
+def test_st_translate_3d(eng, geom, dx, dy, dz, expected):
+ eng = eng.create_or_skip()
+ eng.assert_query_result(
+ f"SELECT ST_Translate({geom_or_null(geom)}, {val_or_null(dx)},
{val_or_null(dy)}, {val_or_null(dz)})",
+ expected,
+ )
diff --git a/rust/sedona-functions/src/st_translate.rs
b/rust/sedona-functions/src/st_translate.rs
index b18c129f..647e29cf 100644
--- a/rust/sedona-functions/src/st_translate.rs
+++ b/rust/sedona-functions/src/st_translate.rs
@@ -14,13 +14,14 @@
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
-use arrow_array::builder::BinaryBuilder;
+use arrow_array::{builder::BinaryBuilder, types::Float64Type, Array,
PrimitiveArray};
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 sedona_common::sedona_internal_err;
use sedona_expr::{
item_crs::ItemCrsKernel,
scalar_udf::{SedonaScalarKernel, SedonaScalarUDF},
@@ -34,7 +35,7 @@ use sedona_schema::{
datatypes::{SedonaType, WKB_GEOMETRY},
matchers::ArgMatcher,
};
-use std::{iter::zip, sync::Arc};
+use std::sync::Arc;
use crate::executor::WkbExecutor;
@@ -42,7 +43,10 @@ use crate::executor::WkbExecutor;
pub fn st_translate_udf() -> SedonaScalarUDF {
SedonaScalarUDF::new(
"st_translate",
- ItemCrsKernel::wrap_impl(vec![Arc::new(STTranslate)]),
+ ItemCrsKernel::wrap_impl(vec![
+ Arc::new(STTranslate { is_3d: true }),
+ Arc::new(STTranslate { is_3d: false }),
+ ]),
Volatility::Immutable,
Some(st_translate_doc()),
)
@@ -62,18 +66,27 @@ fn st_translate_doc() -> Documentation {
}
#[derive(Debug)]
-struct STTranslate;
+struct STTranslate {
+ is_3d: bool,
+}
impl SedonaScalarKernel for STTranslate {
fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
- let matcher = ArgMatcher::new(
+ let matchers = if self.is_3d {
vec![
ArgMatcher::is_geometry(),
ArgMatcher::is_numeric(),
ArgMatcher::is_numeric(),
- ],
- WKB_GEOMETRY,
- );
+ ArgMatcher::is_numeric(),
+ ]
+ } else {
+ vec![
+ ArgMatcher::is_geometry(),
+ ArgMatcher::is_numeric(),
+ ArgMatcher::is_numeric(),
+ ]
+ };
+ let matcher = ArgMatcher::new(matchers, WKB_GEOMETRY);
matcher.match_args(args)
}
@@ -89,21 +102,40 @@ impl SedonaScalarKernel for STTranslate {
WKB_MIN_PROBABLE_BYTES * executor.num_iterations(),
);
- let deltax = args[1]
- .cast_to(&DataType::Float64, None)?
- .to_array(executor.num_iterations())?;
- let deltay = args[2]
- .cast_to(&DataType::Float64, None)?
- .to_array(executor.num_iterations())?;
- let deltax_array = as_float64_array(&deltax)?;
- let deltay_array = as_float64_array(&deltay)?;
- let mut delta_iter = zip(deltax_array, deltay_array);
+ let array_args = args[1..]
+ .iter()
+ .map(|arg| {
+ arg.cast_to(&DataType::Float64, None)?
+ .to_array(executor.num_iterations())
+ })
+ .collect::<Result<Vec<Arc<dyn arrow_array::Array>>>>()?;
+
+ let deltax_array = as_float64_array(&array_args[0])?;
+ let deltay_array = as_float64_array(&array_args[1])?;
+
+ let mut deltas = if self.is_3d {
+ if args.len() != 4 {
+ return sedona_internal_err!("Invalid number of arguments are
passed");
+ }
+
+ let deltaz_array = as_float64_array(&array_args[2])?;
+ Deltas::new(deltax_array, deltay_array, Some(deltaz_array))
+ } else {
+ if args.len() != 3 {
+ return sedona_internal_err!("Invalid number of arguments are
passed");
+ }
+
+ Deltas::new(deltax_array, deltay_array, None)
+ };
executor.execute_wkb_void(|maybe_wkb| {
- let (deltax, deltay) = delta_iter.next().unwrap();
- match (maybe_wkb, deltax, deltay) {
- (Some(wkb), Some(deltax), Some(deltay)) => {
- let trans = Translate { deltax, deltay };
+ match (maybe_wkb, deltas.next().unwrap()) {
+ (Some(wkb), Some((deltax, deltay, deltaz))) => {
+ let trans = Translate {
+ deltax,
+ deltay,
+ deltaz,
+ };
transform(wkb, &trans, &mut builder)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
builder.append_value([]);
@@ -120,10 +152,76 @@ impl SedonaScalarKernel for STTranslate {
}
}
+#[derive(Debug)]
+struct Deltas<'a> {
+ index: usize,
+ x: &'a PrimitiveArray<Float64Type>,
+ y: &'a PrimitiveArray<Float64Type>,
+ z: Option<&'a PrimitiveArray<Float64Type>>,
+ no_null: bool,
+}
+
+impl<'a> Deltas<'a> {
+ fn new(
+ x: &'a PrimitiveArray<Float64Type>,
+ y: &'a PrimitiveArray<Float64Type>,
+ z: Option<&'a PrimitiveArray<Float64Type>>,
+ ) -> Self {
+ let no_null = x.null_count() == 0
+ && y.null_count() == 0
+ && match z {
+ Some(z) => z.null_count() == 0,
+ None => true,
+ };
+
+ Self {
+ index: 0,
+ x,
+ y,
+ z,
+ no_null,
+ }
+ }
+ fn is_null(&self, i: usize) -> bool {
+ if self.no_null {
+ return false;
+ }
+
+ self.x.is_null(i)
+ || self.y.is_null(i)
+ || match self.z {
+ Some(z) => z.is_null(i),
+ None => false,
+ }
+ }
+}
+
+impl<'a> Iterator for Deltas<'a> {
+ type Item = Option<(f64, f64, f64)>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ let i = self.index;
+ self.index += 1;
+
+ if self.is_null(i) {
+ return Some(None);
+ }
+
+ let x = self.x.value(i);
+ let y = self.y.value(i);
+ let z = match self.z {
+ Some(z) => z.value(i),
+ None => 0.0,
+ };
+ Some(Some((x, y, z)))
+ }
+}
+
#[derive(Debug)]
struct Translate {
deltax: f64,
deltay: f64,
+ deltaz: f64,
}
impl CrsTransform for Translate {
@@ -132,6 +230,16 @@ impl CrsTransform for Translate {
coord.1 += self.deltay;
Ok(())
}
+
+ fn transform_coord_3d(
+ &self,
+ coord: &mut (f64, f64, f64),
+ ) -> std::result::Result<(), SedonaGeometryError> {
+ coord.0 += self.deltax;
+ coord.1 += self.deltay;
+ coord.2 += self.deltaz;
+ Ok(())
+ }
}
#[cfg(test)]
@@ -155,7 +263,7 @@ mod tests {
}
#[rstest]
- fn udf(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type: SedonaType)
{
+ fn udf_2d(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type:
SedonaType) {
let tester = ScalarUdfTester::new(
st_translate_udf().into(),
vec![
@@ -225,6 +333,91 @@ mod tests {
assert_array_equal(&result, &expected);
}
+ #[rstest]
+ fn udf_3d(#[values(WKB_GEOMETRY, WKB_VIEW_GEOMETRY)] sedona_type:
SedonaType) {
+ let tester = ScalarUdfTester::new(
+ st_translate_udf().into(),
+ vec![
+ sedona_type.clone(),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ ],
+ );
+ tester.assert_return_type(WKB_GEOMETRY);
+
+ let points = create_array(
+ &[
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT EMPTY"),
+ Some("POINT EMPTY"),
+ Some("POINT EMPTY"),
+ Some("POINT Z EMPTY"),
+ Some("POINT (0 1)"),
+ Some("POINT Z (4 5 6)"),
+ ],
+ &sedona_type,
+ );
+
+ let dx = create_array!(
+ Float64,
+ [
+ Some(1.0),
+ None,
+ Some(1.0),
+ Some(1.0),
+ Some(1.0),
+ Some(1.0),
+ Some(1.0),
+ Some(1.0)
+ ]
+ );
+ let dy = create_array!(
+ Float64,
+ [
+ Some(2.0),
+ Some(2.0),
+ None,
+ Some(2.0),
+ Some(2.0),
+ Some(2.0),
+ Some(2.0),
+ Some(2.0)
+ ]
+ );
+ let dz = create_array!(
+ Float64,
+ [
+ Some(3.0),
+ Some(3.0),
+ Some(3.0),
+ None,
+ Some(3.0),
+ Some(3.0),
+ Some(3.0),
+ Some(3.0)
+ ]
+ );
+
+ let expected = create_array(
+ &[
+ None,
+ None,
+ None,
+ None,
+ Some("POINT EMPTY"),
+ Some("POINT Z EMPTY"),
+ Some("POINT (1 3)"),
+ Some("POINT Z (5 7 9)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+
+ let result = tester.invoke_arrays(vec![points, dx, dy, dz]).unwrap();
+ assert_array_equal(&result, &expected);
+ }
+
#[rstest]
fn udf_invoke_item_crs(#[values(WKB_GEOMETRY_ITEM_CRS.clone())]
sedona_type: SedonaType) {
let tester = ScalarUdfTester::new(
