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 a95d9a48 fix(sedona-raster-functions): Handle array inputs in
coordinate conversion functions (#532)
a95d9a48 is described below
commit a95d9a488aacf17f794065bb0409cf456ca24603
Author: Kristin Cowalcijk <[email protected]>
AuthorDate: Thu Jan 22 15:32:05 2026 +0800
fix(sedona-raster-functions): Handle array inputs in coordinate conversion
functions (#532)
Previously, RS_RasterToWorldCoord and RS_WorldToRasterCoord only handled
scalar inputs and would error with 'Expected scalar ... argument' when array
inputs were passed.
This commit fixes both functions to properly handle array inputs by:
- Using into_array() to expand scalar inputs to arrays
- Processing all values in a loop with proper null handling
Added test cases for array input handling.
---
.../src/rs_rastercoordinate.rs | 102 ++++++++++++++-------
.../src/rs_worldcoordinate.rs | 100 +++++++++++++++-----
2 files changed, 145 insertions(+), 57 deletions(-)
diff --git a/rust/sedona-raster-functions/src/rs_rastercoordinate.rs
b/rust/sedona-raster-functions/src/rs_rastercoordinate.rs
index 5b00a3e0..d3bb59d2 100644
--- a/rust/sedona-raster-functions/src/rs_rastercoordinate.rs
+++ b/rust/sedona-raster-functions/src/rs_rastercoordinate.rs
@@ -19,8 +19,8 @@ use std::{sync::Arc, vec};
use crate::executor::RasterExecutor;
use arrow_array::builder::{BinaryBuilder, Int64Builder};
use arrow_schema::DataType;
-use datafusion_common::ScalarValue;
-use datafusion_common::{error::Result, exec_err};
+use datafusion_common::cast::as_float64_array;
+use datafusion_common::error::Result;
use datafusion_expr::{
scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation,
Volatility,
};
@@ -139,17 +139,29 @@ impl SedonaScalarKernel for RsCoordinateMapper {
let executor = RasterExecutor::new(arg_types, args);
let mut builder =
Int64Builder::with_capacity(executor.num_iterations());
- let coord_opt = extract_scalar_coord(&args[1], &args[2])?;
+ // Expand world x and y coordinate parameters to arrays and cast to
Float64
+ let world_x_array = args[1].clone().cast_to(&DataType::Float64, None)?;
+ let world_x_array =
world_x_array.into_array(executor.num_iterations())?;
+ let world_x_array = as_float64_array(&world_x_array)?;
+ let world_y_array = args[2].clone().cast_to(&DataType::Float64, None)?;
+ let world_y_array =
world_y_array.into_array(executor.num_iterations())?;
+ let world_y_array = as_float64_array(&world_y_array)?;
+ let mut world_x_iter = world_x_array.iter();
+ let mut world_y_iter = world_y_array.iter();
+
executor.execute_raster_void(|_i, raster_opt| {
- match (raster_opt, coord_opt) {
- (Some(raster), (Some(x), Some(y))) => {
+ let x_opt = world_x_iter.next().unwrap();
+ let y_opt = world_y_iter.next().unwrap();
+
+ match (raster_opt, x_opt, y_opt) {
+ (Some(raster), Some(x), Some(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),
};
}
- (_, _) => builder.append_null(),
+ (_, _, _) => builder.append_null(),
}
Ok(())
})?;
@@ -188,16 +200,28 @@ impl SedonaScalarKernel for RsCoordinatePoint {
item.len() * executor.num_iterations(),
);
- let coord_opt = extract_scalar_coord(&args[1], &args[2])?;
+ // Expand world x and y coordinate parameters to arrays and cast to
Float64
+ let world_x_array = args[1].clone().cast_to(&DataType::Float64, None)?;
+ let world_x_array =
world_x_array.into_array(executor.num_iterations())?;
+ let world_x_array = as_float64_array(&world_x_array)?;
+ let world_y_array = args[2].clone().cast_to(&DataType::Float64, None)?;
+ let world_y_array =
world_y_array.into_array(executor.num_iterations())?;
+ let world_y_array = as_float64_array(&world_y_array)?;
+ let mut world_x_iter = world_x_array.iter();
+ let mut world_y_iter = world_y_array.iter();
+
executor.execute_raster_void(|_i, raster_opt| {
- match (raster_opt, coord_opt) {
- (Some(raster), (Some(world_x), Some(world_y))) => {
+ let x_opt = world_x_iter.next().unwrap();
+ let y_opt = world_y_iter.next().unwrap();
+
+ 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)?;
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);
}
- (_, _) => builder.append_null(),
+ (_, _, _) => builder.append_null(),
}
Ok(())
})?;
@@ -206,28 +230,6 @@ impl SedonaScalarKernel for RsCoordinatePoint {
}
}
-fn extract_float_scalar(arg: &ColumnarValue) -> Result<Option<f64>> {
- match arg {
- ColumnarValue::Scalar(scalar) => {
- let f64_val = scalar.cast_to(&DataType::Float64)?;
- match f64_val {
- ScalarValue::Float64(Some(v)) => Ok(Some(v)),
- _ => Ok(None),
- }
- }
- _ => exec_err!("Expected scalar float argument for coordinate"),
- }
-}
-
-fn extract_scalar_coord(
- x_arg: &ColumnarValue,
- y_arg: &ColumnarValue,
-) -> Result<(Option<f64>, Option<f64>)> {
- let x_opt = extract_float_scalar(x_arg)?;
- let y_opt = extract_float_scalar(y_arg)?;
- Ok((x_opt, y_opt))
-}
-
#[cfg(test)]
mod tests {
use super::*;
@@ -330,4 +332,40 @@ mod tests {
.to_string()
.contains("determinant is zero"));
}
+
+ #[rstest]
+ fn udf_invoke_xy_with_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),
+ ],
+ );
+
+ // Use only raster 1 (invertible) with different world coordinates
+ // Raster 1: upper_left=(2,3), scale_x=0.1, scale_y=-0.2, skew_x=0.03,
skew_y=0.04
+ // For world coords (2,3) -> raster (0,0) (the upper left corner)
+ // For world coords (2.1, 2.8) -> need to solve the inverse
+ let rasters = generate_test_rasters(2, Some(0)).unwrap();
+ let world_x = Arc::new(arrow_array::Float64Array::from(vec![2.0,
2.0]));
+ let world_y = Arc::new(arrow_array::Float64Array::from(vec![3.0,
3.0]));
+
+ let expected_values = match coord {
+ Coord::X => vec![None, Some(0_i64)],
+ Coord::Y => vec![None, Some(0_i64)],
+ };
+ let expected: Arc<dyn arrow_array::Array> =
+ Arc::new(arrow_array::Int64Array::from(expected_values));
+
+ let result = tester
+ .invoke_arrays(vec![Arc::new(rasters), world_x, world_y])
+ .unwrap();
+ assert_array_equal(&result, &expected);
+ }
}
diff --git a/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
b/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
index b342a403..4bd264c6 100644
--- a/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
+++ b/rust/sedona-raster-functions/src/rs_worldcoordinate.rs
@@ -19,7 +19,8 @@ use std::{sync::Arc, vec};
use crate::executor::RasterExecutor;
use arrow_array::builder::{BinaryBuilder, Float64Builder};
use arrow_schema::DataType;
-use datafusion_common::{error::Result, exec_err, ScalarValue};
+use datafusion_common::cast::as_int64_array;
+use datafusion_common::error::Result;
use datafusion_expr::{
scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation,
Volatility,
};
@@ -136,9 +137,20 @@ impl SedonaScalarKernel for RsCoordinateMapper {
let executor = RasterExecutor::new(arg_types, args);
let mut builder =
Float64Builder::with_capacity(executor.num_iterations());
- let (x_opt, y_opt) = get_scalar_coord(&args[1], &args[2])?;
+ // Expand x and y coordinate parameters to arrays and cast to Int64
+ let x_array = args[1].clone().cast_to(&DataType::Int64, None)?;
+ let x_array = x_array.into_array(executor.num_iterations())?;
+ let x_array = as_int64_array(&x_array)?;
+ let y_array = args[2].clone().cast_to(&DataType::Int64, None)?;
+ let y_array = y_array.into_array(executor.num_iterations())?;
+ let y_array = as_int64_array(&y_array)?;
+ let mut x_iter = x_array.iter();
+ let mut y_iter = y_array.iter();
executor.execute_raster_void(|_i, raster_opt| {
+ let x_opt = x_iter.next().unwrap();
+ let y_opt = y_iter.next().unwrap();
+
match (raster_opt, x_opt, y_opt) {
(Some(raster), Some(x), Some(y)) => {
let (world_x, world_y) = to_world_coordinate(&raster, x,
y);
@@ -186,9 +198,20 @@ impl SedonaScalarKernel for RsCoordinatePoint {
item.len() * executor.num_iterations(),
);
- let (x_opt, y_opt) = get_scalar_coord(&args[1], &args[2])?;
+ // Expand x and y coordinate parameters to arrays and cast to Int64
+ let x_array = args[1].clone().cast_to(&DataType::Int64, None)?;
+ let x_array = x_array.into_array(executor.num_iterations())?;
+ let x_array = as_int64_array(&x_array)?;
+ let y_array = args[2].clone().cast_to(&DataType::Int64, None)?;
+ let y_array = y_array.into_array(executor.num_iterations())?;
+ let y_array = as_int64_array(&y_array)?;
+ let mut x_iter = x_array.iter();
+ let mut y_iter = y_array.iter();
executor.execute_raster_void(|_i, raster_opt| {
+ let x_opt = x_iter.next().unwrap();
+ let y_opt = y_iter.next().unwrap();
+
match (raster_opt, x_opt, y_opt) {
(Some(raster), Some(x), Some(y)) => {
let (world_x, world_y) = to_world_coordinate(&raster, x,
y);
@@ -205,31 +228,10 @@ impl SedonaScalarKernel for RsCoordinatePoint {
}
}
-fn extract_int_scalar(arg: &ColumnarValue) -> Result<Option<i64>> {
- match arg {
- ColumnarValue::Scalar(scalar) => {
- let i64_val = scalar.cast_to(&DataType::Int64)?;
- match i64_val {
- ScalarValue::Int64(Some(v)) => Ok(Some(v)),
- _ => Ok(None),
- }
- }
- _ => exec_err!("Expected scalar integer argument for coordinate"),
- }
-}
-
-fn get_scalar_coord(
- x_arg: &ColumnarValue,
- y_arg: &ColumnarValue,
-) -> Result<(Option<i64>, Option<i64>)> {
- let x_opt = extract_int_scalar(x_arg)?;
- let y_opt = extract_int_scalar(y_arg)?;
- Ok((x_opt, y_opt))
-}
-
#[cfg(test)]
mod tests {
use super::*;
+ use arrow_array::Array;
use datafusion_expr::ScalarUDF;
use rstest::rstest;
use sedona_schema::datatypes::{RASTER, WKB_GEOMETRY};
@@ -307,4 +309,52 @@ mod tests {
.unwrap();
assert_array_equal(&result, expected);
}
+
+ #[rstest]
+ fn udf_invoke_xy_with_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(3, Some(1)).unwrap();
+ // Test with different pixel coordinates for each raster
+ // Raster 0: upper_left=(1,2), scales=(0,0), so any pixel gives (1,2)
+ // Raster 1: null
+ // Raster 2: upper_left=(3,4), scale_x=0.2, scale_y=-0.4, skew_x=0.06,
skew_y=0.08
+ // At pixel (1,2): x = 3 + 0.2*1 + 0.06*2 = 3.32, y = 4 +
0.08*1 + (-0.4)*2 = 3.28
+ let x_coords = Arc::new(arrow_array::Int32Array::from(vec![0, 0, 1]));
+ let y_coords = Arc::new(arrow_array::Int32Array::from(vec![0, 0, 2]));
+
+ let result = tester
+ .invoke_arrays(vec![Arc::new(rasters), x_coords, y_coords])
+ .unwrap();
+
+ let float_array = result
+ .as_any()
+ .downcast_ref::<arrow_array::Float64Array>()
+ .expect("Expected Float64Array");
+
+ // Check each value with approximate comparison due to floating point
+ match coord {
+ Coord::X => {
+ assert!((float_array.value(0) - 1.0).abs() < 1e-10);
+ assert!(float_array.is_null(1));
+ assert!((float_array.value(2) - 3.32).abs() < 1e-10);
+ }
+ Coord::Y => {
+ assert!((float_array.value(0) - 2.0).abs() < 1e-10);
+ assert!(float_array.is_null(1));
+ assert!((float_array.value(2) - 3.28).abs() < 1e-10);
+ }
+ }
+ }
}
