feat(rust/sedona-raster-functions): Add RS_RasterToWorldCoord as Point (#405)
diff --git a/rust/sedona-raster-functions/benches/native-raster-functions.rs b/rust/sedona-raster-functions/benches/native-raster-functions.rs
index ca7a1f4..8c3df52 100644
--- a/rust/sedona-raster-functions/benches/native-raster-functions.rs
+++ b/rust/sedona-raster-functions/benches/native-raster-functions.rs
@@ -25,6 +25,13 @@
c,
&f,
"native-raster",
+ "rs_rastertoworldcoord",
+ BenchmarkArgs::ArrayScalarScalar(Raster(64, 64), Int32(0, 63), Int32(0, 63)),
+ );
+ benchmark::scalar(
+ c,
+ &f,
+ "native-raster",
"rs_rastertoworldcoordx",
BenchmarkArgs::ArrayScalarScalar(Raster(64, 64), Int32(0, 63), Int32(0, 63)),
);
diff --git a/rust/sedona-raster-functions/src/register.rs b/rust/sedona-raster-functions/src/register.rs
index c850e36..d177caa 100644
--- a/rust/sedona-raster-functions/src/register.rs
+++ b/rust/sedona-raster-functions/src/register.rs
@@ -47,6 +47,7 @@
crate::rs_geotransform::rs_upperlefty_udf,
crate::rs_size::rs_height_udf,
crate::rs_size::rs_width_udf,
+ crate::rs_worldcoordinate::rs_rastertoworldcoord_udf,
crate::rs_worldcoordinate::rs_rastertoworldcoordx_udf,
crate::rs_worldcoordinate::rs_rastertoworldcoordy_udf,
);
diff --git a/rust/sedona-raster-functions/src/rs_worldcoordinate.rs b/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
index 09e9089..b342a40 100644
--- a/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
+++ b/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
@@ -17,7 +17,7 @@
use std::{sync::Arc, vec};
use crate::executor::RasterExecutor;
-use arrow_array::builder::Float64Builder;
+use arrow_array::builder::{BinaryBuilder, Float64Builder};
use arrow_schema::DataType;
use datafusion_common::{error::Result, exec_err, ScalarValue};
use datafusion_expr::{
@@ -25,6 +25,7 @@
};
use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
use sedona_raster::affine_transformation::to_world_coordinate;
+use sedona_schema::datatypes::Edges;
use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher};
/// RS_RasterToWorldCoordY() scalar UDF implementation
@@ -51,6 +52,18 @@
)
}
+/// RS_RasterToWorldCoord() scalar UDF documentation
+///
+/// Converts pixel coordinates to world coordinates
+pub fn rs_rastertoworldcoord_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_rastertoworldcoord",
+ vec![Arc::new(RsCoordinatePoint {})],
+ Volatility::Immutable,
+ Some(rs_rastertoworldcoord_doc()),
+ )
+}
+
fn rs_rastertoworldcoordy_doc() -> Documentation {
Documentation::builder(
DOC_SECTION_OTHER,
@@ -77,6 +90,19 @@
.build()
}
+fn rs_rastertoworldcoord_doc() -> Documentation {
+ Documentation::builder(
+ DOC_SECTION_OTHER,
+ "Returns the upper left X and Y coordinates of the given row and column of the given raster geometric units of the geo-referenced raster as a Point geometry. If any out of bounds values are given, the X and Y coordinates of the assumed point considering existing raster pixel size and skew values will be returned.".to_string(),
+ "RS_RasterToWorldCoord(raster: Raster, x: Integer, y: Integer)".to_string(),
+ )
+ .with_argument("raster", "Raster: Input raster")
+ .with_argument("x", "Integer: Column x into the raster")
+ .with_argument("y", "Integer: Row y into the raster")
+ .with_sql_example("SELECT RS_RasterToWorldCoord(RS_Example(), 0, 0)".to_string())
+ .build()
+}
+
#[derive(Debug, Clone)]
enum Coord {
X,
@@ -130,6 +156,55 @@
}
}
+#[derive(Debug)]
+struct RsCoordinatePoint;
+impl SedonaScalarKernel for RsCoordinatePoint {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let matcher = ArgMatcher::new(
+ vec![
+ ArgMatcher::is_raster(),
+ ArgMatcher::is_integer(),
+ ArgMatcher::is_integer(),
+ ],
+ SedonaType::Wkb(Edges::Planar, None),
+ );
+
+ matcher.match_args(args)
+ }
+
+ fn invoke_batch(
+ &self,
+ arg_types: &[SedonaType],
+ args: &[ColumnarValue],
+ ) -> Result<ColumnarValue> {
+ let executor = RasterExecutor::new(arg_types, args);
+ let mut item: [u8; 21] = [0x00; 21];
+ item[0] = 0x01;
+ item[1] = 0x01;
+ let mut builder = BinaryBuilder::with_capacity(
+ executor.num_iterations(),
+ item.len() * executor.num_iterations(),
+ );
+
+ let (x_opt, y_opt) = get_scalar_coord(&args[1], &args[2])?;
+
+ executor.execute_raster_void(|_i, raster_opt| {
+ match (raster_opt, x_opt, y_opt) {
+ (Some(raster), Some(x), Some(y)) => {
+ let (world_x, world_y) = to_world_coordinate(&raster, x, y);
+ item[5..13].copy_from_slice(&world_x.to_le_bytes());
+ item[13..21].copy_from_slice(&world_y.to_le_bytes());
+ builder.append_value(item);
+ }
+ (_, _, _) => builder.append_null(),
+ }
+ Ok(())
+ })?;
+
+ executor.finish(Arc::new(builder.finish()))
+ }
+}
+
fn extract_int_scalar(arg: &ColumnarValue) -> Result<Option<i64>> {
match arg {
ColumnarValue::Scalar(scalar) => {
@@ -157,8 +232,9 @@
use super::*;
use datafusion_expr::ScalarUDF;
use rstest::rstest;
- use sedona_schema::datatypes::RASTER;
+ use sedona_schema::datatypes::{RASTER, WKB_GEOMETRY};
use sedona_testing::compare::assert_array_equal;
+ use sedona_testing::create::create_array;
use sedona_testing::rasters::generate_test_rasters;
use sedona_testing::testers::ScalarUdfTester;
@@ -171,10 +247,14 @@
let udf: ScalarUDF = rs_rastertoworldcoordx_udf().into();
assert_eq!(udf.name(), "rs_rastertoworldcoordx");
assert!(udf.documentation().is_some());
+
+ let udf: ScalarUDF = rs_rastertoworldcoord_udf().into();
+ assert_eq!(udf.name(), "rs_rastertoworldcoord");
+ assert!(udf.documentation().is_some());
}
#[rstest]
- fn udf_invoke(#[values(Coord::Y, Coord::X)] coord: Coord) {
+ fn udf_invoke_xy(#[values(Coord::Y, Coord::X)] coord: Coord) {
let udf = match coord {
Coord::X => rs_rastertoworldcoordx_udf(),
Coord::Y => rs_rastertoworldcoordy_udf(),
@@ -202,4 +282,29 @@
.unwrap();
assert_array_equal(&result, &expected);
}
+
+ #[rstest]
+ fn udf_invoke_pt() {
+ let udf = rs_rastertoworldcoord_udf();
+ let tester = ScalarUdfTester::new(
+ udf.into(),
+ vec![
+ RASTER,
+ SedonaType::Arrow(DataType::Int32),
+ SedonaType::Arrow(DataType::Int32),
+ ],
+ );
+
+ let rasters = generate_test_rasters(3, Some(1)).unwrap();
+ // At 0,0 expect the upper left corner of the test values
+ let expected = &create_array(
+ &[Some("POINT (1 2)"), None, Some("POINT (3 4)")],
+ &WKB_GEOMETRY,
+ );
+
+ let result = tester
+ .invoke_array_scalar_scalar(Arc::new(rasters), 0_i32, 0_i32)
+ .unwrap();
+ assert_array_equal(&result, expected);
+ }
}
