You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by ja...@apache.org on 2026/08/27 17:26:26 UTC
(sedona-db) branch main updated: refactor(rust/sedona-raster-functions): adopt start_raster_from / copy_into in the structural raster functions (#1193)
This is an automated email from the ASF dual-hosted git repository.
james-willis 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 5c1eada0f refactor(rust/sedona-raster-functions): adopt start_raster_from / copy_into in the structural raster functions (#1193)
5c1eada0f is described below
commit 5c1eada0fc85ac88b2fd89a0809db27ba555f32c
Author: James Willis <ja...@users.noreply.github.com>
AuthorDate: Thu Aug 27 10:26:21 2026 -0700
refactor(rust/sedona-raster-functions): adopt start_raster_from / copy_into in the structural raster functions (#1193)
---
rust/sedona-raster-functions/src/rs_dim_band.rs | 75 ++++++++++++----------
.../src/rs_ensure_loaded.rs | 22 +------
rust/sedona-raster-functions/src/rs_slice.rs | 49 ++++----------
3 files changed, 59 insertions(+), 87 deletions(-)
diff --git a/rust/sedona-raster-functions/src/rs_dim_band.rs b/rust/sedona-raster-functions/src/rs_dim_band.rs
index 34b56ca70..9b2777d84 100644
--- a/rust/sedona-raster-functions/src/rs_dim_band.rs
+++ b/rust/sedona-raster-functions/src/rs_dim_band.rs
@@ -22,10 +22,9 @@ use datafusion_common::cast::as_string_view_array;
use datafusion_common::error::Result;
use datafusion_common::exec_err;
use datafusion_expr::{ColumnarValue, Volatility};
-use sedona_common::sedona_internal_datafusion_err;
use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
-use sedona_raster::builder::{RasterBuilder, StartBandArgs};
-use sedona_raster::traits::RasterRef;
+use sedona_raster::builder::{RasterBuilder, RasterOverrides, StartBandArgs};
+use sedona_raster::traits::{BandOverrides, RasterRef};
use sedona_schema::datatypes::SedonaType;
use sedona_schema::matchers::ArgMatcher;
@@ -90,16 +89,7 @@ impl SedonaScalarKernel for RsDimToBand {
require_any_band_has_dim(raster, name, "RS_DimToBand")?;
- let t: [f64; 6] = raster.transform().try_into().map_err(|_| {
- sedona_internal_datafusion_err!("raster transform is not 6 elements")
- })?;
- let spatial_dims = raster.spatial_dims();
- new_builder.start_raster_nd(
- &t,
- &spatial_dims,
- raster.spatial_shape(),
- raster.crs(),
- )?;
+ new_builder.start_raster_from(raster, RasterOverrides::default())?;
for band_idx in 0..raster.num_bands() {
let band = raster.band(band_idx)?;
@@ -108,16 +98,7 @@ impl SedonaScalarKernel for RsDimToBand {
match maybe_dim_idx {
None => {
// Band doesn't have this dimension -- pass through
- let dim_names = band.dim_names();
- let band_name = raster.band_name(band_idx);
- new_builder.start_band(StartBandArgs {
- name: band_name,
- nodata: band.nodata(),
- ..StartBandArgs::new(&dim_names, band.shape(), band.data_type())
- })?;
- let ndb = band.nd_buffer()?;
- let data = ndb.as_contiguous()?;
- new_builder.band_data_writer().append_value(data);
+ band.copy_into(&mut new_builder, BandOverrides::default())?;
new_builder.finish_band()?;
}
Some(dim_idx) => {
@@ -290,16 +271,7 @@ impl SedonaScalarKernel for RsBandToDim {
let nodata = ref_nodata.as_deref();
- let t: [f64; 6] = raster.transform().try_into().map_err(|_| {
- sedona_internal_datafusion_err!("raster transform is not 6 elements")
- })?;
- let spatial_dims = raster.spatial_dims();
- new_builder.start_raster_nd(
- &t,
- &spatial_dims,
- raster.spatial_shape(),
- raster.crs(),
- )?;
+ new_builder.start_raster_from(raster, RasterOverrides::default())?;
new_builder.start_band(StartBandArgs {
nodata,
..StartBandArgs::new(&new_dim_names, &new_shape, ref_data_type)
@@ -379,6 +351,43 @@ mod tests {
assert_eq!(raster.band_name(2), Some("temp_time_2"));
}
+ #[test]
+ fn dimtoband_passes_through_bands_without_dim_preserving_names() {
+ // Heterogeneous raster: the band carrying `time` expands into one band
+ // per index, while the band without it is passed through untouched.
+ // The pass-through derives the output band via `BandRef::copy_into`,
+ // which inherits the name — assert it survives, since a dropped name is
+ // silent otherwise.
+ let udf: ScalarUDF = rs_dimtoband_udf().into();
+ let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Utf8)]);
+
+ let rasters = RasterSpec::nd(&["time", "y", "x"], &[2, 2, 2])
+ .crs(None)
+ .band_nd(&["y", "x"], &[2, 2], BandDataType::UInt8)
+ .name("elevation")
+ .band(BandDataType::UInt8)
+ .name("temperature")
+ .build();
+
+ let result = tester
+ .invoke_array_scalar(Arc::new(rasters), "time")
+ .unwrap();
+
+ let result_struct = result.as_any().downcast_ref::<StructArray>().unwrap();
+ let raster_array = RasterStructArray::try_new(result_struct).unwrap();
+ let raster = raster_array.get(0).unwrap();
+
+ // Pass-through band keeps its own name; the expanded band contributes
+ // one suffixed band per `time` index.
+ assert_eq!(raster.num_bands(), 3);
+ assert_eq!(raster.band_name(0), Some("elevation"));
+ assert_eq!(raster.band_name(1), Some("temperature_time_0"));
+ assert_eq!(raster.band_name(2), Some("temperature_time_1"));
+ // The pass-through band is emitted unchanged, not expanded.
+ assert_eq!(raster.band(0).unwrap().dim_names(), vec!["y", "x"]);
+ assert_eq!(raster.band(0).unwrap().shape(), &[2, 2]);
+ }
+
#[test]
fn dimtoband_null_raster() {
let udf: ScalarUDF = rs_dimtoband_udf().into();
diff --git a/rust/sedona-raster-functions/src/rs_ensure_loaded.rs b/rust/sedona-raster-functions/src/rs_ensure_loaded.rs
index 20c89f28d..24e0a3f8b 100644
--- a/rust/sedona-raster-functions/src/rs_ensure_loaded.rs
+++ b/rust/sedona-raster-functions/src/rs_ensure_loaded.rs
@@ -40,7 +40,7 @@ use datafusion_expr::{
};
use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
use sedona_raster::array::RasterStructArray;
-use sedona_raster::builder::{RasterBuilder, StartBandArgs};
+use sedona_raster::builder::{RasterBuilder, RasterOverrides, StartBandArgs};
use sedona_raster::raster_loader::{
AsyncRasterLoader, RasterLoadRequest, RasterLoaderConfig, RasterLoaderRegistry,
};
@@ -294,27 +294,11 @@ where
)
})?;
- // Owned per-row metadata so the borrows don't span the per-band
- // `await` points further down.
- let transform: [f64; 6] = raster.transform().try_into().map_err(|_| {
- sedona_internal_datafusion_err!(
- "RS_EnsureLoaded: raster row {raster_idx} transform is not 6 elements"
- )
- })?;
- let spatial_dims_owned: Vec<String> = raster
- .spatial_dims()
- .iter()
- .map(|s| s.to_string())
- .collect();
- let spatial_dims: Vec<&str> = spatial_dims_owned.iter().map(String::as_str).collect();
- let spatial_shape: Vec<i64> = raster.spatial_shape().to_vec();
- let crs: Option<String> = raster.crs().map(|s| s.to_string());
-
builder
- .start_raster_nd(&transform, &spatial_dims, &spatial_shape, crs.as_deref())
+ .start_raster_from(&raster, RasterOverrides::default())
.map_err(|e| {
sedona_internal_datafusion_err!(
- "RS_EnsureLoaded: start_raster_nd failed at row {raster_idx}: {e}"
+ "RS_EnsureLoaded: start_raster_from failed at row {raster_idx}: {e}"
)
})?;
diff --git a/rust/sedona-raster-functions/src/rs_slice.rs b/rust/sedona-raster-functions/src/rs_slice.rs
index 6ead2ed24..42eaeeddc 100644
--- a/rust/sedona-raster-functions/src/rs_slice.rs
+++ b/rust/sedona-raster-functions/src/rs_slice.rs
@@ -22,10 +22,9 @@ use datafusion_common::cast::{as_int64_array, as_string_array};
use datafusion_common::error::Result;
use datafusion_common::exec_err;
use datafusion_expr::{ColumnarValue, Volatility};
-use sedona_common::sedona_internal_datafusion_err;
use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
-use sedona_raster::builder::{RasterBuilder, StartBandArgs};
-use sedona_raster::traits::{BandRef, RasterRef};
+use sedona_raster::builder::{RasterBuilder, RasterOverrides, StartBandArgs};
+use sedona_raster::traits::{BandOverrides, BandRef, RasterRef};
use sedona_schema::datatypes::SedonaType;
use sedona_schema::matchers::ArgMatcher;
@@ -98,11 +97,7 @@ impl SedonaScalarKernel for RsSlice {
}
validate_not_spatial(raster, name, "RS_Slice")?;
- let t: [f64; 6] = raster.transform().try_into().map_err(|_| {
- sedona_internal_datafusion_err!("raster transform is not 6 elements")
- })?;
- let spatial_dims = raster.spatial_dims();
- new_builder.start_raster_nd(&t, &spatial_dims, raster.spatial_shape(), raster.crs())?;
+ new_builder.start_raster_from(raster, RasterOverrides::default())?;
require_any_band_has_dim(raster, name, "RS_Slice")?;
@@ -114,16 +109,7 @@ impl SedonaScalarKernel for RsSlice {
// RS_DimToBand, and matches xarray's `isel` — variables
// without the indexed dim are left alone.
let Some(dim_idx) = band.dim_index(name) else {
- let dim_names = band.dim_names();
- let band_name = raster.band_name(band_idx);
- new_builder.start_band(StartBandArgs {
- name: band_name,
- nodata: band.nodata(),
- ..StartBandArgs::new(&dim_names, band.shape(), band.data_type())
- })?;
- let ndb = band.nd_buffer()?;
- let data = ndb.as_contiguous()?;
- new_builder.band_data_writer().append_value(data);
+ band.copy_into(&mut new_builder, BandOverrides::default())?;
new_builder.finish_band()?;
continue;
};
@@ -257,11 +243,7 @@ impl SedonaScalarKernel for RsSliceRange {
);
}
- let t: [f64; 6] = raster.transform().try_into().map_err(|_| {
- sedona_internal_datafusion_err!("raster transform is not 6 elements")
- })?;
- let spatial_dims = raster.spatial_dims();
- new_builder.start_raster_nd(&t, &spatial_dims, raster.spatial_shape(), raster.crs())?;
+ new_builder.start_raster_from(raster, RasterOverrides::default())?;
require_any_band_has_dim(raster, name, "RS_SliceRange")?;
@@ -272,16 +254,7 @@ impl SedonaScalarKernel for RsSliceRange {
// dimension are emitted unchanged. Same convention as
// RS_Slice and RS_DimToBand.
let Some(dim_idx) = band.dim_index(name) else {
- let dim_names = band.dim_names();
- let band_name = raster.band_name(band_idx);
- new_builder.start_band(StartBandArgs {
- name: band_name,
- nodata: band.nodata(),
- ..StartBandArgs::new(&dim_names, band.shape(), band.data_type())
- })?;
- let ndb = band.nd_buffer()?;
- let data = ndb.as_contiguous()?;
- new_builder.band_data_writer().append_value(data);
+ band.copy_into(&mut new_builder, BandOverrides::default())?;
new_builder.finish_band()?;
continue;
};
@@ -420,7 +393,9 @@ mod tests {
RasterSpec::nd(&["time", "y", "x"], &[3, 2, 3])
.crs(None)
.band_nd(&["y", "x"], &[2, 3], BandDataType::UInt8)
+ .name("elevation")
.band(BandDataType::UInt8)
+ .name("temperature")
.build()
}
@@ -449,7 +424,9 @@ mod tests {
let expected = RasterSpec::nd(&["time", "y", "x"], &[3, 2, 3])
.crs(None)
.band_values_nd(&["y", "x"], &[2, 3], &(0u8..6).collect::<Vec<u8>>())
- .band_values_nd(&["y", "x"], &[2, 3], &(6u8..12).collect::<Vec<u8>>());
+ .name("elevation")
+ .band_values_nd(&["y", "x"], &[2, 3], &(6u8..12).collect::<Vec<u8>>())
+ .name("temperature");
assert_rasters_equal(&result, &[Some(expected)]);
}
@@ -481,11 +458,13 @@ mod tests {
let expected = RasterSpec::nd(&["time", "y", "x"], &[3, 2, 3])
.crs(None)
.band_values_nd(&["y", "x"], &[2, 3], &(0u8..6).collect::<Vec<u8>>())
+ .name("elevation")
.band_values_nd(
&["time", "y", "x"],
&[2, 2, 3],
&(6u8..18).collect::<Vec<u8>>(),
- );
+ )
+ .name("temperature");
assert_rasters_equal(&result, &[Some(expected)]);
}