This is an automated email from the ASF dual-hosted git repository.
kontinuation 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 402dcb31 fix(rust/sedona-raster-functions): RasterExecutor always call
the closure only once for scalar raster input (#559)
402dcb31 is described below
commit 402dcb318f010b7db943f28e5bd496a5d96f91a1
Author: Kristin Cowalcijk <[email protected]>
AuthorDate: Sun Feb 1 16:31:52 2026 +0800
fix(rust/sedona-raster-functions): RasterExecutor always call the closure
only once for scalar raster input (#559)
## Summary
- Fix RasterExecutor::execute_raster_void() to broadcast scalar rasters
across num_iterations() when other args are arrays.
- Add regression tests for scalar raster + array coordinate inputs in
rs_worldcoordinate and rs_rastercoordinate.
## Why
DataFusion can invoke these scalar UDFs with a scalar raster and array
coordinate arguments. Previously the executor only invoked the per-row closure
once for scalar rasters, producing output arrays of length 1 instead of the
expected batch length.
## Testing
- cargo test -p sedona-raster-functions
---
rust/sedona-raster-functions/src/executor.rs | 22 ++++--
rust/sedona-raster-functions/src/rs_envelope.rs | 2 +-
.../sedona-raster-functions/src/rs_geotransform.rs | 2 +-
.../src/rs_rastercoordinate.rs | 91 +++++++++++++++++++++-
.../src/rs_worldcoordinate.rs | 90 ++++++++++++++++++++-
5 files changed, 194 insertions(+), 13 deletions(-)
diff --git a/rust/sedona-raster-functions/src/executor.rs
b/rust/sedona-raster-functions/src/executor.rs
index e7bbbe9d..4f995c48 100644
--- a/rust/sedona-raster-functions/src/executor.rs
+++ b/rust/sedona-raster-functions/src/executor.rs
@@ -59,7 +59,7 @@ impl<'a, 'b> RasterExecutor<'a, 'b> {
/// 4. Calling the provided function with each raster
pub fn execute_raster_void<F>(&self, mut func: F) -> Result<()>
where
- F: FnMut(usize, Option<RasterRefImpl<'_>>) -> Result<()>,
+ F: FnMut(usize, Option<&RasterRefImpl<'_>>) -> Result<()>,
{
if self.arg_types[0] != RASTER {
return sedona_internal_err!("First argument must be a raster
type");
@@ -85,7 +85,7 @@ impl<'a, 'b> RasterExecutor<'a, 'b> {
continue;
}
let raster = raster_array.get(i)?;
- func(i, Some(raster))?;
+ func(i, Some(&raster))?;
}
Ok(())
@@ -93,14 +93,22 @@ impl<'a, 'b> RasterExecutor<'a, 'b> {
ColumnarValue::Scalar(scalar_value) => match scalar_value {
ScalarValue::Struct(arc_struct) => {
let raster_array =
RasterStructArray::new(arc_struct.as_ref());
- if raster_array.is_null(0) {
- func(0, None)
+ let raster_opt = if raster_array.is_null(0) {
+ None
} else {
- let raster = raster_array.get(0)?;
- func(0, Some(raster))
+ Some(raster_array.get(0)?)
+ };
+ for i in 0..self.num_iterations {
+ func(i, raster_opt.as_ref())?;
}
+ Ok(())
+ }
+ ScalarValue::Null => {
+ for i in 0..self.num_iterations {
+ func(i, None)?;
+ }
+ Ok(())
}
- ScalarValue::Null => func(0, None),
_ => sedona_internal_err!("Expected Struct scalar for raster"),
},
}
diff --git a/rust/sedona-raster-functions/src/rs_envelope.rs
b/rust/sedona-raster-functions/src/rs_envelope.rs
index 1ce7d157..8ecb60cb 100644
--- a/rust/sedona-raster-functions/src/rs_envelope.rs
+++ b/rust/sedona-raster-functions/src/rs_envelope.rs
@@ -82,7 +82,7 @@ impl SedonaScalarKernel for RsEnvelope {
executor.execute_raster_void(|_i, raster_opt| {
match raster_opt {
Some(raster) => {
- create_envelope_wkb(&raster, &mut builder)?;
+ create_envelope_wkb(raster, &mut builder)?;
builder.append_value([]);
}
None => builder.append_null(),
diff --git a/rust/sedona-raster-functions/src/rs_geotransform.rs
b/rust/sedona-raster-functions/src/rs_geotransform.rs
index 559cb3d8..0d98ee2b 100644
--- a/rust/sedona-raster-functions/src/rs_geotransform.rs
+++ b/rust/sedona-raster-functions/src/rs_geotransform.rs
@@ -251,7 +251,7 @@ impl SedonaScalarKernel for RsGeoTransform {
let metadata = raster.metadata();
match self.param {
GeoTransformParam::Rotation => {
- let rotation = rotation(&raster);
+ let rotation = rotation(raster);
builder.append_value(rotation);
}
GeoTransformParam::ScaleX =>
builder.append_value(metadata.scale_x()),
diff --git a/rust/sedona-raster-functions/src/rs_rastercoordinate.rs
b/rust/sedona-raster-functions/src/rs_rastercoordinate.rs
index d3bb59d2..bdfcf48f 100644
--- a/rust/sedona-raster-functions/src/rs_rastercoordinate.rs
+++ b/rust/sedona-raster-functions/src/rs_rastercoordinate.rs
@@ -155,7 +155,7 @@ impl SedonaScalarKernel for RsCoordinateMapper {
match (raster_opt, x_opt, y_opt) {
(Some(raster), Some(x), Some(y)) => {
- let (raster_x, raster_y) = to_raster_coordinate(&raster,
x, y)?;
+ let (raster_x, raster_y) = to_raster_coordinate(raster, x,
y)?;
match self.coord {
Coord::X => builder.append_value(raster_x),
Coord::Y => builder.append_value(raster_y),
@@ -216,7 +216,7 @@ impl SedonaScalarKernel for RsCoordinatePoint {
match (raster_opt, x_opt, y_opt) {
(Some(raster), Some(world_x), Some(world_y)) => {
- let (raster_x, raster_y) = to_raster_coordinate(&raster,
world_x, world_y)?;
+ let (raster_x, raster_y) = to_raster_coordinate(raster,
world_x, world_y)?;
item[5..13].copy_from_slice(&(raster_x as
f64).to_le_bytes());
item[13..21].copy_from_slice(&(raster_y as
f64).to_le_bytes());
builder.append_value(item);
@@ -234,6 +234,7 @@ impl SedonaScalarKernel for RsCoordinatePoint {
mod tests {
use super::*;
use arrow_schema::DataType;
+ use datafusion_common::ScalarValue;
use datafusion_expr::ScalarUDF;
use rstest::rstest;
use sedona_schema::datatypes::{RASTER, WKB_GEOMETRY};
@@ -368,4 +369,90 @@ mod tests {
.unwrap();
assert_array_equal(&result, &expected);
}
+
+ #[rstest]
+ fn udf_invoke_xy_with_scalar_raster_array_coords(#[values(Coord::Y,
Coord::X)] coord: Coord) {
+ let udf = match coord {
+ Coord::X => rs_worldtorastercoordx_udf(),
+ Coord::Y => rs_worldtorastercoordy_udf(),
+ };
+ let tester = ScalarUdfTester::new(
+ udf.into(),
+ vec![
+ RASTER,
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ ],
+ );
+
+ let rasters = generate_test_rasters(2, Some(0)).unwrap();
+ let scalar_raster = ScalarValue::try_from_array(&rasters, 1).unwrap();
+
+ let world_x = Arc::new(arrow_array::Float64Array::from(vec![2.0, 2.5,
3.25]));
+ let world_y = Arc::new(arrow_array::Float64Array::from(vec![3.0, 2.5,
1.75]));
+
+ let expected_coords = match coord {
+ Coord::X => vec![Some(0_i64), Some(4_i64), Some(10_i64)],
+ Coord::Y => vec![Some(0_i64), Some(3_i64), Some(8_i64)],
+ };
+
+ let result = tester
+ .invoke(vec![
+ ColumnarValue::Scalar(scalar_raster),
+ ColumnarValue::Array(world_x.clone()),
+ ColumnarValue::Array(world_y.clone()),
+ ])
+ .unwrap();
+
+ let array = match result {
+ ColumnarValue::Array(array) => array,
+ ColumnarValue::Scalar(_) => panic!("Expected array result"),
+ };
+
+ let expected: Arc<dyn arrow_array::Array> =
+ Arc::new(arrow_array::Int64Array::from(expected_coords));
+ assert_array_equal(&array, &expected);
+ }
+
+ #[test]
+ fn udf_invoke_pt_with_scalar_raster_array_coords() {
+ let udf = rs_worldtorastercoord_udf();
+ let tester = ScalarUdfTester::new(
+ udf.into(),
+ vec![
+ RASTER,
+ SedonaType::Arrow(DataType::Float64),
+ SedonaType::Arrow(DataType::Float64),
+ ],
+ );
+
+ let rasters = generate_test_rasters(2, Some(0)).unwrap();
+ let scalar_raster = ScalarValue::try_from_array(&rasters, 1).unwrap();
+
+ let world_x = Arc::new(arrow_array::Float64Array::from(vec![2.0, 2.5,
3.25]));
+ let world_y = Arc::new(arrow_array::Float64Array::from(vec![3.0, 2.5,
1.75]));
+
+ let result = tester
+ .invoke(vec![
+ ColumnarValue::Scalar(scalar_raster),
+ ColumnarValue::Array(world_x.clone()),
+ ColumnarValue::Array(world_y.clone()),
+ ])
+ .unwrap();
+
+ let array = match result {
+ ColumnarValue::Array(array) => array,
+ ColumnarValue::Scalar(_) => panic!("Expected array result"),
+ };
+
+ let expected = create_array(
+ &[
+ Some("POINT (0 0)"),
+ Some("POINT (4 3)"),
+ Some("POINT (10 8)"),
+ ],
+ &WKB_GEOMETRY,
+ );
+ assert_array_equal(&array, &expected);
+ }
}
diff --git a/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
b/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
index 4bd264c6..617e4fcd 100644
--- a/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
+++ b/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
@@ -153,7 +153,7 @@ impl SedonaScalarKernel for RsCoordinateMapper {
match (raster_opt, x_opt, y_opt) {
(Some(raster), Some(x), Some(y)) => {
- let (world_x, world_y) = to_world_coordinate(&raster, x,
y);
+ let (world_x, world_y) = to_world_coordinate(raster, x, y);
match self.coord {
Coord::X => builder.append_value(world_x),
Coord::Y => builder.append_value(world_y),
@@ -214,7 +214,7 @@ impl SedonaScalarKernel for RsCoordinatePoint {
match (raster_opt, x_opt, y_opt) {
(Some(raster), Some(x), Some(y)) => {
- let (world_x, world_y) = to_world_coordinate(&raster, x,
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);
@@ -232,6 +232,7 @@ impl SedonaScalarKernel for RsCoordinatePoint {
mod tests {
use super::*;
use arrow_array::Array;
+ use datafusion_common::ScalarValue;
use datafusion_expr::ScalarUDF;
use rstest::rstest;
use sedona_schema::datatypes::{RASTER, WKB_GEOMETRY};
@@ -357,4 +358,89 @@ mod tests {
}
}
}
+
+ #[rstest]
+ fn udf_invoke_xy_with_scalar_raster_array_coords(#[values(Coord::Y,
Coord::X)] coord: Coord) {
+ let udf = match coord {
+ Coord::X => rs_rastertoworldcoordx_udf(),
+ Coord::Y => rs_rastertoworldcoordy_udf(),
+ };
+ let tester = ScalarUdfTester::new(
+ udf.into(),
+ vec![
+ RASTER,
+ SedonaType::Arrow(DataType::Int32),
+ SedonaType::Arrow(DataType::Int32),
+ ],
+ );
+
+ let rasters = generate_test_rasters(2, None).unwrap();
+ let scalar_raster = ScalarValue::try_from_array(&rasters, 1).unwrap();
+
+ let x_vals = vec![0_i32, 1_i32];
+ let y_vals = vec![0_i32, 1_i32];
+ let x_coords: Arc<dyn Array> =
Arc::new(arrow_array::Int32Array::from(x_vals.clone()));
+ let y_coords: Arc<dyn Array> =
Arc::new(arrow_array::Int32Array::from(y_vals.clone()));
+
+ let result = tester
+ .invoke(vec![
+ ColumnarValue::Scalar(scalar_raster),
+ ColumnarValue::Array(x_coords),
+ ColumnarValue::Array(y_coords),
+ ])
+ .unwrap();
+
+ let array = match result {
+ ColumnarValue::Array(array) => array,
+ ColumnarValue::Scalar(_) => panic!("Expected array result"),
+ };
+
+ let expected_values = match coord {
+ Coord::X => vec![Some(2.0), Some(2.13)],
+ Coord::Y => vec![Some(3.0), Some(2.84)],
+ };
+ let expected: Arc<dyn arrow_array::Array> =
+ Arc::new(arrow_array::Float64Array::from(expected_values));
+ assert_array_equal(&array, &expected);
+ }
+
+ #[test]
+ fn udf_invoke_pt_with_scalar_raster_array_coords() {
+ 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(2, None).unwrap();
+ let scalar_raster = ScalarValue::try_from_array(&rasters, 1).unwrap();
+
+ let x_vals = vec![0_i32, 1_i32];
+ let y_vals = vec![0_i32, 1_i32];
+ let x_coords: Arc<dyn Array> =
Arc::new(arrow_array::Int32Array::from(x_vals.clone()));
+ let y_coords: Arc<dyn Array> =
Arc::new(arrow_array::Int32Array::from(y_vals.clone()));
+
+ let result = tester
+ .invoke(vec![
+ ColumnarValue::Scalar(scalar_raster),
+ ColumnarValue::Array(x_coords),
+ ColumnarValue::Array(y_coords),
+ ])
+ .unwrap();
+
+ let array = match result {
+ ColumnarValue::Array(array) => array,
+ ColumnarValue::Scalar(_) => panic!("Expected array result"),
+ };
+
+ let expected = create_array(
+ &[Some("POINT (2 3)"), Some("POINT (2.13 2.84)")],
+ &WKB_GEOMETRY,
+ );
+ assert_array_equal(&array, &expected);
+ }
}
