diff --git a/Cargo.lock b/Cargo.lock index 95b434345..59326e150 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6150,6 +6150,7 @@ dependencies = [ "sedona-expr", "sedona-functions", "sedona-gdal", + "sedona-proj", "sedona-raster", "sedona-raster-functions", "sedona-schema", diff --git a/docs/reference/sql/rs_clip.qmd b/docs/reference/sql/rs_clip.qmd new file mode 100644 index 000000000..ee08039dc --- /dev/null +++ b/docs/reference/sql/rs_clip.qmd @@ -0,0 +1,131 @@ +--- +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +title: RS_Clip +description: > + Clips a raster to a geometry, setting pixels outside the geometry to nodata + and (by default) cropping to the geometry's bounding box. +kernels: + - returns: raster + args: + - raster + - {name: band, type: integer} + - {name: geom, type: geometry} + - returns: raster + args: + - raster + - {name: band, type: integer} + - {name: geom, type: geometry} + - {name: all_touched, type: boolean} + - returns: raster + args: + - raster + - {name: band, type: integer} + - {name: geom, type: geometry} + - {name: all_touched, type: boolean} + - {name: no_data_value, type: double} + - returns: raster + args: + - raster + - {name: band, type: integer} + - {name: geom, type: geometry} + - {name: all_touched, type: boolean} + - {name: no_data_value, type: double} + - {name: crop, type: boolean} + - returns: raster + args: + - raster + - name: band + type: integer + description: > + 1-based index of the band to clip, or 0 to clip every band. A NULL + band yields a NULL result. + - name: geom + type: geometry + description: > + Clip geometry. Reprojected into the raster's CRS when both carry one; + it is an error for exactly one side to have a CRS. + - name: all_touched + type: boolean + description: > + If true, keep every pixel the geometry touches; otherwise keep only + pixels whose center falls inside it. Defaults to false. + - name: no_data_value + type: double + description: > + Value written to pixels outside the geometry. Defaults to the band's + own nodata value, or the band data type's minimum if it has none. It is + an error if the value does not fit the band's data type. + - name: crop + type: boolean + description: Crop the output to the geometry's bounding box. Defaults to true. + - name: lenient + type: boolean + description: > + If true, a geometry that does not intersect the raster yields NULL; + if false, it raises an error. Defaults to true. +--- + +## Description + +`RS_Clip` clips a raster to a geometry, in the spirit of PostGIS `ST_Clip`. +Pixels of the selected band that fall outside the geometry are set to the +band's nodata value (or `no_data_value` if given); pixels inside are kept. When +`crop` is true (the default) the output is cropped to the geometry's bounding +box intersected with the raster extent, snapped to the pixel grid; unselected +pixels within that window remain as nodata padding. Otherwise the original +extent is preserved. + +The geometry is reprojected into the raster's CRS when both carry one. If +exactly one side has a CRS, the call errors; if neither does, the geometry is +used as-is. + +A pixel is selected only when its center falls inside the geometry. A geometry +smaller than a pixel, or a thin sliver, can therefore select no pixels even +though it overlaps the raster; set `all_touched` to true to keep every pixel the +geometry touches instead. + +The clip is a 2-D `(y, x)` operation. For an N-D raster (bands with extra +dimensions such as `time`), the same mask and crop window are broadcast across +every non-spatial plane, so the non-spatial dimensions are preserved and only +the `(y, x)` extent changes. + +## Examples + +```sql +SELECT RS_Width( + RS_Clip( + RS_Example(), + 1, + ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84') + ) +); +``` + +```sql +SELECT RS_Width( + RS_Clip( + RS_Example(), + 1, + ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), + false, + 0.0, + false + ) +); +``` diff --git a/rust/sedona-raster-gdal/Cargo.toml b/rust/sedona-raster-gdal/Cargo.toml index a8899f4e4..4d3889c05 100644 --- a/rust/sedona-raster-gdal/Cargo.toml +++ b/rust/sedona-raster-gdal/Cargo.toml @@ -42,6 +42,7 @@ sedona-common = { workspace = true } sedona-expr = { workspace = true } sedona-functions = { workspace = true } sedona-gdal = { workspace = true } +sedona-proj = { workspace = true } sedona-raster = { workspace = true } sedona-raster-functions = { workspace = true } sedona-schema = { workspace = true } @@ -54,10 +55,16 @@ bindgen = ["sedona-gdal/bindgen"] [dev-dependencies] criterion = { workspace = true } sedona-gdal = { workspace = true, features = ["gdal-sys"] } -sedona-testing = { workspace = true } +sedona-proj = { workspace = true, features = ["proj-sys"] } +sedona-testing = { workspace = true, features = ["criterion"] } tempfile = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } +[[bench]] +harness = false +name = "rs_clip" +path = "benches/rs_clip.rs" + [[bench]] harness = false name = "rs_frompath" diff --git a/rust/sedona-raster-gdal/benches/rs_clip.rs b/rust/sedona-raster-gdal/benches/rs_clip.rs new file mode 100644 index 000000000..d499bec14 --- /dev/null +++ b/rust/sedona-raster-gdal/benches/rs_clip.rs @@ -0,0 +1,140 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for the RS_Clip UDF. +//! +//! RS_Clip rasterizes the clip geometry into a mask, sets pixels outside it to +//! nodata, and (by default) crops to the geometry's bounding box. +//! +//! Each case builds a raster whose world extent is exactly the clip-polygon +//! generator's `[-10, 10]²` bounds at the requested resolution, so every +//! generated polygon lands on the raster and the full mask/crop path runs. +//! `all_touched = true` guarantees an overlapping polygon burns at least one +//! pixel even when it is smaller than a cell (otherwise a sub-pixel polygon +//! would produce an empty mask and hit the no-intersection early return — which +//! is what an earlier version of this benchmark accidentally measured, because +//! it placed the rasters far outside the polygons entirely). +//! +//! Axes: +//! - **Raster resolution** (`64²`, `256²`, `1024²`) with a small polygon: cost +//! is dominated by O(width × height) mask handling (rasterize + mask scan). +//! - **Clip polygon complexity** (vertex count) at a fixed resolution, driving +//! the GDAL rasterization cost. +//! - **Large clip**: a polygon covering most of the raster, so the mask/crop +//! copy (`apply_mask_and_crop`) dominates rather than mask rasterization. + +use std::sync::Arc; + +use arrow_array::{ArrayRef, BinaryArray, BooleanArray, Int32Array}; +use arrow_schema::DataType; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion_expr::ScalarUDF; +use sedona_schema::datatypes::{SedonaType, RASTER, WKB_GEOMETRY}; +use sedona_testing::{ + benchmark_util::BenchmarkArgSpec, create::make_wkb, raster_spec::RasterSpec, + testers::ScalarUdfTester, +}; + +fn criterion_benchmark(c: &mut Criterion) { + let f = sedona_raster_gdal::register::default_function_set(); + let udf: ScalarUDF = f + .scalar_udf("rs_clip") + .expect("rs_clip is registered") + .clone() + .into(); + + // RS_Clip(raster, band, geom, all_touched). + let tester = ScalarUdfTester::new( + udf, + vec![ + RASTER, + SedonaType::Arrow(DataType::Int32), + WKB_GEOMETRY, + SedonaType::Arrow(DataType::Boolean), + ], + ); + + let band: ArrayRef = Arc::new(Int32Array::from(vec![1])); + let all_touched: ArrayRef = Arc::new(BooleanArray::from(vec![true])); + + // A north-up raster covering exactly the polygon generator's [-10, 10]² + // bounds at the requested resolution, so every generated polygon overlaps. + let build_raster = |w: i64, h: i64| -> ArrayRef { + let transform = [-10.0, 20.0 / w as f64, 0.0, 10.0, 0.0, -20.0 / h as f64]; + let values = vec![1u8; (w * h) as usize]; + Arc::new( + RasterSpec::d2(w, h) + .band_values(&values) + .crs(None) + .transform(transform) + .build(), + ) + }; + + // `c` is passed in rather than captured, so several cases can share it. + let run = |c: &mut Criterion, label: &str, raster: ArrayRef, geom: ArrayRef| { + c.bench_function(label, |b| { + b.iter(|| { + tester + .invoke_arrays(vec![ + raster.clone(), + band.clone(), + geom.clone(), + all_touched.clone(), + ]) + .unwrap() + }) + }); + }; + + let gen_polygon = |vertices: usize| -> ArrayRef { + BenchmarkArgSpec::Polygon(vertices) + .build_arrays(0, 1, 1) + .expect("build clip polygon") + .remove(0) + }; + + // Resolution sweep (simple 8-vertex polygon). + for (w, h) in [(64i64, 64i64), (256, 256), (1024, 1024)] { + let label = format!("raster-gdal rs_clip Clip(Raster({w}x{h}), Polygon(8))"); + run(c, &label, build_raster(w, h), gen_polygon(8)); + } + + // Polygon-complexity axis at a fixed 64×64 resolution. + run( + c, + "raster-gdal rs_clip Clip(Raster(64x64), Polygon(50))", + build_raster(64, 64), + gen_polygon(50), + ); + + // Large clip: the polygon covers nearly the whole raster, so the crop copy + // (apply_mask_and_crop) dominates — this is where its row memcpy matters. + let big_geom: ArrayRef = Arc::new(BinaryArray::from_iter_values([make_wkb( + "POLYGON ((-9.5 -9.5, 9.5 -9.5, 9.5 9.5, -9.5 9.5, -9.5 -9.5))", + ) + .as_slice()])); + run( + c, + "raster-gdal rs_clip Clip(Raster(1024x1024), Polygon(large))", + build_raster(1024, 1024), + big_geom, + ); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/rust/sedona-raster-gdal/src/lib.rs b/rust/sedona-raster-gdal/src/lib.rs index 6682d41f6..ffcfc5b3e 100644 --- a/rust/sedona-raster-gdal/src/lib.rs +++ b/rust/sedona-raster-gdal/src/lib.rs @@ -33,6 +33,7 @@ mod gdal_common; mod gdal_dataset_provider; mod raster_loader; +mod rs_clip; mod rs_frompath; mod rs_metadata; mod source_uri; @@ -44,6 +45,7 @@ pub use gdal_common::{ nodata_bytes_to_f64, nodata_f64_to_bytes, GdalBandLayout, GdalBandPlan, }; pub use raster_loader::{GdalLoader, GDAL_FORMAT}; +pub use rs_clip::rs_clip_udf; pub use rs_frompath::rs_frompath_udf; pub use rs_metadata::rs_metadata_udf; pub use utils::{ diff --git a/rust/sedona-raster-gdal/src/register.rs b/rust/sedona-raster-gdal/src/register.rs index f4e792268..8362f0e52 100644 --- a/rust/sedona-raster-gdal/src/register.rs +++ b/rust/sedona-raster-gdal/src/register.rs @@ -20,6 +20,7 @@ use sedona_expr::function_set::FunctionSet; /// Export the set of GDAL-backed functions defined in this crate. pub fn default_function_set() -> FunctionSet { let mut function_set = FunctionSet::new(); + function_set.insert_scalar_udf(crate::rs_clip::rs_clip_udf()); function_set.insert_scalar_udf(crate::rs_frompath::rs_frompath_udf()); function_set.insert_scalar_udf(crate::rs_metadata::rs_metadata_udf()); function_set diff --git a/rust/sedona-raster-gdal/src/rs_clip.rs b/rust/sedona-raster-gdal/src/rs_clip.rs new file mode 100644 index 000000000..0999823ad --- /dev/null +++ b/rust/sedona-raster-gdal/src/rs_clip.rs @@ -0,0 +1,1491 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! RS_Clip UDF - Clip a raster to a geometry boundary +//! +//! Similar to PostGIS ST_Clip, this function clips a raster to the extent of a geometry. +//! Pixels outside the geometry are set to a nodata value: the explicit `no_data_value` +//! argument if given, otherwise the band's own nodata value, otherwise the minimum value +//! of the band's data type (so masked pixels stay distinguishable from real data). + +use std::sync::Arc; + +use arrow_array::{Array, ArrayRef}; +use datafusion_common::cast::{as_boolean_array, as_float64_array, as_int32_array}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::error::Result; +use datafusion_common::exec_err; +use datafusion_common::{exec_datafusion_err, ScalarValue}; +use datafusion_expr::{ColumnarValue, Volatility}; +use sedona_common::sedona_internal_err; +use sedona_gdal::gdal::Gdal; +use sedona_gdal::geo_transform::{GeoTransform, GeoTransformEx}; +use sedona_gdal::mem::MemDatasetBuilder; +use sedona_gdal::raster::types::GdalDataType; +use sedona_gdal::vector::geometry::Geometry; +use sedona_proj::transform::with_global_proj_engine; + +use arrow_schema::DataType; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_raster::array::RasterRefImpl; +use sedona_raster::builder::RasterBuilder; +use sedona_raster::traits::{is_spatial_dim_pair, RasterRef}; +use sedona_raster_functions::crs_utils::{crs_transform_wkb, resolve_crs}; +use sedona_raster_functions::rs_ensure_loaded::{ + NEEDS_PIXELS_METADATA_KEY, RETURNS_BYTES_METADATA_KEY, +}; +use sedona_raster_functions::RasterExecutor; +use sedona_schema::datatypes::{SedonaType, RASTER}; +use sedona_schema::matchers::ArgMatcher; +use sedona_schema::raster::BandDataType; + +use crate::gdal_common::with_gdal; +use crate::gdal_common::{nodata_bytes_to_f64, nodata_f64_to_bytes}; +use crate::gdal_dataset_provider::configure_thread_local_options; + +/// RS_Clip() scalar UDF implementation +/// +/// Clips a raster to a geometry boundary. +/// +/// Signatures: +/// - `RS_Clip(raster, band, geom)` — 3 args +/// - `RS_Clip(raster, band, geom, allTouched)` — 4 args +/// - `RS_Clip(raster, band, geom, allTouched, noDataValue)` — 5 args +/// - `RS_Clip(raster, band, geom, allTouched, noDataValue, crop)` — 6 args +/// - `RS_Clip(raster, band, geom, allTouched, noDataValue, crop, lenient)` — 7 args +pub fn rs_clip_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_clip", + vec![ + Arc::new(RsClip { arg_count: 3 }), // (raster, band, geom) + Arc::new(RsClip { arg_count: 4 }), // (raster, band, geom, allTouched) + Arc::new(RsClip { arg_count: 5 }), // (raster, band, geom, allTouched, noDataValue) + Arc::new(RsClip { arg_count: 6 }), // (raster, band, geom, allTouched, noDataValue, crop) + Arc::new(RsClip { arg_count: 7 }), // (raster, band, geom, allTouched, noDataValue, crop, lenient) + ], + Volatility::Immutable, + ) + // Reads band pixels (so the planner materializes OutDb rasters via + // RS_EnsureLoaded first) and emits a fresh InDb raster (so its output is + // already loaded and isn't wrapped again). + .with_metadata(NEEDS_PIXELS_METADATA_KEY, "true") + .with_metadata(RETURNS_BYTES_METADATA_KEY, "true") +} + +/// Kernel implementation for RS_Clip +#[derive(Debug)] +struct RsClip { + /// Number of arguments in the matched signature (3..=7) + arg_count: usize, +} + +impl SedonaScalarKernel for RsClip { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let matchers = match self.arg_count { + 3 => vec![ + // RS_Clip(raster, band, geom) + ArgMatcher::is_raster(), + ArgMatcher::is_integer(), + ArgMatcher::is_geometry_or_geography(), + ], + 4 => vec![ + // RS_Clip(raster, band, geom, allTouched) + ArgMatcher::is_raster(), + ArgMatcher::is_integer(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_boolean(), + ], + 5 => vec![ + // RS_Clip(raster, band, geom, allTouched, noDataValue) + ArgMatcher::is_raster(), + ArgMatcher::is_integer(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_boolean(), + ArgMatcher::is_numeric(), + ], + 6 => vec![ + // RS_Clip(raster, band, geom, allTouched, noDataValue, crop) + ArgMatcher::is_raster(), + ArgMatcher::is_integer(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_boolean(), + ArgMatcher::is_numeric(), + ArgMatcher::is_boolean(), + ], + 7 => vec![ + // RS_Clip(raster, band, geom, allTouched, noDataValue, crop, lenient) + ArgMatcher::is_raster(), + ArgMatcher::is_integer(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_boolean(), + ArgMatcher::is_numeric(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ], + _ => { + return sedona_internal_err!("RS_Clip: unexpected arg_count {}", self.arg_count); + } + }; + + let matcher = ArgMatcher::new(matchers, RASTER); + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + self.invoke_batch_from_args(arg_types, args, &SedonaType::Arrow(DataType::Null), 0, None) + } + + fn invoke_batch_from_args( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + _return_type: &SedonaType, + _num_rows: usize, + config_options: Option<&ConfigOptions>, + ) -> Result { + let num_iterations = calc_num_iterations(args); + + // Band is always at index 1, geom is always at index 2. + let geom_arg_idx: usize = 2; + + // Expand band to array + let band_array = args[1] + .clone() + .cast_to(&arrow_schema::DataType::Int32, None)? + .into_array(num_iterations)?; + let band_array = as_int32_array(&band_array)?.clone(); + + // allTouched at index 3 (when arg_count >= 4) + let all_touched_array = if self.arg_count >= 4 { + args[3] + .clone() + .cast_to(&arrow_schema::DataType::Boolean, None)? + .into_array(num_iterations)? + } else { + ScalarValue::Boolean(Some(false)).to_array_of_size(num_iterations)? + }; + let all_touched_array = as_boolean_array(&all_touched_array)?.clone(); + + // noDataValue at index 4 (when arg_count >= 5) + let nodata_array = if self.arg_count >= 5 { + args[4] + .clone() + .cast_to(&arrow_schema::DataType::Float64, None)? + .into_array(num_iterations)? + } else { + ScalarValue::Float64(None).to_array_of_size(num_iterations)? + }; + let nodata_array = as_float64_array(&nodata_array)?.clone(); + + // crop at index 5 (when arg_count >= 6), default true + let crop_array = if self.arg_count >= 6 { + args[5] + .clone() + .cast_to(&arrow_schema::DataType::Boolean, None)? + .into_array(num_iterations)? + } else { + ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)? + }; + let crop_array = as_boolean_array(&crop_array)?.clone(); + + // lenient at index 6 (when arg_count >= 7), default true + let lenient_array = if self.arg_count >= 7 { + args[6] + .clone() + .cast_to(&arrow_schema::DataType::Boolean, None)? + .into_array(num_iterations)? + } else { + ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)? + }; + let lenient_array = as_boolean_array(&lenient_array)?.clone(); + + let mut band_iter = band_array.iter(); + let mut all_touched_iter = all_touched_array.iter(); + let mut nodata_iter = nodata_array.iter(); + let mut crop_iter = crop_array.iter(); + let mut lenient_iter = lenient_array.iter(); + + // Build output rasters + let mut builder = RasterBuilder::new(num_iterations); + + let exec_arg_types = vec![arg_types[0].clone(), arg_types[geom_arg_idx].clone()]; + let exec_args = vec![args[0].clone(), args[geom_arg_idx].clone()]; + let executor = + RasterExecutor::new_with_num_iterations(&exec_arg_types, &exec_args, num_iterations); + + with_gdal(|gdal| { + configure_thread_local_options(gdal, config_options)?; + with_global_proj_engine(|engine| { + executor.execute_raster_wkb_crs_void(|raster_opt, wkb_opt, geom_crs| { + // Advance every option iterator in lockstep. A NULL in the band + // or any boolean flag propagates to a NULL output row (SQL + // semantics) — in particular a NULL band must NOT fall through + // to the "band 0 = all bands" mode. `no_data_value` is the + // exception: its NULL is the meaningful "not supplied" sentinel + // (the 3/4-arg signatures pass NULL here), so it stays an Option. + let band = band_iter.next().flatten(); + let all_touched = all_touched_iter.next().flatten(); + let nodata_value = nodata_iter.next().flatten(); + let crop = crop_iter.next().flatten(); + let lenient = lenient_iter.next().flatten(); + let (Some(band), Some(all_touched), Some(crop), Some(lenient)) = + (band, all_touched, crop, lenient) + else { + builder.append_null()?; + return Ok(()); + }; + + let (raster, geom_wkb) = match (raster_opt, wkb_opt) { + (Some(r), Some(w)) => (r, w), + _ => { + builder.append_null()?; + return Ok(()); + } + }; + + let raster_crs = resolve_crs(raster.crs())?; + let geom_wkb = match (geom_crs, raster_crs.as_deref()) { + (Some(geom_crs), Some(raster_crs)) => { + crs_transform_wkb(geom_wkb, geom_crs, raster_crs, engine)? + } + (None, None) => geom_wkb.to_vec(), + (Some(_), None) => { + return exec_err!( + "Cannot operate on geometry and raster: raster has no CRS but geometry does" + ) + } + (None, Some(_)) => { + return exec_err!( + "Cannot operate on geometry and raster: geometry has no CRS but raster does" + ) + } + }; + + // Band 0 means "all bands" (handled in clip_raster, which also + // range-checks the upper bound). A negative band is an error, + // not a silent clamp to band 1. + if band < 0 { + return exec_err!( + "RS_Clip: band must be >= 0 (0 = all bands), got {band}" + ); + } + let band_index = band as usize; + match clip_raster( + gdal, + raster, + &geom_wkb, + band_index, + nodata_value, + all_touched, + crop, + ) { + Ok(Some(clipped_data)) => { + build_clipped_raster(&mut builder, raster, &clipped_data)? + } + Ok(None) => { + // The clip mask is empty — no pixel was selected. `lenient` + // yields NULL either way; strict errors, with a message + // conditioned on `all_touched`: when it is already true an + // empty mask means the geometry is genuinely disjoint, + // whereas with the default `all_touched = false` the + // geometry may still overlap but fall between pixel centers. + if lenient { + builder.append_null()?; + } else if all_touched { + return exec_err!("RS_Clip: raster and geometry do not intersect"); + } else { + return exec_err!( + "RS_Clip: geometry selects no pixels; it may fall between \ + pixel centers — pass all_touched => true to keep any pixel it touches" + ); + } + } + // A genuine failure (malformed WKB, GDAL error, …) always + // propagates — it is not the no-intersection case `lenient` + // is meant to soften. + Err(e) => return Err(e), + } + + Ok(()) + }) + })?; + + // Decide array-vs-scalar over *all* args, not just the raster/geom + // the executor was given: a per-row band/option column over a scalar + // raster+geom must still yield an N-row array. (`executor.finish` + // only inspects its two exec args and would collapse to row 0.) + let out: ArrayRef = Arc::new(builder.finish()?); + if args.iter().any(|a| matches!(a, ColumnarValue::Array(_))) { + Ok(ColumnarValue::Array(out)) + } else { + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(&out, 0)?)) + } + }) + } +} + +/// One clipped band: masked/cropped bytes plus the N-D layout needed to rebuild +/// it. The clip is a 2-D `(y, x)` operation broadcast across every non-spatial +/// plane, so `dim_names` are unchanged from the source and only the trailing +/// `(y, x)` extent of `shape` shrinks when cropping. +struct ClippedBand { + /// Masked/cropped bytes, plane-major in the band's dim order. + data: Vec, + /// Visible dim names, unchanged from the source band (e.g. `["time","y","x"]`). + dim_names: Vec, + /// Output shape: leading non-spatial dims unchanged, trailing `(y, x)` clipped. + shape: Vec, + data_type: BandDataType, + /// nodata sentinel bytes written for masked-out pixels. + nodata: Vec, + /// Source band name, preserved so clipping a named/N-D band keeps it + /// addressable by name on the output. + name: Option, +} + +/// Data for a clipped raster +struct ClippedRasterData { + /// One entry per processed band. + bands: Vec, + /// Crop window in pixel coordinates (col_off, row_off, width, height): + /// the geometry's envelope intersected with the raster extent, snapped + /// outward to the pixel grid. `None` means the full original raster + /// extent was kept (crop=false). + crop_window: Option, +} + +/// A rectangular crop window in pixel coordinates. +#[derive(Debug, Clone, Copy)] +struct CropWindow { + col_off: usize, + row_off: usize, + width: usize, + height: usize, +} + +/// Clip a raster to a geometry. +/// +/// Returns `Ok(None)` when the geometry does not intersect the raster extent +/// (caller decides how to handle based on `lenient`). +fn clip_raster( + gdal: &Gdal, + raster: &RasterRefImpl<'_>, + geom_wkb: &[u8], + band_num: usize, + custom_nodata: Option, + all_touched: bool, + crop: bool, +) -> Result> { + let metadata = raster.metadata(); + let bands = raster.bands(); + let width = metadata.width() as usize; + let height = metadata.height() as usize; + + // Parse geometry from WKB + let geometry = gdal + .geometry_from_wkb(geom_wkb) + .map_err(|e| exec_datafusion_err!("Failed to parse geometry from WKB: {}", e))?; + + let geotransform = [ + metadata.upper_left_x(), + metadata.scale_x(), + metadata.skew_x(), + metadata.upper_left_y(), + metadata.skew_y(), + metadata.scale_y(), + ]; + + // The clip window is the geometry's envelope intersected with the raster + // extent, snapped outward to the pixel grid — the window PostGIS ST_Clip, + // `gdalwarp -crop_to_cutline`, and Sedona Spark's RS_Clip use. Knowing it + // up front rejects disjoint geometries before any GDAL work and bounds the + // mask to the window instead of the full raster. + let Some(window) = envelope_window(&geometry, &geotransform, width, height)? else { + return Ok(None); + }; + + // Create a mask raster covering only the clip window, with the geotransform + // shifted to the window's upper-left corner. + let mask_dataset = + MemDatasetBuilder::create(gdal, window.width, window.height, 1, GdalDataType::UInt8) + .map_err(|e| exec_datafusion_err!("Failed to create mask dataset: {}", e))?; + let (window_ulx, window_uly) = geotransform.apply(window.col_off as f64, window.row_off as f64); + let mask_geotransform = [ + window_ulx, + geotransform[1], + geotransform[2], + window_uly, + geotransform[4], + geotransform[5], + ]; + mask_dataset + .set_geo_transform(&mask_geotransform) + .map_err(|e| exec_datafusion_err!("Failed to set geotransform: {}", e))?; + + // GDAL's MEM driver zero-fills owned band buffers at creation, so the mask + // already reads 0 (outside) everywhere; rasterize_affine burns 1 inside the + // geometry. No explicit zero-init write needed. + gdal.rasterize_affine( + &mask_dataset, + &[1], // band 1 + &[geometry], + &[1.0], // burn value = 1 (inside) + all_touched, + ) + .map_err(|e| exec_datafusion_err!("Failed to rasterize geometry: {}", e))?; + + // Read the (window-sized) mask + let mask_band = mask_dataset + .rasterband(1) + .map_err(|e| exec_datafusion_err!("Failed to get mask band: {}", e))?; + let mask_buffer = mask_band + .read_as::( + (0, 0), + (window.width, window.height), + (window.width, window.height), + None, + ) + .map_err(|e| exec_datafusion_err!("Failed to read mask: {}", e))?; + let mask = mask_buffer.data(); + + // The envelope may overlap the raster while the geometry itself selects no + // pixel (e.g. it falls between pixel centers); that is still the + // no-intersection case. + let has_intersection = mask.iter().any(|&v| v != 0); + if !has_intersection { + return Ok(None); + } + + let crop_window = if crop { Some(window) } else { None }; + + // Determine which bands to process + let band_indices: Vec = if band_num == 0 { + (1..=bands.len()).collect() + } else { + if band_num > bands.len() { + return exec_err!("Band {} is out of range (1-{})", band_num, bands.len()); + } + vec![band_num] + }; + + // Process each band. The clip is a 2-D (y, x) operation; for an N-D band + // (extra leading dims such as time) the same mask and crop window are + // broadcast across every non-spatial plane. + let mut clipped_bands = Vec::with_capacity(band_indices.len()); + + for &band_idx in &band_indices { + let band = bands + .band(band_idx) + .map_err(|e| exec_datafusion_err!("Failed to get band {}: {}", band_idx, e))?; + // `band_idx` is 1-based; the `band_name` accessor is 0-based. + let band_name = raster.band_name(band_idx - 1).map(|s| s.to_string()); + + let band_metadata = band.metadata(); + let data_type = band_metadata.data_type()?; + + // The trailing two axes are the spatial (y, x) plane; anything before + // them is a stack of planes the 2-D clip is broadcast over. + let shape = band.shape().to_vec(); + let dim_names: Vec = band.dim_names().iter().map(|s| s.to_string()).collect(); + let ndim = shape.len(); + if ndim < 2 { + return exec_err!( + "RS_Clip: band {} has {} dimension(s); a 2-D (y, x) plane is required", + band_idx, + ndim + ); + } + // The trailing two dims must actually be the spatial (y, x) pair. A band + // whose trailing dims are e.g. ["x","y"] is legal at the format level and + // would otherwise be masked/cropped transposed on a square raster; mirror + // the `is_spatial_dim_pair` guard the GDAL bridge (gdal_common) performs. + if !is_spatial_dim_pair(&dim_names[ndim - 2], &dim_names[ndim - 1]) { + return exec_err!( + "RS_Clip: band {} trailing dims {:?} are not a (y, x) spatial pair", + band_idx, + &dim_names[ndim - 2..] + ); + } + let (plane_h, plane_w) = (shape[ndim - 2] as usize, shape[ndim - 1] as usize); + if plane_w != width || plane_h != height { + return exec_err!( + "RS_Clip: band {} spatial extent {}x{} does not match the raster {}x{}", + band_idx, + plane_w, + plane_h, + width, + height + ); + } + let n_planes: usize = shape[..ndim - 2].iter().map(|&d| d as usize).product(); + + // `as_contiguous` borrows the band bytes; we only ever read them here + // (the mask/crop helpers allocate their own output), so no copy is needed. + let nd_buffer = band.nd_buffer().map_err(|e| { + exec_datafusion_err!("RS_Clip: failed to read band {}: {}", band_idx, e) + })?; + let original_data = nd_buffer.as_contiguous().map_err(|e| { + exec_datafusion_err!("RS_Clip: band {} is not contiguous: {}", band_idx, e) + })?; + + // A custom nodata that isn't representable in the band's data type would + // be silently saturated by `nodata_f64_to_bytes` (e.g. -9999 -> 0 on + // UInt8, colliding with real data), so reject it explicitly here. + if let Some(cn) = custom_nodata { + let (lo, hi) = ( + band_data_type_min(&data_type), + band_data_type_max(&data_type), + ); + if cn < lo || cn > hi { + return exec_err!( + "RS_Clip: no_data_value {cn} is outside the representable range [{lo}, {hi}] of band {band_idx}'s data type {data_type:?}" + ); + } + } + + // nodata precedence: the explicit argument, then the band's own nodata, + // then the band data type's minimum value — never a silent 0.0, which + // would be indistinguishable from real zero-valued pixels. + let nodata = custom_nodata + .or_else(|| nodata_bytes_to_f64(band_metadata.nodata_value(), &data_type)) + .unwrap_or_else(|| band_data_type_min(&data_type)); + + let byte_size = data_type.byte_size(); + let in_plane_bytes = width * height * byte_size; + if original_data.len() != n_planes * in_plane_bytes { + return exec_err!( + "RS_Clip: band {} byte length {} does not match {} planes of {}x{}", + band_idx, + original_data.len(), + n_planes, + width, + height + ); + } + + // Apply the (shared) mask/crop to each plane, then concatenate. + let out_plane_bytes = match crop_window { + Some(cw) => cw.width * cw.height * byte_size, + None => in_plane_bytes, + }; + let mut clipped_data = Vec::with_capacity(n_planes * out_plane_bytes); + for plane in 0..n_planes { + let plane_bytes = &original_data[plane * in_plane_bytes..(plane + 1) * in_plane_bytes]; + let clipped_plane = if let Some(cw) = crop_window { + apply_mask_and_crop(plane_bytes, mask, width, &data_type, nodata, &cw)? + } else { + apply_mask_to_band( + plane_bytes, + mask, + width, + height, + &data_type, + nodata, + &window, + )? + }; + clipped_data.extend_from_slice(&clipped_plane); + } + + // Output shape: leading dims unchanged; trailing (y, x) becomes the crop + // window when cropping, else the original plane extent. + let (out_h, out_w) = match crop_window { + Some(cw) => (cw.height as i64, cw.width as i64), + None => (height as i64, width as i64), + }; + let mut out_shape = shape[..ndim - 2].to_vec(); + out_shape.push(out_h); + out_shape.push(out_w); + + clipped_bands.push(ClippedBand { + data: clipped_data, + dim_names, + shape: out_shape, + data_type, + nodata: nodata_f64_to_bytes(nodata, &data_type), + name: band_name, + }); + } + + Ok(Some(ClippedRasterData { + bands: clipped_bands, + crop_window, + })) +} + +/// Compute the clip window: the geometry's envelope intersected with the +/// raster extent, snapped outward to the pixel grid. Returns `None` when the +/// envelope is disjoint from the raster extent (no clipping possible). +/// +/// The envelope corners are mapped through the inverse geotransform (all four, +/// so a skewed/rotated raster still gets a correct superset window) and the +/// resulting pixel-space bbox is floored/ceiled to whole pixels. A degenerate +/// envelope (point/line) landing exactly on a grid line is widened to one +/// pixel so the rasterizer — not the snapping — decides whether it burns. +fn envelope_window( + geometry: &Geometry, + geotransform: &GeoTransform, + width: usize, + height: usize, +) -> Result> { + let env = geometry.envelope(); + let inverse = geotransform + .invert() + .map_err(|e| exec_datafusion_err!("RS_Clip: geotransform is not invertible: {}", e))?; + + let corners = [ + (env.MinX, env.MinY), + (env.MinX, env.MaxY), + (env.MaxX, env.MinY), + (env.MaxX, env.MaxY), + ]; + let mut min_col = f64::INFINITY; + let mut max_col = f64::NEG_INFINITY; + let mut min_row = f64::INFINITY; + let mut max_row = f64::NEG_INFINITY; + for (x, y) in corners { + let (col, row) = inverse.apply(x, y); + min_col = min_col.min(col); + max_col = max_col.max(col); + min_row = min_row.min(row); + max_row = max_row.max(row); + } + + let col0 = min_col.floor(); + let row0 = min_row.floor(); + let col1 = max_col.ceil().max(col0 + 1.0); + let row1 = max_row.ceil().max(row0 + 1.0); + + // Intersect with the raster extent. `>=` also rejects the NaN envelope of + // an empty geometry. + let col0 = col0.max(0.0); + let row0 = row0.max(0.0); + let col1 = col1.min(width as f64); + let row1 = row1.min(height as f64); + if !(col0 < col1 && row0 < row1) { + return Ok(None); + } + + Ok(Some(CropWindow { + col_off: col0 as usize, + row_off: row0 as usize, + width: (col1 - col0) as usize, + height: (row1 - row0) as usize, + })) +} + +/// Apply mask to band data (no cropping — preserves original dimensions). +/// The mask covers only `window`; every pixel outside it is outside the +/// geometry's envelope and therefore nodata. +fn apply_mask_to_band( + original_data: &[u8], + mask: &[u8], + width: usize, + height: usize, + data_type: &BandDataType, + nodata: f64, + window: &CropWindow, +) -> Result> { + let byte_size = data_type.byte_size(); + let nodata_bytes = nodata_f64_to_bytes(nodata, data_type); + let nodata_row: Vec = nodata_bytes.repeat(width); + let mut result = vec![0u8; width * height * byte_size]; + + let row_bytes = width * byte_size; + for row in 0..height { + let dst_row = &mut result[row * row_bytes..(row + 1) * row_bytes]; + if row < window.row_off || row >= window.row_off + window.height { + dst_row.copy_from_slice(&nodata_row); + continue; + } + // Within a window row: nodata left and right of the window, source + // bytes inside it (masked-out pixels overwritten below). + let win_start = window.col_off * byte_size; + let win_end = (window.col_off + window.width) * byte_size; + dst_row[..win_start].copy_from_slice(&nodata_row[..win_start]); + dst_row[win_end..].copy_from_slice(&nodata_row[win_end..]); + let src_start = row * row_bytes + win_start; + dst_row[win_start..win_end] + .copy_from_slice(&original_data[src_start..src_start + (win_end - win_start)]); + + let mask_row_start = (row - window.row_off) * window.width; + for col in 0..window.width { + if mask[mask_row_start + col] == 0 { + let dst = win_start + col * byte_size; + dst_row[dst..dst + byte_size].copy_from_slice(&nodata_bytes); + } + } + } + + Ok(result) +} + +/// Apply mask AND crop to the given crop window in one pass. The mask is +/// window-sized (row-major over `cw.width`×`cw.height`); the source data is +/// the full raster plane. +fn apply_mask_and_crop( + original_data: &[u8], + mask: &[u8], + full_width: usize, + data_type: &BandDataType, + nodata: f64, + cw: &CropWindow, +) -> Result> { + let byte_size = data_type.byte_size(); + let nodata_bytes = nodata_f64_to_bytes(nodata, data_type); + let row_bytes = cw.width * byte_size; + let mut result = vec![0u8; cw.height * row_bytes]; + + // The crop-window columns of a source row are contiguous, so copy each row + // in one bulk memcpy (which vectorizes) rather than per pixel, then overwrite + // only the masked-out pixels with nodata. Masked pixels are written twice — + // once by the bulk copy, once here — but the bulk memcpy is far cheaper than + // the dynamic-width per-pixel copy it replaces. + for row in 0..cw.height { + let src_row = cw.row_off + row; + let src_start = (src_row * full_width + cw.col_off) * byte_size; + let dst_start = row * row_bytes; + result[dst_start..dst_start + row_bytes] + .copy_from_slice(&original_data[src_start..src_start + row_bytes]); + + let mask_row_start = row * cw.width; + for col in 0..cw.width { + if mask[mask_row_start + col] == 0 { + let dst = dst_start + col * byte_size; + result[dst..dst + byte_size].copy_from_slice(&nodata_bytes); + } + } + } + + Ok(result) +} + +/// Build the clipped raster via the N-D builder. A 2-D raster is just the +/// `["y", "x"]` case; an N-D raster keeps its non-spatial dims and only its +/// `(y, x)` extent changes when cropping. +fn build_clipped_raster( + builder: &mut RasterBuilder, + original_raster: &RasterRefImpl<'_>, + clipped_data: &ClippedRasterData, +) -> Result<()> { + // Geotransform is 2-D and shared across all planes. A crop shifts the + // upper-left by the pixel offset; scale/skew are unchanged. + // Layout: [upper_left_x, scale_x, skew_x, upper_left_y, skew_y, scale_y]. + let src = original_raster.transform(); + let transform: [f64; 6] = if let Some(cw) = clipped_data.crop_window { + let new_ulx = src[0] + cw.col_off as f64 * src[1] + cw.row_off as f64 * src[2]; + let new_uly = src[3] + cw.row_off as f64 * src[5] + cw.col_off as f64 * src[4]; + [new_ulx, src[1], src[2], new_uly, src[4], src[5]] + } else { + [src[0], src[1], src[2], src[3], src[4], src[5]] + }; + + // Spatial extent after the clip. `spatial_dims`/`spatial_shape` are kept in + // the raster's own axis order (X-first, as the readers emit), so map each + // spatial dim to its clipped size by name rather than assuming an order. + let spatial_dims = original_raster.spatial_dims(); + let spatial_shape: Vec = match clipped_data.crop_window { + None => original_raster.spatial_shape().to_vec(), + Some(cw) => { + let x_dim = original_raster.x_dim(); + spatial_dims + .iter() + .map(|&d| { + if d == x_dim { + cw.width as i64 + } else { + cw.height as i64 + } + }) + .collect() + } + }; + + builder + .start_raster_nd( + &transform, + &spatial_dims, + &spatial_shape, + original_raster.crs(), + ) + .map_err(|e| exec_datafusion_err!("Failed to start raster: {}", e))?; + + for band in &clipped_data.bands { + let dim_names: Vec<&str> = band.dim_names.iter().map(String::as_str).collect(); + builder + .start_band_nd( + band.name.as_deref(), + &dim_names, + &band.shape, + band.data_type, + Some(&band.nodata), + None, + None, + ) + .map_err(|e| exec_datafusion_err!("Failed to start band: {}", e))?; + builder.band_data_writer().append_value(&band.data); + builder + .finish_band() + .map_err(|e| exec_datafusion_err!("Failed to finish band: {}", e))?; + } + + builder + .finish_raster() + .map_err(|e| exec_datafusion_err!("Failed to finish raster: {}", e))?; + + Ok(()) +} + +/// The minimum representable value of a band data type, as `f64` — the default +/// nodata sentinel when neither an explicit value nor the band's own nodata is +/// available, so masked-out pixels stay distinguishable from real data. +fn band_data_type_min(data_type: &BandDataType) -> f64 { + match data_type { + BandDataType::UInt8 + | BandDataType::UInt16 + | BandDataType::UInt32 + | BandDataType::UInt64 => 0.0, + BandDataType::Int8 => i8::MIN as f64, + BandDataType::Int16 => i16::MIN as f64, + BandDataType::Int32 => i32::MIN as f64, + BandDataType::Int64 => i64::MIN as f64, + BandDataType::Float32 => f32::MIN as f64, + BandDataType::Float64 => f64::MIN, + } +} + +/// The maximum representable value of a band data type, as `f64` — paired with +/// [`band_data_type_min`] to range-check a custom `no_data_value` before it is +/// narrowed to the band's type (which would otherwise saturate silently). +fn band_data_type_max(data_type: &BandDataType) -> f64 { + match data_type { + BandDataType::UInt8 => u8::MAX as f64, + BandDataType::UInt16 => u16::MAX as f64, + BandDataType::UInt32 => u32::MAX as f64, + BandDataType::UInt64 => u64::MAX as f64, + BandDataType::Int8 => i8::MAX as f64, + BandDataType::Int16 => i16::MAX as f64, + BandDataType::Int32 => i32::MAX as f64, + BandDataType::Int64 => i64::MAX as f64, + BandDataType::Float32 => f32::MAX as f64, + BandDataType::Float64 => f64::MAX, + } +} + +fn calc_num_iterations(args: &[ColumnarValue]) -> usize { + for arg in args { + if let ColumnarValue::Array(array) = arg { + return array.len(); + } + } + 1 +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::StructArray; + use sedona_expr::scalar_udf::SedonaScalarKernel; + use sedona_raster::array::RasterStructArray; + use sedona_schema::crs::deserialize_crs; + use sedona_schema::datatypes::Edges; + use sedona_testing::raster_spec::RasterSpec; + + fn wkb_from_wkt(gdal: &sedona_gdal::gdal::Gdal, wkt: &str) -> Result> { + let geometry = gdal.geometry_from_wkt(wkt).unwrap(); + geometry.wkb().map_err(|e| exec_datafusion_err!("{e}")) + } + + /// A 4×2 EPSG:4326 raster, origin (0, 2), 1×1 north-up pixels — world extent + /// x ∈ [0, 4], y ∈ [0, 2]. One UInt8 band with values 1..=8 (row-major). + fn test_raster_array() -> StructArray { + RasterSpec::d2(4, 2) + .band_values(&[1u8, 2, 3, 4, 5, 6, 7, 8]) + .crs(Some("EPSG:4326")) + .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]) + .build() + } + + #[test] + fn test_rs_clip_basic() { + // crop=false: the output keeps the original extent and band byte length; + // pixels outside the clip polygon are set to nodata. + let array = test_raster_array(); + with_gdal(|gdal| { + let rasters = RasterStructArray::try_new(&array).unwrap(); + let raster = rasters.get(0).unwrap(); + + // Left half of the raster: x ∈ [0, 2], y ∈ [0, 2]. + let geom_wkb = wkb_from_wkt(gdal, "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))")?; + let clipped = clip_raster(gdal, &raster, &geom_wkb, 0, None, false, false)? + .expect("Should have intersection"); + + let original_len = raster + .bands() + .band(1) + .unwrap() + .nd_buffer() + .unwrap() + .as_contiguous() + .unwrap() + .len(); + assert!(!clipped.bands.is_empty(), "Should have at least one band"); + assert_eq!( + clipped.bands[0].data.len(), + original_len, + "Clipped band should have same size as original when crop=false" + ); + // The polygon covers cols 0-1 of both rows; cols 2-3 lie outside + // the geometry envelope entirely and must also read nodata (0, the + // UInt8 minimum) — this pins the outside-window fill. + assert_eq!(clipped.bands[0].data, vec![1u8, 2, 0, 0, 5, 6, 0, 0]); + assert!( + clipped.crop_window.is_none(), + "crop_window should be None when crop=false" + ); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn test_rs_clip_crop() { + // crop=true: the output shrinks to the geometry's bbox window. + let array = test_raster_array(); + with_gdal(|gdal| { + let rasters = RasterStructArray::try_new(&array).unwrap(); + let raster = rasters.get(0).unwrap(); + let metadata = raster.metadata(); + + // Top-left quadrant: x ∈ [0, 2], y ∈ [1, 2] — covers pixel centers + // (0.5, 1.5) and (1.5, 1.5), i.e. a 2×1 window. + let geom_wkb = wkb_from_wkt(gdal, "POLYGON((0 1, 2 1, 2 2, 0 2, 0 1))")?; + let clipped = clip_raster(gdal, &raster, &geom_wkb, 0, None, false, true)? + .expect("Should have intersection"); + let cw = clipped + .crop_window + .expect("crop_window should be set when crop=true"); + + let byte_size = clipped.bands[0].data_type.byte_size(); + assert_eq!( + clipped.bands[0].data.len(), + cw.width * cw.height * byte_size, + "Cropped band data size should match crop window" + ); + // Both window pixels are inside the polygon, so the cropped band + // holds the source values for that 2×1 window (row-major) — this + // pins the row-copy offsets in `apply_mask_and_crop`. + assert_eq!( + clipped.bands[0].data, + vec![1u8, 2], + "cropped band should keep the source pixel values in the window" + ); + assert!( + (cw.width as i64) < metadata.width(), + "Cropped width should be smaller" + ); + assert!( + (cw.height as i64) < metadata.height(), + "Cropped height should be smaller" + ); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn crop_uses_envelope_window_not_tight_mask_bbox() { + // PostGIS ST_Clip / gdalwarp -crop_to_cutline semantics: the crop + // window is the geometry's envelope ∩ raster extent snapped to the + // grid, keeping unselected envelope pixels as nodata padding — not the + // tight bbox of the selected mask pixels. + let array = test_raster_array(); + with_gdal(|gdal| { + let rasters = RasterStructArray::try_new(&array).unwrap(); + let raster = rasters.get(0).unwrap(); + + // Right triangle with envelope x ∈ [0, 3], y ∈ [0, 2] — a 3×2 + // pixel window. Only three pixel centers fall inside it: (0.5, 1.5) + // -> 1, (0.5, 0.5) -> 5, (1.5, 0.5) -> 6. The tight mask bbox + // would be 2×2; the envelope window must be 3×2 with nodata (0) + // padding on the unselected pixels. + let geom_wkb = wkb_from_wkt(gdal, "POLYGON((0 0, 3 0, 0 2, 0 0))")?; + let clipped = clip_raster(gdal, &raster, &geom_wkb, 0, None, false, true)? + .expect("Should have intersection"); + + let cw = clipped.crop_window.expect("crop_window should be set"); + assert_eq!( + (cw.col_off, cw.row_off, cw.width, cw.height), + (0, 0, 3, 2), + "crop window should be the snapped envelope, not the mask bbox" + ); + assert_eq!(clipped.bands[0].data, vec![1u8, 0, 0, 5, 6, 0]); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn test_rs_clip_no_intersection() { + let array = test_raster_array(); + with_gdal(|gdal| { + let rasters = RasterStructArray::try_new(&array).unwrap(); + let raster = rasters.get(0).unwrap(); + // Far outside the raster's [0,4]×[0,2] extent. + let geom_wkb = wkb_from_wkt( + gdal, + "POLYGON((100 100, 101 100, 101 101, 100 101, 100 100))", + )?; + let result = clip_raster(gdal, &raster, &geom_wkb, 0, None, false, true)?; + assert!(result.is_none(), "Should return None for no intersection"); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn test_rs_clip_crs_mismatch() { + // A geometry given in EPSG:3857 must be reprojected to the raster's + // EPSG:4326 before clipping, yielding the same result as the equivalent + // EPSG:4326 geometry. + let array = test_raster_array(); + + // Build the EPSG:3857 geometry with the same PROJ engine the UDF uses, + // so the test is robust to axis-order / normalization across builds. + let crs_4326 = deserialize_crs("EPSG:4326").unwrap().unwrap(); + let crs_3857 = deserialize_crs("EPSG:3857").unwrap().unwrap(); + + let geom_wkb_4326 = + with_gdal(|gdal| wkb_from_wkt(gdal, "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))")).unwrap(); + let geom_wkb_3857 = with_global_proj_engine(|engine| { + crs_transform_wkb(&geom_wkb_4326, crs_4326.as_ref(), crs_3857.as_ref(), engine) + }) + .unwrap(); + + // 3-arg variant: RS_Clip(raster, band, geom). + let kernel = RsClip { arg_count: 3 }; + let raster_scalar = ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(array))); + let band_type = SedonaType::Arrow(DataType::Int32); + let band_val = ColumnarValue::Scalar(ScalarValue::Int32(Some(1))); + + let clip_band1 = |geom_type: SedonaType, geom_wkb: Vec| -> Vec { + let result = kernel + .invoke_batch( + &[RASTER, band_type.clone(), geom_type], + &[ + raster_scalar.clone(), + band_val.clone(), + ColumnarValue::Scalar(ScalarValue::Binary(Some(geom_wkb))), + ], + ) + .unwrap(); + match result { + ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => { + let array = RasterStructArray::try_new(struct_array.as_ref()).unwrap(); + array + .get(0) + .unwrap() + .bands() + .band(1) + .unwrap() + .nd_buffer() + .unwrap() + .as_contiguous() + .unwrap() + .to_vec() + } + _ => panic!("Expected raster scalar result"), + } + }; + + let band_data_4326 = clip_band1( + SedonaType::Wkb(Edges::Planar, Some(crs_4326)), + geom_wkb_4326, + ); + let band_data_3857 = clip_band1( + SedonaType::Wkb(Edges::Planar, Some(crs_3857)), + geom_wkb_3857, + ); + + assert_eq!(band_data_4326, band_data_3857); + } + + #[test] + fn test_rs_clip_nd_broadcasts_across_planes() { + // A [time=2, y=2, x=4] raster: clipping with a 2-D polygon crops the + // (y, x) plane and broadcasts the same mask across both time planes, + // preserving the time dimension. Values 1..=16 (C order): time 0 is + // rows [1,2,3,4] / [5,6,7,8], time 1 is [9..12] / [13..16]. + let array = RasterSpec::nd(&["time", "y", "x"], &[2, 2, 4]) + .band_values(&[1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]) + .crs(Some("EPSG:4326")) + .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]) + .build(); + + let crs_4326 = deserialize_crs("EPSG:4326").unwrap().unwrap(); + let geom_wkb = + with_gdal(|gdal| wkb_from_wkt(gdal, "POLYGON((0 1, 2 1, 2 2, 0 2, 0 1))")).unwrap(); + + // RS_Clip(raster, band, geom, allTouched, noData, crop=true). + let kernel = RsClip { arg_count: 6 }; + let result = kernel + .invoke_batch( + &[ + RASTER, + SedonaType::Arrow(DataType::Int32), + SedonaType::Wkb(Edges::Planar, Some(crs_4326)), + SedonaType::Arrow(DataType::Boolean), + SedonaType::Arrow(DataType::Float64), + SedonaType::Arrow(DataType::Boolean), + ], + &[ + ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(array))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ColumnarValue::Scalar(ScalarValue::Binary(Some(geom_wkb))), + ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))), + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))), + ], + ) + .unwrap(); + + let out = match result { + ColumnarValue::Scalar(ScalarValue::Struct(s)) => s, + _ => panic!("Expected raster scalar result"), + }; + let rasters = RasterStructArray::try_new(out.as_ref()).unwrap(); + let raster = rasters.get(0).unwrap(); + let band = raster.bands().band(1).unwrap(); + + // The time dim is preserved; (y, x) is cropped to the 1×2 mask window. + assert_eq!(band.dim_names(), vec!["time", "y", "x"]); + assert_eq!(band.shape(), &[2, 1, 2]); + + // The crop window is cols 0-1 of row 0, applied to both planes: + // time 0 -> [1, 2], time 1 -> [9, 10]. + let bytes = band.nd_buffer().unwrap().as_contiguous().unwrap(); + assert_eq!(bytes, &[1u8, 2, 9, 10]); + } + + #[test] + fn test_band_data_type_min() { + // Unsigned types floor at 0; signed/float at their most-negative value. + assert_eq!(band_data_type_min(&BandDataType::UInt8), 0.0); + assert_eq!(band_data_type_min(&BandDataType::UInt64), 0.0); + assert_eq!(band_data_type_min(&BandDataType::Int8), -128.0); + assert_eq!(band_data_type_min(&BandDataType::Int16), i16::MIN as f64); + assert_eq!(band_data_type_min(&BandDataType::Int32), i32::MIN as f64); + assert_eq!(band_data_type_min(&BandDataType::Float32), f32::MIN as f64); + assert_eq!(band_data_type_min(&BandDataType::Float64), f64::MIN); + } + + /// A 2×1, two-band EPSG:4326 raster (band 1 = [1,2], band 2 = [10,20]) — the + /// two bands differ so a per-band clip is observably distinct. + fn two_band_raster() -> StructArray { + RasterSpec::d2(2, 1) + .band_values(&[1u8, 2]) + .band_values(&[10u8, 20]) + .crs(Some("EPSG:4326")) + .transform([0.0, 1.0, 0.0, 1.0, 0.0, -1.0]) + .build() + } + + #[test] + fn scalar_raster_geom_with_band_column_yields_all_rows() { + // Regression: a constant raster+geom with a per-row band column must + // produce an N-row array, not collapse to row 0. (The executor only sees + // [raster, geom], so output packaging must consider all args.) + let crs_4326 = deserialize_crs("EPSG:4326").unwrap().unwrap(); + let geom_wkb = + with_gdal(|gdal| wkb_from_wkt(gdal, "POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))")).unwrap(); + + let kernel = RsClip { arg_count: 3 }; + let result = kernel + .invoke_batch( + &[ + RASTER, + SedonaType::Arrow(DataType::Int32), + SedonaType::Wkb(Edges::Planar, Some(crs_4326)), + ], + &[ + ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(two_band_raster()))), + ColumnarValue::Array(Arc::new(arrow_array::Int32Array::from(vec![1, 2]))), + ColumnarValue::Scalar(ScalarValue::Binary(Some(geom_wkb))), + ], + ) + .unwrap(); + + // Must be a 2-row array (not a broadcast scalar), with row 0 clipping + // band 1 and row 1 clipping band 2 — distinct outputs. + let arr = match result { + ColumnarValue::Array(a) => a, + ColumnarValue::Scalar(_) => panic!("expected an array; the batch collapsed to row 0"), + }; + let rasters = + RasterStructArray::try_new(arr.as_any().downcast_ref::().unwrap()) + .unwrap(); + assert_eq!(rasters.len(), 2); + let band_data = |row: usize| { + rasters + .get(row) + .unwrap() + .bands() + .band(1) + .unwrap() + .nd_buffer() + .unwrap() + .as_contiguous() + .unwrap() + .to_vec() + }; + // Row 0 clipped band 1 (values 1,2); row 1 clipped band 2 (values 10,20). + assert_ne!(band_data(0), band_data(1)); + } + + #[test] + fn band_zero_clips_all_bands() { + // Band 0 means "all bands" — it must reach clip_raster as 0, not be + // clamped to 1. + let crs_4326 = deserialize_crs("EPSG:4326").unwrap().unwrap(); + let geom_wkb = + with_gdal(|gdal| wkb_from_wkt(gdal, "POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))")).unwrap(); + let kernel = RsClip { arg_count: 3 }; + let result = kernel + .invoke_batch( + &[ + RASTER, + SedonaType::Arrow(DataType::Int32), + SedonaType::Wkb(Edges::Planar, Some(crs_4326)), + ], + &[ + ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(two_band_raster()))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(0))), + ColumnarValue::Scalar(ScalarValue::Binary(Some(geom_wkb))), + ], + ) + .unwrap(); + let out = match result { + ColumnarValue::Scalar(ScalarValue::Struct(s)) => s, + _ => panic!("expected raster scalar"), + }; + let raster = RasterStructArray::try_new(out.as_ref()) + .unwrap() + .get(0) + .unwrap(); + assert_eq!(raster.bands().len(), 2, "band 0 should clip all bands"); + } + + #[test] + fn negative_band_errors() { + let crs_4326 = deserialize_crs("EPSG:4326").unwrap().unwrap(); + let geom_wkb = + with_gdal(|gdal| wkb_from_wkt(gdal, "POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))")).unwrap(); + let kernel = RsClip { arg_count: 3 }; + let err = kernel + .invoke_batch( + &[ + RASTER, + SedonaType::Arrow(DataType::Int32), + SedonaType::Wkb(Edges::Planar, Some(crs_4326)), + ], + &[ + ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(two_band_raster()))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(-1))), + ColumnarValue::Scalar(ScalarValue::Binary(Some(geom_wkb))), + ], + ) + .unwrap_err() + .to_string(); + assert!(err.contains("band must be >= 0"), "unexpected error: {err}"); + } + + #[test] + fn malformed_geometry_errors_even_when_lenient() { + // `lenient` (default true) softens only the no-intersection case; a + // genuine failure (garbage WKB) must still error, not become NULL. + let kernel = RsClip { arg_count: 3 }; + let err = kernel + .invoke_batch( + &[ + RASTER, + SedonaType::Arrow(DataType::Int32), + SedonaType::Wkb(Edges::Planar, None), + ], + &[ + // No CRS on raster or geom, so we reach rasterization with the + // garbage WKB rather than erroring on a CRS mismatch first. + ColumnarValue::Scalar(ScalarValue::Struct(Arc::new( + RasterSpec::d2(2, 1) + .band_values(&[1u8, 2]) + .crs(None) + .transform([0.0, 1.0, 0.0, 1.0, 0.0, -1.0]) + .build(), + ))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ColumnarValue::Scalar(ScalarValue::Binary(Some(vec![0xff, 0xff, 0xff, 0xff]))), + ], + ) + .unwrap_err(); + // The point is it errored rather than returning a NULL raster. + let _ = err; + } + + #[test] + fn null_band_yields_null() { + // A NULL band must propagate to NULL, not silently clip every band. + let crs_4326 = deserialize_crs("EPSG:4326").unwrap().unwrap(); + let geom_wkb = + with_gdal(|gdal| wkb_from_wkt(gdal, "POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))")).unwrap(); + let kernel = RsClip { arg_count: 3 }; + let result = kernel + .invoke_batch( + &[ + RASTER, + SedonaType::Arrow(DataType::Int32), + SedonaType::Wkb(Edges::Planar, Some(crs_4326)), + ], + &[ + ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(two_band_raster()))), + ColumnarValue::Scalar(ScalarValue::Int32(None)), + ColumnarValue::Scalar(ScalarValue::Binary(Some(geom_wkb))), + ], + ) + .unwrap(); + match result { + ColumnarValue::Scalar(sv) => { + assert!( + sv.is_null(), + "NULL band should yield NULL, not clip all bands" + ) + } + other => panic!("expected scalar, got {other:?}"), + } + } + + #[test] + fn custom_nodata_out_of_range_errors() { + // A UInt8 band can't represent -9999; reject it rather than saturating + // it to 0 (which would collide with real zero-valued data). + let array = test_raster_array(); // one UInt8 band + with_gdal(|gdal| { + let rasters = RasterStructArray::try_new(&array).unwrap(); + let raster = rasters.get(0).unwrap(); + let geom_wkb = wkb_from_wkt(gdal, "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))")?; + + let err = match clip_raster(gdal, &raster, &geom_wkb, 1, Some(-9999.0), false, false) { + Err(e) => e.to_string(), + Ok(_) => panic!("expected out-of-range nodata to error"), + }; + assert!( + err.contains("outside the representable range"), + "unexpected error: {err}" + ); + + // An in-range custom nodata is accepted. + let ok = clip_raster(gdal, &raster, &geom_wkb, 1, Some(200.0), false, false)?; + assert!(ok.is_some(), "in-range custom nodata should be accepted"); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn preserves_band_name() { + // Clipping a named band keeps the name on the output band. + let array = RasterSpec::d2(4, 2) + .band_values(&[1u8, 2, 3, 4, 5, 6, 7, 8]) + .name("elevation") + .crs(Some("EPSG:4326")) + .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]) + .build(); + with_gdal(|gdal| { + let rasters = RasterStructArray::try_new(&array).unwrap(); + let raster = rasters.get(0).unwrap(); + let geom_wkb = wkb_from_wkt(gdal, "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))")?; + let clipped = clip_raster(gdal, &raster, &geom_wkb, 1, None, false, false)? + .expect("Should have intersection"); + assert_eq!( + clipped.bands[0].name.as_deref(), + Some("elevation"), + "clipped band should keep the source band name" + ); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn strict_no_pixel_message_depends_on_all_touched() { + // With lenient=false and an empty clip mask, the strict error is + // conditioned on all_touched: a plain "do not intersect" when every + // touched pixel was already considered, and a sub-pixel hint otherwise. + let geom_wkb = with_gdal(|gdal| { + wkb_from_wkt( + gdal, + "POLYGON((100 100, 101 100, 101 101, 100 101, 100 100))", + ) + }) + .unwrap(); + // No CRS on raster or geom, so we reach the mask check rather than a + // CRS-mismatch error first. + let raster = RasterSpec::d2(4, 2) + .band_values(&[1u8, 2, 3, 4, 5, 6, 7, 8]) + .crs(None) + .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]) + .build(); + let arg_types = [ + RASTER, + SedonaType::Arrow(DataType::Int32), + SedonaType::Wkb(Edges::Planar, None), + SedonaType::Arrow(DataType::Boolean), + SedonaType::Arrow(DataType::Float64), + SedonaType::Arrow(DataType::Boolean), + SedonaType::Arrow(DataType::Boolean), + ]; + let kernel = RsClip { arg_count: 7 }; + let invoke = |all_touched: bool| { + kernel + .invoke_batch( + &arg_types, + &[ + ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(raster.clone()))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ColumnarValue::Scalar(ScalarValue::Binary(Some(geom_wkb.clone()))), + ColumnarValue::Scalar(ScalarValue::Boolean(Some(all_touched))), + ColumnarValue::Scalar(ScalarValue::Float64(None)), + ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))), // crop + ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))), // lenient + ], + ) + .unwrap_err() + .to_string() + }; + assert!( + invoke(true).contains("do not intersect"), + "all_touched=true should give the disjoint message" + ); + let msg = invoke(false); + assert!( + msg.contains("selects no pixels") && msg.contains("all_touched"), + "all_touched=false should hint at the sub-pixel case: {msg}" + ); + } +}