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 b64611ec feat(rust/sedona-functions): Add item SRID support to
geometry constructors (#574)
b64611ec is described below
commit b64611ecf7465d1c5a0564d08208dbe7483a3e60
Author: Dewey Dunnington <[email protected]>
AuthorDate: Mon Feb 9 16:42:34 2026 -0600
feat(rust/sedona-functions): Add item SRID support to geometry constructors
(#574)
Co-authored-by: Copilot <[email protected]>
---
rust/sedona-functions/src/st_geomfromwkb.rs | 79 +++++++++++++++++++++---
rust/sedona-functions/src/st_geomfromwkt.rs | 63 ++++++++++++++++++-
rust/sedona-functions/src/st_point.rs | 59 ++++++++++++++++++
rust/sedona-functions/src/st_setsrid.rs | 95 +++++++++++++++++++++--------
4 files changed, 262 insertions(+), 34 deletions(-)
diff --git a/rust/sedona-functions/src/st_geomfromwkb.rs
b/rust/sedona-functions/src/st_geomfromwkb.rs
index b218747c..38c4239b 100644
--- a/rust/sedona-functions/src/st_geomfromwkb.rs
+++ b/rust/sedona-functions/src/st_geomfromwkb.rs
@@ -27,18 +27,20 @@ use sedona_schema::{
matchers::ArgMatcher,
};
-use crate::executor::WkbExecutor;
+use crate::{executor::WkbExecutor, st_setsrid::SRIDifiedKernel};
/// ST_GeomFromWKB() scalar UDF implementation
///
/// An implementation of WKB reading using GeoRust's wkb crate.
pub fn st_geomfromwkb_udf() -> SedonaScalarUDF {
+ let kernel = Arc::new(STGeomFromWKB {
+ validate: true,
+ out_type: WKB_VIEW_GEOMETRY,
+ });
+ let sridified_kernel = Arc::new(SRIDifiedKernel::new(kernel.clone()));
SedonaScalarUDF::new(
"st_geomfromwkb",
- vec![Arc::new(STGeomFromWKB {
- validate: true,
- out_type: WKB_VIEW_GEOMETRY,
- })],
+ vec![sridified_kernel, kernel],
Volatility::Immutable,
Some(doc("ST_GeomFromWKB", "Geometry")),
)
@@ -155,14 +157,13 @@ impl SedonaScalarKernel for STGeomFromWKB {
#[cfg(test)]
mod tests {
- use arrow_array::{ArrayRef, BinaryArray, BinaryViewArray};
+ use arrow_array::{create_array, ArrayRef, BinaryArray, BinaryViewArray};
use datafusion_common::scalar::ScalarValue;
use datafusion_expr::ScalarUDF;
use rstest::rstest;
use sedona_testing::{
compare::{assert_array_equal, assert_scalar_equal},
- create::create_array,
- create::create_scalar,
+ create::{create_array, create_array_item_crs, create_array_storage,
create_scalar},
testers::ScalarUdfTester,
};
@@ -306,6 +307,68 @@ mod tests {
}
}
+ #[test]
+ fn udf_invoke_with_array_crs() {
+ let udf = st_geomfromwkb_udf();
+ let tester = ScalarUdfTester::new(
+ udf.into(),
+ vec![
+ SedonaType::Arrow(DataType::Binary),
+ SedonaType::Arrow(DataType::Utf8),
+ ],
+ );
+
+ let return_type = tester.return_type().unwrap();
+ assert_eq!(
+ return_type,
+ SedonaType::new_item_crs(&WKB_VIEW_GEOMETRY).unwrap()
+ );
+
+ let sedona_type = WKB_GEOMETRY;
+ let wkb_array = create_array_storage(
+ &[
+ Some("POINT (0 1)"),
+ Some("POINT (2 3)"),
+ Some("POINT (4 5)"),
+ Some("POINT (6 7)"),
+ None,
+ ],
+ &sedona_type,
+ );
+ let crs_array = create_array!(
+ Utf8,
+ [
+ Some("EPSG:4326"),
+ Some("EPSG:3857"),
+ Some("EPSG:3857"),
+ Some("0"),
+ None
+ ]
+ ) as ArrayRef;
+
+ let result = tester.invoke_arrays(vec![wkb_array, crs_array]).unwrap();
+ assert_eq!(
+ &result,
+ &create_array_item_crs(
+ &[
+ Some("POINT (0 1)"),
+ Some("POINT (2 3)"),
+ Some("POINT (4 5)"),
+ Some("POINT (6 7)"),
+ None
+ ],
+ [
+ Some("OGC:CRS84"),
+ Some("EPSG:3857"),
+ Some("EPSG:3857"),
+ None,
+ None
+ ],
+ &WKB_VIEW_GEOMETRY
+ )
+ );
+ }
+
#[test]
fn geog() {
let udf = st_geogfromwkb_udf();
diff --git a/rust/sedona-functions/src/st_geomfromwkt.rs
b/rust/sedona-functions/src/st_geomfromwkt.rs
index 01eef0fe..4bb5655f 100644
--- a/rust/sedona-functions/src/st_geomfromwkt.rs
+++ b/rust/sedona-functions/src/st_geomfromwkt.rs
@@ -276,7 +276,7 @@ fn invoke_scalar_with_srid(
#[cfg(test)]
mod tests {
- use arrow_array::ArrayRef;
+ use arrow_array::{create_array, ArrayRef};
use arrow_schema::DataType;
use datafusion_common::scalar::ScalarValue;
use datafusion_expr::{Literal, ScalarUDF};
@@ -406,6 +406,67 @@ mod tests {
assert_array_equal(&tester.invoke_array(array_in).unwrap(), &expected);
}
+ #[test]
+ fn udf_invoke_with_array_crs() {
+ let udf = st_geomfromwkt_udf();
+ let tester = ScalarUdfTester::new(
+ udf.into(),
+ vec![
+ SedonaType::Arrow(DataType::Utf8),
+ SedonaType::Arrow(DataType::Utf8),
+ ],
+ );
+
+ let return_type = tester.return_type().unwrap();
+ assert_eq!(
+ return_type,
+ SedonaType::new_item_crs(&WKB_GEOMETRY).unwrap()
+ );
+
+ let wkt_array: ArrayRef = create_array!(
+ Utf8,
+ [
+ Some("POINT (0 1)"),
+ Some("POINT (2 3)"),
+ Some("POINT (4 5)"),
+ Some("POINT (6 7)"),
+ None
+ ]
+ );
+ let crs_array = create_array!(
+ Utf8,
+ [
+ Some("EPSG:4326"),
+ Some("EPSG:3857"),
+ Some("EPSG:3857"),
+ Some("0"),
+ None
+ ]
+ ) as ArrayRef;
+
+ let result = tester.invoke_arrays(vec![wkt_array, crs_array]).unwrap();
+ assert_eq!(
+ &result,
+ &create_array_item_crs(
+ &[
+ Some("POINT (0 1)"),
+ Some("POINT (2 3)"),
+ Some("POINT (4 5)"),
+ Some("POINT (6 7)"),
+ None
+ ],
+ [
+ Some("OGC:CRS84"),
+ Some("EPSG:3857"),
+ Some("EPSG:3857"),
+ None,
+ None
+ ],
+ &WKB_GEOMETRY
+ )
+ );
+ }
+
#[test]
fn invalid_wkt() {
let udf = st_geomfromwkt_udf();
diff --git a/rust/sedona-functions/src/st_point.rs
b/rust/sedona-functions/src/st_point.rs
index aedf5e34..de3e279e 100644
--- a/rust/sedona-functions/src/st_point.rs
+++ b/rust/sedona-functions/src/st_point.rs
@@ -170,6 +170,7 @@ mod tests {
use sedona_schema::crs::lnglat;
use sedona_schema::datatypes::Edges;
use sedona_testing::compare::assert_array_equal;
+ use sedona_testing::create::create_array_item_crs;
use sedona_testing::{create::create_array, testers::ScalarUdfTester};
use super::*;
@@ -284,6 +285,64 @@ mod tests {
tester.assert_scalar_result_equals_with_return_type(result, "POINT (1
2)", return_type);
}
+ #[test]
+ fn udf_invoke_with_array_crs() {
+ let udf = st_point_udf();
+ let tester = ScalarUdfTester::new(
+ udf.into(),
+ vec![
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Utf8),
+ ],
+ );
+
+ let return_type = tester.return_type().unwrap();
+ assert_eq!(
+ return_type,
+ SedonaType::new_item_crs(&WKB_GEOMETRY).unwrap()
+ );
+
+ let x_array: ArrayRef =
+ create_array!(Float64, [Some(0.0), Some(2.0), Some(4.0),
Some(6.0), None]);
+ let y_array: ArrayRef =
+ create_array!(Float64, [Some(1.0), Some(3.0), Some(5.0),
Some(7.0), None]);
+ let crs_array = create_array!(
+ Utf8,
+ [
+ Some("EPSG:4326"),
+ Some("EPSG:3857"),
+ Some("EPSG:3857"),
+ Some("0"),
+ None
+ ]
+ ) as ArrayRef;
+
+ let result = tester
+ .invoke_arrays(vec![x_array, y_array, crs_array])
+ .unwrap();
+ assert_eq!(
+ &result,
+ &create_array_item_crs(
+ &[
+ Some("POINT (0 1)"),
+ Some("POINT (2 3)"),
+ Some("POINT (4 5)"),
+ Some("POINT (6 7)"),
+ None
+ ],
+ [
+ Some("OGC:CRS84"),
+ Some("EPSG:3857"),
+ Some("EPSG:3857"),
+ None,
+ None
+ ],
+ &WKB_GEOMETRY
+ )
+ );
+ }
+
#[test]
fn udf_invoke_with_invalid_srid() {
let udf = st_point_udf();
diff --git a/rust/sedona-functions/src/st_setsrid.rs
b/rust/sedona-functions/src/st_setsrid.rs
index 0d0e09fb..78fe9ae6 100644
--- a/rust/sedona-functions/src/st_setsrid.rs
+++ b/rust/sedona-functions/src/st_setsrid.rs
@@ -36,7 +36,10 @@ use datafusion_expr::{
};
use sedona_common::sedona_internal_err;
use sedona_expr::{
- item_crs::{make_item_crs, parse_item_crs_arg,
parse_item_crs_arg_type_strip_crs},
+ item_crs::{
+ make_item_crs, parse_item_crs_arg, parse_item_crs_arg_type,
+ parse_item_crs_arg_type_strip_crs,
+ },
scalar_udf::{ScalarKernelRef, SedonaScalarKernel, SedonaScalarUDF},
};
use sedona_geometry::transform::CrsEngine;
@@ -317,38 +320,61 @@ impl SedonaScalarKernel for SRIDifiedKernel {
None => return Ok(None),
};
- let crs = match scalar_args[orig_args_len] {
- Some(crs) => crs,
- None => return Ok(None),
- };
- let new_crs = match crs.cast_to(&DataType::Utf8) {
- Ok(ScalarValue::Utf8(Some(crs))) => {
- if crs == "0" {
- None
- } else {
- validate_crs(&crs, None)?;
- deserialize_crs(&crs)?
+ // If we have a scalar CRS, the output type has a type-level CRS
+ if let Some(scalar_crs) = scalar_args[orig_args_len] {
+ let new_crs = match scalar_crs.cast_to(&DataType::Utf8) {
+ Ok(ScalarValue::Utf8(Some(crs))) => {
+ if crs == "0" {
+ None
+ } else {
+ validate_crs(&crs, None)?;
+ deserialize_crs(&crs)?
+ }
+ }
+ Ok(ScalarValue::Utf8(None)) => None,
+ Ok(_) | Err(_) => {
+ return sedona_internal_err!("Can't cast Crs {scalar_crs:?}
to Utf8")
+ }
+ };
+
+ match &mut inner_result {
+ SedonaType::Wkb(_, crs) => *crs = new_crs,
+ SedonaType::WkbView(_, crs) => *crs = new_crs,
+ _ => {
+ return sedona_internal_err!("Return type must be Wkb or
WkbView");
}
}
- Ok(ScalarValue::Utf8(None)) => None,
- Ok(_) | Err(_) => return sedona_internal_err!("Can't cast Crs
{crs:?} to Utf8"),
- };
- match &mut inner_result {
- SedonaType::Wkb(_, crs) => *crs = new_crs,
- SedonaType::WkbView(_, crs) => *crs = new_crs,
- _ => {
- return sedona_internal_err!("Return type must be Wkb or
WkbView");
+ Ok(Some(inner_result))
+ } else {
+ // If we have a column CRS, the output has an item level CRS
+
+ // Assert that the output type had a Crs of None. If the inner
kernel returned
+ // a specific output CRS it is likely that the SRIDified kernel is
not appropriate.
+ match &inner_result {
+ SedonaType::Wkb(_, crs) | SedonaType::WkbView(_, crs) => {
+ if crs.is_some() {
+ return sedona_internal_err!(
+ "Return type of SRIDifiedKernel inner specified an
explicit CRS"
+ );
+ }
+ }
+ _ => {
+ return sedona_internal_err!("Return type must be Wkb or
WkbView");
+ }
}
- }
- Ok(Some(inner_result))
+ Ok(Some(SedonaType::new_item_crs(&inner_result)?))
+ }
}
- fn invoke_batch(
+ fn invoke_batch_from_args(
&self,
arg_types: &[SedonaType],
args: &[ColumnarValue],
+ return_type: &SedonaType,
+ _num_rows: usize,
+ _config_options: Option<&ConfigOptions>,
) -> Result<ColumnarValue> {
let orig_args_len = arg_types.len() - 1;
let orig_arg_types = &arg_types[..orig_args_len];
@@ -358,8 +384,10 @@ impl SedonaScalarKernel for SRIDifiedKernel {
// Note that, this behavior is different from PostGIS.
let result = self.inner.invoke_batch(orig_arg_types, orig_args)?;
- // If the specified SRID is NULL, the result is also NULL.
+ // If the CRS input is a scalar, we can return the inner result as-is
except
+ // for the NULL case.
if let ColumnarValue::Scalar(sc) = &args[orig_args_len] {
+ // If the specified SRID is NULL, the result is also NULL.
if sc.is_null() {
// Create the same length of NULLs as the original result.
let len = match &result {
@@ -374,9 +402,26 @@ impl SedonaScalarKernel for SRIDifiedKernel {
let new_array = builder.finish();
return Ok(ColumnarValue::Array(Arc::new(new_array)));
}
+
+ Ok(result)
+ } else {
+ let (item_type, _) = parse_item_crs_arg_type(return_type)?;
+ let normalized_crs_value =
normalize_crs_array(&args[orig_args_len], None)?;
+ make_item_crs(
+ &item_type,
+ result,
+ &ColumnarValue::Array(normalized_crs_value),
+ crs_input_nulls(&args[orig_args_len]),
+ )
}
+ }
- Ok(result)
+ fn invoke_batch(
+ &self,
+ _arg_types: &[SedonaType],
+ _args: &[ColumnarValue],
+ ) -> Result<ColumnarValue> {
+ sedona_internal_err!("Should not be called because
invoke_batch_from_args() is implemented")
}
fn return_type(&self, _args: &[SedonaType]) -> Result<Option<SedonaType>> {
