paleolimbot commented on code in PR #1066: URL: https://github.com/apache/sedona-db/pull/1066#discussion_r3690860237
########## rust/sedona-raster-gdal/src/rs_zonal_stats.rs: ########## @@ -0,0 +1,1749 @@ +// 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_ZonalStats / RS_ZonalStatsAll UDFs — summary statistics of the raster +//! pixels covered by a roi geometry. +//! +//! Both mirror Apache Sedona Spark's positional overloads verbatim so that +//! Spark SQL tends to run unchanged. `RS_ZonalStats` returns one statistic as a +//! `Float64`: +//! +//! - `RS_ZonalStats(raster, roi, stat)` +//! - `RS_ZonalStats(raster, roi, band, stat)` +//! - `RS_ZonalStats(raster, roi, band, stat, all_touched)` +//! - `RS_ZonalStats(raster, roi, band, stat, all_touched, exclude_no_data)` +//! - `RS_ZonalStats(raster, roi, band, stat, all_touched, exclude_no_data, lenient)` +//! +//! `RS_ZonalStatsAll` returns every statistic as a struct, with the same ladder +//! minus `stat`: +//! +//! - `RS_ZonalStatsAll(raster, roi)` +//! - `RS_ZonalStatsAll(raster, roi, band)` +//! - `RS_ZonalStatsAll(raster, roi, band, all_touched)` +//! - `RS_ZonalStatsAll(raster, roi, band, all_touched, exclude_no_data)` +//! - `RS_ZonalStatsAll(raster, roi, band, all_touched, exclude_no_data, lenient)` +//! +//! A pixel is included when its centre falls inside the roi (or that the roi +//! merely touches, with `all_touched`), optionally excluding the band's nodata +//! value. `all_touched` defaults to false, `exclude_no_data` to true, and +//! `lenient` to true. Unlike Sedona Spark, the band-less overloads do not +//! default to band 1 on a multiband raster: naming the band is required there +//! (a single-band raster resolves unambiguously). +//! +//! These functions operate on 2-D `(y, x)` bands. A band that is not a 2-D +//! spatial grid is rejected; computing a statistic per non-spatial plane of an +//! N-D band is not supported. + +use std::sync::Arc; + +use arrow_array::builder::{Float64Builder, Int64Builder}; +use arrow_array::{ArrayRef, BooleanArray, Int64Array, StringArray, StructArray}; +use arrow_buffer::{BooleanBufferBuilder, NullBuffer}; +use arrow_schema::{DataType, Field, Fields}; +use datafusion_common::cast::{as_boolean_array, as_int64_array, as_string_array}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::error::Result; +use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue}; +use datafusion_expr::{ColumnarValue, Volatility}; + +use sedona_common::sedona_internal_err; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_gdal::gdal::Gdal; +use sedona_raster::array::RasterRefImpl; +use sedona_raster::traits::RasterRef; +use sedona_raster_functions::crs_utils::{align_wkb_to_crs, resolve_crs, with_crs_engine}; +use sedona_raster_functions::rs_ensure_loaded::NEEDS_PIXELS_METADATA_KEY; +use sedona_raster_functions::rs_spatial_predicates::raster_intersects_geom_wkb; +use sedona_raster_functions::RasterExecutor; +use sedona_schema::datatypes::SedonaType; +use sedona_schema::matchers::ArgMatcher; +use sedona_schema::raster::BandDataType; + +use crate::gdal_common::{raster_geo_transform, with_gdal}; +use crate::gdal_dataset_provider::configure_thread_local_options; +use crate::mask::{envelope_window, rasterize_geometry_mask, PixelWindow}; + +/// The statistics RS_ZonalStatsAll returns, in the order Sedona Spark reports +/// them. RS_ZonalStats selects one of these by name. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StatType { + Count, + Sum, + Mean, + Median, + Mode, + StdDev, + Variance, + Min, + Max, +} + +impl StatType { + /// Parse a statistic name (case-insensitive). Aliases match Sedona Spark + /// (`avg`/`average` for mean, `sd` for stddev). + fn from_str(s: &str) -> Option<StatType> { + match s.to_lowercase().as_str() { + "count" => Some(StatType::Count), + "sum" => Some(StatType::Sum), + "mean" | "avg" | "average" => Some(StatType::Mean), + "median" => Some(StatType::Median), + "mode" => Some(StatType::Mode), + "stddev" | "sd" => Some(StatType::StdDev), + "variance" => Some(StatType::Variance), + "min" => Some(StatType::Min), + "max" => Some(StatType::Max), + _ => None, + } + } +} + +/// Defaults for the trailing flags, applied by the narrower overloads that omit +/// them (matching Sedona Spark). +const DEFAULT_ALL_TOUCHED: bool = false; +const DEFAULT_EXCLUDE_NODATA: bool = true; +const DEFAULT_LENIENT: bool = true; + +/// The resolved parameters for one row's zonal-stats computation, assembled from +/// the positional arguments the matched overload carried. +#[derive(Debug, Clone)] +struct ZonalStatsParams { + /// 1-based band to compute over. `None` means "resolve the implicit band": + /// band 1 for a single-band raster, an error for a multiband raster (naming + /// the band is required rather than silently getting band 1). Only the + /// band-less overloads leave this `None`. + band: Option<i64>, + /// Include every pixel the roi touches, not only those whose centre it + /// covers. + all_touched: bool, + /// Skip pixels equal to the band's nodata value. + exclude_no_data: bool, + /// Return NULL when the roi does not intersect the raster, rather than + /// erroring. Only the no-intersection case is softened; malformed geometry + /// or an unreadable band always errors. + lenient: bool, +} + +/// Every statistic for a roi. `count` is always present (0 when the roi +/// selects no pixels); the remaining fields are `None` in exactly that +/// no-pixel case and `Some` otherwise, mirroring Sedona Spark (which returns +/// `count = 0` and NULL for the rest). +#[derive(Debug, Clone, PartialEq)] +struct ZonalStatistics { + count: i64, + sum: Option<f64>, + mean: Option<f64>, + median: Option<f64>, + mode: Option<f64>, + stddev: Option<f64>, + variance: Option<f64>, + min: Option<f64>, + max: Option<f64>, +} + +/// Whether a roi geometry intersects the raster. `NoIntersection` mirrors Sedona +/// Spark's `rsIntersects` gate (the caller turns it into NULL when `lenient`, an +/// error otherwise); `Collected` means the selected pixel values are in the +/// caller's scratch buffer, possibly empty for a roi that intersects the +/// footprint but selects no pixel centre (a `count = 0` result). +enum RoiCoverage { + NoIntersection, + Collected, +} + +// ============================================================================= +// RS_ZonalStats +// ============================================================================= + +/// `RS_ZonalStats` — one statistic as a `Float64`. `stat` is a statistic name +/// (`count`, `sum`, `mean`, `median`, `mode`, `stddev`, `variance`, `min`, +/// `max`). See the module docs for the full positional overload ladder. +pub fn rs_zonal_stats_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_zonalstats", + vec![ + Arc::new(RsZonalStats { arg_count: 3 }), // (raster, roi, stat) + Arc::new(RsZonalStats { arg_count: 4 }), // (raster, roi, band, stat) + Arc::new(RsZonalStats { arg_count: 5 }), // + all_touched + Arc::new(RsZonalStats { arg_count: 6 }), // + exclude_no_data + Arc::new(RsZonalStats { arg_count: 7 }), // + lenient + ], + Volatility::Immutable, + ) + // Reads band pixels, so the planner materializes OutDb rasters via + // RS_EnsureLoaded first. + .with_metadata(NEEDS_PIXELS_METADATA_KEY, "true") +} + +#[derive(Debug)] +struct RsZonalStats { + /// Number of arguments in the matched signature (3..=7). + arg_count: usize, +} + +impl SedonaScalarKernel for RsZonalStats { + fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> { + // Argument order mirrors Sedona Spark: (raster, roi, [band,] stat, + // [all_touched, [exclude_no_data, [lenient]]]). The 3-arg overload omits + // band (its stat is at index 2); the 4+-arg overloads carry band at + // index 2 and stat at index 3. + let matchers = match self.arg_count { + 3 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_string(), + ], + 4 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ], + 5 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ArgMatcher::is_boolean(), + ], + 6 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ], + 7 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ], + _ => { + return sedona_internal_err!( + "RS_ZonalStats: unexpected arg_count {}", + self.arg_count + ); + } + }; + let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(DataType::Float64)); + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + 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<ColumnarValue> { + let num_iterations = RasterExecutor::num_iterations_over(args); + + // band (index 2) only exists in the 4+-arg overloads; the 3-arg overload + // leaves it implicit. stat is at index 2 (3-arg) or 3 (4+-arg). + let has_band = self.arg_count >= 4; + let stat_idx = if has_band { 3 } else { 2 }; + let stat_array = expand_string_arg(&args[stat_idx], num_iterations)?; + let mut stat_iter = stat_array.iter(); + + let band_array = has_band + .then(|| expand_int64_arg(&args[2], num_iterations)) + .transpose()?; + let mut band_iter = band_array.as_ref().map(|a| a.iter()); + + // all_touched (index 4), exclude_no_data (index 5), lenient (index 6): + // read from the column when the overload carries it, else the default. + let all_touched_array = expand_flag( + args, + 4, + self.arg_count >= 5, + DEFAULT_ALL_TOUCHED, + num_iterations, + )?; + let exclude_no_data_array = expand_flag( + args, + 5, + self.arg_count >= 6, + DEFAULT_EXCLUDE_NODATA, + num_iterations, + )?; + let lenient_array = expand_flag( + args, + 6, + self.arg_count >= 7, + DEFAULT_LENIENT, + num_iterations, + )?; + let mut all_touched_iter = all_touched_array.iter(); + let mut exclude_no_data_iter = exclude_no_data_array.iter(); + let mut lenient_iter = lenient_array.iter(); + + let mut builder = Float64Builder::with_capacity(num_iterations); + let mut scratch: Vec<f64> = Vec::new(); + let mut mask_scratch: Vec<u8> = Vec::new(); + + // The executor only sees (raster, roi); the option columns are advanced + // in lockstep below. + let exec_arg_types = [arg_types[0].clone(), arg_types[1].clone()]; + let exec_args = [args[0].clone(), args[1].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_crs_engine(config_options, |engine| { + executor.execute_raster_wkb_crs_void(|raster_opt, wkb_opt, geom_crs| { + let stat_str = stat_iter.next().flatten(); + let Some(params) = next_params( + &mut band_iter, + &mut all_touched_iter, + &mut exclude_no_data_iter, + &mut lenient_iter, + ) else { + builder.append_null(); + return Ok(()); + }; + + // A NULL stat, raster, or roi propagates to a NULL row. + let (Some(stat_str), Some(raster), Some(wkb)) = (stat_str, raster_opt, wkb_opt) + else { + builder.append_null(); + return Ok(()); + }; + let stat_type = StatType::from_str(stat_str).ok_or_else(|| { + exec_datafusion_err!("RS_ZonalStats: unknown statistic {stat_str:?}") + })?; + + // Reproject the roi into the raster's CRS, borrowing it + // unchanged when the CRSes already match; a CRS on exactly + // one side is an error, since it would mislocate the roi. + let raster_crs = resolve_crs(raster.crs())?; + let geom_wkb = align_wkb_to_crs( + wkb, + geom_crs, + raster_crs.as_deref(), + "geometry", + "raster", + engine, + )?; + match collect_zonal_values( + gdal, + raster, + &geom_wkb, + ¶ms, + &mut scratch, + &mut mask_scratch, + )? { + // Compute only the requested statistic, not all of them. + RoiCoverage::Collected => { + match compute_single_statistic(&mut scratch, stat_type) { + Some(value) => builder.append_value(value), + None => builder.append_null(), + } + } + // The roi does not intersect the raster: NULL when + // lenient (the default), an error otherwise. + RoiCoverage::NoIntersection if params.lenient => builder.append_null(), + RoiCoverage::NoIntersection => return no_intersection_err(), + } + Ok(()) + }) + })?; + + let out: ArrayRef = Arc::new(builder.finish()); + RasterExecutor::finish_over(args, out) + }) + } +} + +// ============================================================================= +// RS_ZonalStatsAll +// ============================================================================= + +/// `RS_ZonalStatsAll` — every statistic as a struct with fields `count, sum, +/// mean, median, mode, stddev, variance, min, max`. See the module docs for the +/// full positional overload ladder. +pub fn rs_zonal_stats_all_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_zonalstatsall", + vec![ + Arc::new(RsZonalStatsAll { arg_count: 2 }), // (raster, roi) + Arc::new(RsZonalStatsAll { arg_count: 3 }), // (raster, roi, band) + Arc::new(RsZonalStatsAll { arg_count: 4 }), // + all_touched + Arc::new(RsZonalStatsAll { arg_count: 5 }), // + exclude_no_data + Arc::new(RsZonalStatsAll { arg_count: 6 }), // + lenient + ], + Volatility::Immutable, + ) + .with_metadata(NEEDS_PIXELS_METADATA_KEY, "true") +} + +#[derive(Debug)] +struct RsZonalStatsAll { + /// Number of arguments in the matched signature (2..=6). + arg_count: usize, +} + +impl SedonaScalarKernel for RsZonalStatsAll { + fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> { + // Argument order mirrors Sedona Spark: (raster, roi, [band, + // [all_touched, [exclude_no_data, [lenient]]]]). The 2-arg overload omits + // band; the 3+-arg overloads carry it at index 2. + let mut matchers = vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ]; + if self.arg_count >= 3 { + matchers.push(ArgMatcher::is_integer()); // band + } + for _ in 4..=self.arg_count { + matchers.push(ArgMatcher::is_boolean()); // all_touched, exclude_no_data, lenient + } + if self.arg_count < 2 || self.arg_count > 6 { + return sedona_internal_err!( + "RS_ZonalStatsAll: unexpected arg_count {}", + self.arg_count + ); + } + let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(zonal_stats_struct_type())); + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + 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<ColumnarValue> { + let num_iterations = RasterExecutor::num_iterations_over(args); + + // band (index 2) only exists in the 3+-arg overloads; the 2-arg overload + // leaves it implicit. all_touched (index 3), exclude_no_data (index 4), + // and lenient (index 5) follow. + let band_array = (self.arg_count >= 3) + .then(|| expand_int64_arg(&args[2], num_iterations)) + .transpose()?; + let mut band_iter = band_array.as_ref().map(|a| a.iter()); + + let all_touched_array = expand_flag( + args, + 3, + self.arg_count >= 4, + DEFAULT_ALL_TOUCHED, + num_iterations, + )?; + let exclude_no_data_array = expand_flag( + args, + 4, + self.arg_count >= 5, + DEFAULT_EXCLUDE_NODATA, + num_iterations, + )?; + let lenient_array = expand_flag( + args, + 5, + self.arg_count >= 6, + DEFAULT_LENIENT, + num_iterations, + )?; + let mut all_touched_iter = all_touched_array.iter(); + let mut exclude_no_data_iter = exclude_no_data_array.iter(); + let mut lenient_iter = lenient_array.iter(); + + let mut builders = ZonalStatsBuilders::with_capacity(num_iterations); + let mut scratch: Vec<f64> = Vec::new(); + let mut mask_scratch: Vec<u8> = Vec::new(); + + let exec_arg_types = [arg_types[0].clone(), arg_types[1].clone()]; + let exec_args = [args[0].clone(), args[1].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_crs_engine(config_options, |engine| { + executor.execute_raster_wkb_crs_void(|raster_opt, wkb_opt, geom_crs| { + let Some(params) = next_params( + &mut band_iter, + &mut all_touched_iter, + &mut exclude_no_data_iter, + &mut lenient_iter, + ) else { + builders.push_null(); + return Ok(()); + }; + + let (Some(raster), Some(wkb)) = (raster_opt, wkb_opt) else { + builders.push_null(); + return Ok(()); + }; + + let raster_crs = resolve_crs(raster.crs())?; + let geom_wkb = align_wkb_to_crs( + wkb, + geom_crs, + raster_crs.as_deref(), + "geometry", + "raster", + engine, + )?; + match collect_zonal_values( + gdal, + raster, + &geom_wkb, + ¶ms, + &mut scratch, + &mut mask_scratch, + )? { + RoiCoverage::Collected => { + builders.push_stats(&compute_statistics(&mut scratch)) + } + RoiCoverage::NoIntersection if params.lenient => builders.push_null(), + RoiCoverage::NoIntersection => return no_intersection_err(), + } + Ok(()) + }) + })?; + + let out: ArrayRef = Arc::new(builders.finish()); + RasterExecutor::finish_over(args, out) + }) + } +} + +/// Struct data type RS_ZonalStatsAll returns. +fn zonal_stats_struct_type() -> DataType { + DataType::Struct(zonal_stats_struct_fields()) +} + +/// Fields of the RS_ZonalStatsAll struct, in Sedona Spark order. `count` is an +/// `Int64` (a whole pixel count); every other statistic is a `Float64`. +fn zonal_stats_struct_fields() -> Fields { + Fields::from(vec![ + Field::new("count", DataType::Int64, true), + Field::new("sum", DataType::Float64, true), + Field::new("mean", DataType::Float64, true), + Field::new("median", DataType::Float64, true), + Field::new("mode", DataType::Float64, true), + Field::new("stddev", DataType::Float64, true), + Field::new("variance", DataType::Float64, true), + Field::new("min", DataType::Float64, true), + Field::new("max", DataType::Float64, true), + ]) +} + +/// Column builders for the RS_ZonalStatsAll struct output: one typed builder per +/// field plus an outer struct-level validity buffer. +/// +/// Building the columns directly and assembling the [`StructArray`] once at the +/// end avoids downcasting a `StructBuilder`'s boxed field builders on every row. +struct ZonalStatsBuilders { + count: Int64Builder, + /// `sum, mean, median, mode, stddev, variance, min, max` — Sedona Spark field + /// order, matching [`zonal_stats_struct_fields`] after `count`. + floats: [Float64Builder; 8], + /// Struct-level null bitmap: `false` for a fully-NULL row (a NULL input or a + /// non-intersecting roi under `lenient`). + validity: BooleanBufferBuilder, +} + +impl ZonalStatsBuilders { + fn with_capacity(capacity: usize) -> Self { + Self { + count: Int64Builder::with_capacity(capacity), + floats: std::array::from_fn(|_| Float64Builder::with_capacity(capacity)), + validity: BooleanBufferBuilder::new(capacity), + } + } + + /// Append one computed-stats row (the struct itself is valid). The float + /// fields carry through the `Option`, so an empty roi records `count = 0` + /// with the rest NULL. + fn push_stats(&mut self, stats: &ZonalStatistics) { + self.count.append_value(stats.count); + let values = [ + stats.sum, + stats.mean, + stats.median, + stats.mode, + stats.stddev, + stats.variance, + stats.min, + stats.max, + ]; + for (builder, value) in self.floats.iter_mut().zip(values) { + builder.append_option(value); + } + self.validity.append(true); + } + + /// Append a fully-NULL struct row (a NULL input, or a non-intersecting roi + /// when `lenient`). Every field is null and the struct itself is null. + fn push_null(&mut self) { + self.count.append_null(); + for builder in &mut self.floats { + builder.append_null(); + } + self.validity.append(false); + } + + /// Assemble the accumulated columns into the struct array. + fn finish(mut self) -> StructArray { + let mut arrays: Vec<ArrayRef> = Vec::with_capacity(9); + arrays.push(Arc::new(self.count.finish())); + for builder in &mut self.floats { + arrays.push(Arc::new(builder.finish())); + } + let nulls = NullBuffer::new(self.validity.finish()); + StructArray::new(zonal_stats_struct_fields(), arrays, Some(nulls)) + } +} + +// ============================================================================= +// Core computation +// ============================================================================= + +/// Collect the pixel values a roi geometry selects on one band into `scratch`. +/// +/// Returns [`RoiCoverage::NoIntersection`] when the roi geometry does not +/// intersect the raster's footprint — a true geometry intersection (matching +/// Sedona Spark's `rsIntersects` gate), not a bounding-box overlap: a roi whose +/// envelope overlaps the raster but whose geometry is disjoint is a +/// no-intersection case. The caller turns that into NULL when `lenient`, an +/// error otherwise. A roi that intersects the footprint but whose selected +/// pixels are all outside the geometry or all nodata returns +/// [`RoiCoverage::Collected`] with `scratch` left empty (a `count = 0` result). +/// +/// The caller computes the statistic(s) it needs from `scratch` — every one for +/// RS_ZonalStatsAll, only the requested one for RS_ZonalStats — so this shared +/// collection never computes statistics the caller would discard. +/// +/// `scratch` is a reused buffer for the selected pixel values and `mask_scratch` +/// for the rasterized roi mask, both reused so the per-row computation does not +/// allocate a fresh `Vec` each call. +fn collect_zonal_values( + gdal: &Gdal, + raster: &RasterRefImpl<'_>, + geom_wkb: &[u8], + params: &ZonalStatsParams, + scratch: &mut Vec<f64>, + mask_scratch: &mut Vec<u8>, +) -> Result<RoiCoverage> { + let num_bands = raster.num_bands(); + let band_num = resolve_band(params.band, num_bands)?; + + let band = raster + .bands() + .band(band_num) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to read band {band_num}: {e}"))?; + if !band.is_spatial_2d() { + return exec_err!( + "RS_ZonalStats supports 2-D rasters only; band {band_num} is not a 2-D (y, x) grid" + ); + } + let data_type = band.data_type(); + let byte_size = data_type.byte_size(); + + let metadata = raster.metadata(); + let transform = raster_geo_transform(raster)?; + let width = usize::try_from(metadata.width()) + .map_err(|_| exec_datafusion_err!("RS_ZonalStats: negative raster width"))?; + let height = usize::try_from(metadata.height()) + .map_err(|_| exec_datafusion_err!("RS_ZonalStats: negative raster height"))?; + + // No-intersection gate: a true geometry intersection between the roi and + // the raster footprint (matching Sedona Spark's rsIntersects gate), not a + // bounding-box overlap. A roi whose envelope overlaps the raster but whose + // geometry is disjoint is a no-intersection case, not a count-0 case. The + // roi is already in the raster's CRS here, so no transform is needed. + if !raster_intersects_geom_wkb(raster, geom_wkb)? { + return Ok(RoiCoverage::NoIntersection); + } + + // Parse the roi and clamp its envelope to the raster grid for the pixel + // window to rasterize. The gate above already established overlap; a + // degenerate window (the roi only touches the raster boundary) selects no + // pixels, so it is count 0 rather than no-intersection. + let geometry = gdal + .geometry_from_wkb(geom_wkb) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to parse geometry: {e}"))?; + let Some(window) = envelope_window(&geometry, &transform, width, height)? else { + scratch.clear(); + return Ok(RoiCoverage::Collected); + }; + + // Rasterize the roi into a window-sized 0/1 mask (moves `geometry`, whose + // only remaining use is the burn). The mask reuses `mask_scratch` across + // rows rather than allocating a fresh buffer each call. + rasterize_geometry_mask( + gdal, + geometry, + &transform, + &window, + params.all_touched, + mask_scratch, + )?; + + // Read the band once (zero-copy borrow) and collect the selected values. + let nd_buffer = band + .nd_buffer() + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to read band {band_num}: {e}"))?; + let band_bytes = nd_buffer.as_contiguous().map_err(|e| { + exec_datafusion_err!("RS_ZonalStats: band {band_num} is not contiguous: {e}") + })?; + let expected = width + .checked_mul(height) + .and_then(|n| n.checked_mul(byte_size)) + .ok_or_else(|| exec_datafusion_err!("RS_ZonalStats: raster dimensions overflow"))?; + if band_bytes.len() != expected { + return sedona_internal_err!( + "RS_ZonalStats: band {band_num} byte length {} does not match {width}x{height} of {data_type:?}", + band_bytes.len() + ); + } + + // Nodata is compared in the band's own byte representation, never through + // f64 — an Int64/UInt64 nodata beyond 2^53 must not alias a nearby pixel. + let nodata = if params.exclude_no_data { + band.nodata() + } else { + None + }; + if let Some(nd) = nodata { + if nd.len() != byte_size { + return sedona_internal_err!( + "RS_ZonalStats: band {band_num} nodata is {} bytes, expected {byte_size} for {data_type:?}", + nd.len() + ); + } + } + + scratch.clear(); + collect_masked_values( + band_bytes, + data_type, + width, + &window, + mask_scratch, + nodata, + scratch, + ); + + Ok(RoiCoverage::Collected) +} + +/// Resolve the 1-based band to use. `Some(b)` must be a valid 1-based index; +/// `None` defaults to band 1 for a single-band raster and errors for a +/// multiband raster (matching the codebase's `default_band` convention, which +/// refuses to silently pick band 1 when the choice is ambiguous). +fn resolve_band(band: Option<i64>, num_bands: usize) -> Result<usize> { + match band { + Some(b) => { + if b < 1 { + return exec_err!("RS_ZonalStats: band must be >= 1, got {b}"); + } + let b = b as usize; + if b > num_bands { + return exec_err!("RS_ZonalStats: band {b} is out of range (1-{num_bands})"); + } + Ok(b) + } + None => { + if num_bands == 1 { + Ok(1) + } else { + exec_err!( + "RS_ZonalStats: raster has {num_bands} bands; pass the band argument to \ + choose one (only a single-band raster may omit it)" + ) + } + } + } +} + +/// Append every selected pixel value (masked in, and — when `nodata` is set — +/// not byte-equal to the nodata sentinel) to `out` as `f64`. +/// +/// The data type is dispatched once, outside the loop, so the per-pixel body is +/// a fixed-width little-endian read plus the mask/nodata comparisons rather +/// than a per-pixel type match. +fn collect_masked_values( + band_bytes: &[u8], + data_type: BandDataType, + width: usize, + window: &PixelWindow, + mask: &[u8], + nodata: Option<&[u8]>, + out: &mut Vec<f64>, +) { + macro_rules! collect { + ($t:ty, $n:literal) => {{ + for row in 0..window.height { + let src_row = window.row_off + row; + let mask_row = row * window.width; + for col in 0..window.width { + if mask[mask_row + col] == 0 { + continue; + } + let idx = (src_row * width + window.col_off + col) * $n; + let px = &band_bytes[idx..idx + $n]; + if let Some(nd) = nodata { + if px == nd { + continue; + } + } + let mut arr = [0u8; $n]; + arr.copy_from_slice(px); + out.push(<$t>::from_le_bytes(arr) as f64); + } + } + }}; + } + + match data_type { + BandDataType::UInt8 => collect!(u8, 1), + BandDataType::Int8 => collect!(i8, 1), + BandDataType::UInt16 => collect!(u16, 2), + BandDataType::Int16 => collect!(i16, 2), + BandDataType::UInt32 => collect!(u32, 4), + BandDataType::Int32 => collect!(i32, 4), + BandDataType::UInt64 => collect!(u64, 8), + BandDataType::Int64 => collect!(i64, 8), + BandDataType::Float32 => collect!(f32, 4), + BandDataType::Float64 => collect!(f64, 8), + } +} + +/// Compute every statistic from the selected pixel values (for +/// RS_ZonalStatsAll, which returns all of them). +/// +/// An empty slice yields `count = 0` and NULL for the rest (Sedona Spark's +/// empty-roi shortcut). Variance is the sample (n-1) variance, matching Spark; +/// for a single pixel it is 0. Median is the linear-interpolated 50th +/// percentile, which reduces to the middle element (odd n) or the mean of the +/// two central elements (even n). Mode is the most frequent value, breaking ties +/// toward the larger value. +/// +/// `values` is sorted in place (for the median and mode); the caller owns it as +/// reusable scratch. +fn compute_statistics(values: &mut [f64]) -> ZonalStatistics { + let count = values.len() as i64; + if values.is_empty() { + return ZonalStatistics { + count: 0, + sum: None, + mean: None, + median: None, + mode: None, + stddev: None, + variance: None, + min: None, + max: None, + }; + } + + // A NaN pixel (e.g. a float band whose NaN nodata was not excluded) poisons + // every statistic under numpy semantics. Return NaN for all of them rather + // than letting f64::min / f64::max silently skip NaN while sum and mean + // propagate it — an internally inconsistent, reference-diverging result. + if values.iter().any(|v| v.is_nan()) { + return ZonalStatistics { + count, + sum: Some(f64::NAN), + mean: Some(f64::NAN), + median: Some(f64::NAN), + mode: Some(f64::NAN), + stddev: Some(f64::NAN), + variance: Some(f64::NAN), + min: Some(f64::NAN), + max: Some(f64::NAN), + }; + } + + let sum: f64 = values.iter().sum(); + let mean = sum / count as f64; + let min = values.iter().copied().fold(f64::INFINITY, f64::min); + let max = values.iter().copied().fold(f64::NEG_INFINITY, f64::max); + let variance = sample_variance(values); + let stddev = variance.sqrt(); + + // Median and mode both read the values in sorted order; sort once in place. + sort_values(values); + let median = median_of_sorted(values); + let mode = mode_of_sorted(values); + + ZonalStatistics { + count, + sum: Some(sum), + mean: Some(mean), + median: Some(median), + mode: Some(mode), + stddev: Some(stddev), + variance: Some(variance), + min: Some(min), + max: Some(max), + } +} + +/// Compute only `stat` from the selected pixel values (for RS_ZonalStats, which +/// returns a single statistic). `values` is sorted in place only when the +/// requested statistic — median or mode — needs ordered data, so the simple +/// statistics do not pay for a sort. +/// +/// Mirrors [`compute_statistics`]' empty/NaN semantics: `count` is always +/// defined (0 for an empty roi); every other statistic is NULL (`None`) for an +/// empty roi, and NaN when a NaN pixel is present. +fn compute_single_statistic(values: &mut [f64], stat: StatType) -> Option<f64> { + if stat == StatType::Count { + return Some(values.len() as f64); + } + if values.is_empty() { + return None; + } + if values.iter().any(|v| v.is_nan()) { + return Some(f64::NAN); + } + Some(match stat { + // Count is handled before the value checks above. + StatType::Count => unreachable!("count returns early"), + StatType::Sum => values.iter().sum(), + StatType::Mean => values.iter().sum::<f64>() / values.len() as f64, + StatType::Min => values.iter().copied().fold(f64::INFINITY, f64::min), + StatType::Max => values.iter().copied().fold(f64::NEG_INFINITY, f64::max), + StatType::Variance => sample_variance(values), + StatType::StdDev => sample_variance(values).sqrt(), + StatType::Median => { + sort_values(values); + median_of_sorted(values) + } + StatType::Mode => { + sort_values(values); + mode_of_sorted(values) + } + }) +} + +/// The sample (n-1) variance of `values`, matching Sedona Spark; 0 for a single +/// value. The two-pass form (mean first, then squared deviations) avoids the +/// catastrophic cancellation of the naive sum-of-squares formula. `values` must +/// be non-empty and NaN-free. +fn sample_variance(values: &[f64]) -> f64 { + let n = values.len(); + if n <= 1 { + return 0.0; + } + let mean = values.iter().sum::<f64>() / n as f64; Review Comment: Can we use a previously computed sum or mean here? ########## rust/sedona-raster-gdal/src/rs_zonal_stats.rs: ########## @@ -0,0 +1,1749 @@ +// 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_ZonalStats / RS_ZonalStatsAll UDFs — summary statistics of the raster +//! pixels covered by a roi geometry. +//! +//! Both mirror Apache Sedona Spark's positional overloads verbatim so that +//! Spark SQL tends to run unchanged. `RS_ZonalStats` returns one statistic as a +//! `Float64`: +//! +//! - `RS_ZonalStats(raster, roi, stat)` +//! - `RS_ZonalStats(raster, roi, band, stat)` +//! - `RS_ZonalStats(raster, roi, band, stat, all_touched)` +//! - `RS_ZonalStats(raster, roi, band, stat, all_touched, exclude_no_data)` +//! - `RS_ZonalStats(raster, roi, band, stat, all_touched, exclude_no_data, lenient)` +//! +//! `RS_ZonalStatsAll` returns every statistic as a struct, with the same ladder +//! minus `stat`: +//! +//! - `RS_ZonalStatsAll(raster, roi)` +//! - `RS_ZonalStatsAll(raster, roi, band)` +//! - `RS_ZonalStatsAll(raster, roi, band, all_touched)` +//! - `RS_ZonalStatsAll(raster, roi, band, all_touched, exclude_no_data)` +//! - `RS_ZonalStatsAll(raster, roi, band, all_touched, exclude_no_data, lenient)` +//! +//! A pixel is included when its centre falls inside the roi (or that the roi +//! merely touches, with `all_touched`), optionally excluding the band's nodata +//! value. `all_touched` defaults to false, `exclude_no_data` to true, and +//! `lenient` to true. Unlike Sedona Spark, the band-less overloads do not +//! default to band 1 on a multiband raster: naming the band is required there +//! (a single-band raster resolves unambiguously). +//! +//! These functions operate on 2-D `(y, x)` bands. A band that is not a 2-D +//! spatial grid is rejected; computing a statistic per non-spatial plane of an +//! N-D band is not supported. + +use std::sync::Arc; + +use arrow_array::builder::{Float64Builder, Int64Builder}; +use arrow_array::{ArrayRef, BooleanArray, Int64Array, StringArray, StructArray}; +use arrow_buffer::{BooleanBufferBuilder, NullBuffer}; +use arrow_schema::{DataType, Field, Fields}; +use datafusion_common::cast::{as_boolean_array, as_int64_array, as_string_array}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::error::Result; +use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue}; +use datafusion_expr::{ColumnarValue, Volatility}; + +use sedona_common::sedona_internal_err; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_gdal::gdal::Gdal; +use sedona_raster::array::RasterRefImpl; +use sedona_raster::traits::RasterRef; +use sedona_raster_functions::crs_utils::{align_wkb_to_crs, resolve_crs, with_crs_engine}; +use sedona_raster_functions::rs_ensure_loaded::NEEDS_PIXELS_METADATA_KEY; +use sedona_raster_functions::rs_spatial_predicates::raster_intersects_geom_wkb; +use sedona_raster_functions::RasterExecutor; +use sedona_schema::datatypes::SedonaType; +use sedona_schema::matchers::ArgMatcher; +use sedona_schema::raster::BandDataType; + +use crate::gdal_common::{raster_geo_transform, with_gdal}; +use crate::gdal_dataset_provider::configure_thread_local_options; +use crate::mask::{envelope_window, rasterize_geometry_mask, PixelWindow}; + +/// The statistics RS_ZonalStatsAll returns, in the order Sedona Spark reports +/// them. RS_ZonalStats selects one of these by name. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StatType { + Count, + Sum, + Mean, + Median, + Mode, + StdDev, + Variance, + Min, + Max, +} + +impl StatType { + /// Parse a statistic name (case-insensitive). Aliases match Sedona Spark + /// (`avg`/`average` for mean, `sd` for stddev). + fn from_str(s: &str) -> Option<StatType> { + match s.to_lowercase().as_str() { + "count" => Some(StatType::Count), + "sum" => Some(StatType::Sum), + "mean" | "avg" | "average" => Some(StatType::Mean), + "median" => Some(StatType::Median), + "mode" => Some(StatType::Mode), + "stddev" | "sd" => Some(StatType::StdDev), + "variance" => Some(StatType::Variance), + "min" => Some(StatType::Min), + "max" => Some(StatType::Max), + _ => None, + } + } +} + +/// Defaults for the trailing flags, applied by the narrower overloads that omit +/// them (matching Sedona Spark). +const DEFAULT_ALL_TOUCHED: bool = false; +const DEFAULT_EXCLUDE_NODATA: bool = true; +const DEFAULT_LENIENT: bool = true; + +/// The resolved parameters for one row's zonal-stats computation, assembled from +/// the positional arguments the matched overload carried. +#[derive(Debug, Clone)] +struct ZonalStatsParams { + /// 1-based band to compute over. `None` means "resolve the implicit band": + /// band 1 for a single-band raster, an error for a multiband raster (naming + /// the band is required rather than silently getting band 1). Only the + /// band-less overloads leave this `None`. + band: Option<i64>, + /// Include every pixel the roi touches, not only those whose centre it + /// covers. + all_touched: bool, + /// Skip pixels equal to the band's nodata value. + exclude_no_data: bool, + /// Return NULL when the roi does not intersect the raster, rather than + /// erroring. Only the no-intersection case is softened; malformed geometry + /// or an unreadable band always errors. + lenient: bool, +} + +/// Every statistic for a roi. `count` is always present (0 when the roi +/// selects no pixels); the remaining fields are `None` in exactly that +/// no-pixel case and `Some` otherwise, mirroring Sedona Spark (which returns +/// `count = 0` and NULL for the rest). +#[derive(Debug, Clone, PartialEq)] +struct ZonalStatistics { + count: i64, + sum: Option<f64>, + mean: Option<f64>, + median: Option<f64>, + mode: Option<f64>, + stddev: Option<f64>, + variance: Option<f64>, + min: Option<f64>, + max: Option<f64>, +} + +/// Whether a roi geometry intersects the raster. `NoIntersection` mirrors Sedona +/// Spark's `rsIntersects` gate (the caller turns it into NULL when `lenient`, an +/// error otherwise); `Collected` means the selected pixel values are in the +/// caller's scratch buffer, possibly empty for a roi that intersects the +/// footprint but selects no pixel centre (a `count = 0` result). +enum RoiCoverage { + NoIntersection, + Collected, +} + +// ============================================================================= +// RS_ZonalStats +// ============================================================================= + +/// `RS_ZonalStats` — one statistic as a `Float64`. `stat` is a statistic name +/// (`count`, `sum`, `mean`, `median`, `mode`, `stddev`, `variance`, `min`, +/// `max`). See the module docs for the full positional overload ladder. +pub fn rs_zonal_stats_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_zonalstats", + vec![ + Arc::new(RsZonalStats { arg_count: 3 }), // (raster, roi, stat) + Arc::new(RsZonalStats { arg_count: 4 }), // (raster, roi, band, stat) + Arc::new(RsZonalStats { arg_count: 5 }), // + all_touched + Arc::new(RsZonalStats { arg_count: 6 }), // + exclude_no_data + Arc::new(RsZonalStats { arg_count: 7 }), // + lenient + ], + Volatility::Immutable, + ) + // Reads band pixels, so the planner materializes OutDb rasters via + // RS_EnsureLoaded first. + .with_metadata(NEEDS_PIXELS_METADATA_KEY, "true") +} + +#[derive(Debug)] +struct RsZonalStats { + /// Number of arguments in the matched signature (3..=7). + arg_count: usize, +} + +impl SedonaScalarKernel for RsZonalStats { + fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> { + // Argument order mirrors Sedona Spark: (raster, roi, [band,] stat, + // [all_touched, [exclude_no_data, [lenient]]]). The 3-arg overload omits + // band (its stat is at index 2); the 4+-arg overloads carry band at + // index 2 and stat at index 3. + let matchers = match self.arg_count { + 3 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_string(), + ], + 4 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ], + 5 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ArgMatcher::is_boolean(), + ], + 6 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ], + 7 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ], + _ => { + return sedona_internal_err!( + "RS_ZonalStats: unexpected arg_count {}", + self.arg_count + ); + } + }; + let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(DataType::Float64)); + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + 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<ColumnarValue> { + let num_iterations = RasterExecutor::num_iterations_over(args); + + // band (index 2) only exists in the 4+-arg overloads; the 3-arg overload + // leaves it implicit. stat is at index 2 (3-arg) or 3 (4+-arg). + let has_band = self.arg_count >= 4; + let stat_idx = if has_band { 3 } else { 2 }; + let stat_array = expand_string_arg(&args[stat_idx], num_iterations)?; + let mut stat_iter = stat_array.iter(); + + let band_array = has_band + .then(|| expand_int64_arg(&args[2], num_iterations)) + .transpose()?; + let mut band_iter = band_array.as_ref().map(|a| a.iter()); + + // all_touched (index 4), exclude_no_data (index 5), lenient (index 6): + // read from the column when the overload carries it, else the default. + let all_touched_array = expand_flag( + args, + 4, + self.arg_count >= 5, + DEFAULT_ALL_TOUCHED, + num_iterations, + )?; + let exclude_no_data_array = expand_flag( + args, + 5, + self.arg_count >= 6, + DEFAULT_EXCLUDE_NODATA, + num_iterations, + )?; + let lenient_array = expand_flag( + args, + 6, + self.arg_count >= 7, + DEFAULT_LENIENT, + num_iterations, + )?; + let mut all_touched_iter = all_touched_array.iter(); + let mut exclude_no_data_iter = exclude_no_data_array.iter(); + let mut lenient_iter = lenient_array.iter(); + + let mut builder = Float64Builder::with_capacity(num_iterations); + let mut scratch: Vec<f64> = Vec::new(); + let mut mask_scratch: Vec<u8> = Vec::new(); + + // The executor only sees (raster, roi); the option columns are advanced + // in lockstep below. + let exec_arg_types = [arg_types[0].clone(), arg_types[1].clone()]; + let exec_args = [args[0].clone(), args[1].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_crs_engine(config_options, |engine| { + executor.execute_raster_wkb_crs_void(|raster_opt, wkb_opt, geom_crs| { + let stat_str = stat_iter.next().flatten(); + let Some(params) = next_params( + &mut band_iter, + &mut all_touched_iter, + &mut exclude_no_data_iter, + &mut lenient_iter, + ) else { + builder.append_null(); + return Ok(()); + }; + + // A NULL stat, raster, or roi propagates to a NULL row. + let (Some(stat_str), Some(raster), Some(wkb)) = (stat_str, raster_opt, wkb_opt) + else { + builder.append_null(); + return Ok(()); + }; + let stat_type = StatType::from_str(stat_str).ok_or_else(|| { + exec_datafusion_err!("RS_ZonalStats: unknown statistic {stat_str:?}") + })?; + + // Reproject the roi into the raster's CRS, borrowing it + // unchanged when the CRSes already match; a CRS on exactly + // one side is an error, since it would mislocate the roi. + let raster_crs = resolve_crs(raster.crs())?; + let geom_wkb = align_wkb_to_crs( + wkb, + geom_crs, + raster_crs.as_deref(), + "geometry", + "raster", + engine, + )?; + match collect_zonal_values( + gdal, + raster, + &geom_wkb, + ¶ms, + &mut scratch, + &mut mask_scratch, + )? { + // Compute only the requested statistic, not all of them. + RoiCoverage::Collected => { + match compute_single_statistic(&mut scratch, stat_type) { + Some(value) => builder.append_value(value), + None => builder.append_null(), + } + } + // The roi does not intersect the raster: NULL when + // lenient (the default), an error otherwise. + RoiCoverage::NoIntersection if params.lenient => builder.append_null(), + RoiCoverage::NoIntersection => return no_intersection_err(), + } + Ok(()) + }) + })?; + + let out: ArrayRef = Arc::new(builder.finish()); + RasterExecutor::finish_over(args, out) + }) + } +} + +// ============================================================================= +// RS_ZonalStatsAll +// ============================================================================= + +/// `RS_ZonalStatsAll` — every statistic as a struct with fields `count, sum, +/// mean, median, mode, stddev, variance, min, max`. See the module docs for the +/// full positional overload ladder. +pub fn rs_zonal_stats_all_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_zonalstatsall", + vec![ + Arc::new(RsZonalStatsAll { arg_count: 2 }), // (raster, roi) + Arc::new(RsZonalStatsAll { arg_count: 3 }), // (raster, roi, band) + Arc::new(RsZonalStatsAll { arg_count: 4 }), // + all_touched + Arc::new(RsZonalStatsAll { arg_count: 5 }), // + exclude_no_data + Arc::new(RsZonalStatsAll { arg_count: 6 }), // + lenient + ], + Volatility::Immutable, + ) + .with_metadata(NEEDS_PIXELS_METADATA_KEY, "true") +} + +#[derive(Debug)] +struct RsZonalStatsAll { + /// Number of arguments in the matched signature (2..=6). + arg_count: usize, +} + +impl SedonaScalarKernel for RsZonalStatsAll { + fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> { + // Argument order mirrors Sedona Spark: (raster, roi, [band, + // [all_touched, [exclude_no_data, [lenient]]]]). The 2-arg overload omits + // band; the 3+-arg overloads carry it at index 2. + let mut matchers = vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ]; + if self.arg_count >= 3 { + matchers.push(ArgMatcher::is_integer()); // band + } + for _ in 4..=self.arg_count { + matchers.push(ArgMatcher::is_boolean()); // all_touched, exclude_no_data, lenient + } + if self.arg_count < 2 || self.arg_count > 6 { + return sedona_internal_err!( + "RS_ZonalStatsAll: unexpected arg_count {}", + self.arg_count + ); + } + let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(zonal_stats_struct_type())); + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + 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<ColumnarValue> { + let num_iterations = RasterExecutor::num_iterations_over(args); + + // band (index 2) only exists in the 3+-arg overloads; the 2-arg overload + // leaves it implicit. all_touched (index 3), exclude_no_data (index 4), + // and lenient (index 5) follow. + let band_array = (self.arg_count >= 3) + .then(|| expand_int64_arg(&args[2], num_iterations)) + .transpose()?; + let mut band_iter = band_array.as_ref().map(|a| a.iter()); + + let all_touched_array = expand_flag( + args, + 3, + self.arg_count >= 4, + DEFAULT_ALL_TOUCHED, + num_iterations, + )?; + let exclude_no_data_array = expand_flag( + args, + 4, + self.arg_count >= 5, + DEFAULT_EXCLUDE_NODATA, + num_iterations, + )?; + let lenient_array = expand_flag( + args, + 5, + self.arg_count >= 6, + DEFAULT_LENIENT, + num_iterations, + )?; + let mut all_touched_iter = all_touched_array.iter(); + let mut exclude_no_data_iter = exclude_no_data_array.iter(); + let mut lenient_iter = lenient_array.iter(); + + let mut builders = ZonalStatsBuilders::with_capacity(num_iterations); + let mut scratch: Vec<f64> = Vec::new(); + let mut mask_scratch: Vec<u8> = Vec::new(); + + let exec_arg_types = [arg_types[0].clone(), arg_types[1].clone()]; + let exec_args = [args[0].clone(), args[1].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_crs_engine(config_options, |engine| { + executor.execute_raster_wkb_crs_void(|raster_opt, wkb_opt, geom_crs| { + let Some(params) = next_params( + &mut band_iter, + &mut all_touched_iter, + &mut exclude_no_data_iter, + &mut lenient_iter, + ) else { + builders.push_null(); + return Ok(()); + }; + + let (Some(raster), Some(wkb)) = (raster_opt, wkb_opt) else { + builders.push_null(); + return Ok(()); + }; + + let raster_crs = resolve_crs(raster.crs())?; + let geom_wkb = align_wkb_to_crs( + wkb, + geom_crs, + raster_crs.as_deref(), + "geometry", + "raster", + engine, + )?; + match collect_zonal_values( + gdal, + raster, + &geom_wkb, + ¶ms, + &mut scratch, + &mut mask_scratch, + )? { + RoiCoverage::Collected => { + builders.push_stats(&compute_statistics(&mut scratch)) + } + RoiCoverage::NoIntersection if params.lenient => builders.push_null(), + RoiCoverage::NoIntersection => return no_intersection_err(), + } + Ok(()) + }) + })?; + + let out: ArrayRef = Arc::new(builders.finish()); + RasterExecutor::finish_over(args, out) + }) + } +} + +/// Struct data type RS_ZonalStatsAll returns. +fn zonal_stats_struct_type() -> DataType { + DataType::Struct(zonal_stats_struct_fields()) +} + +/// Fields of the RS_ZonalStatsAll struct, in Sedona Spark order. `count` is an +/// `Int64` (a whole pixel count); every other statistic is a `Float64`. +fn zonal_stats_struct_fields() -> Fields { + Fields::from(vec![ + Field::new("count", DataType::Int64, true), + Field::new("sum", DataType::Float64, true), + Field::new("mean", DataType::Float64, true), + Field::new("median", DataType::Float64, true), + Field::new("mode", DataType::Float64, true), + Field::new("stddev", DataType::Float64, true), + Field::new("variance", DataType::Float64, true), + Field::new("min", DataType::Float64, true), + Field::new("max", DataType::Float64, true), + ]) +} + +/// Column builders for the RS_ZonalStatsAll struct output: one typed builder per +/// field plus an outer struct-level validity buffer. +/// +/// Building the columns directly and assembling the [`StructArray`] once at the +/// end avoids downcasting a `StructBuilder`'s boxed field builders on every row. +struct ZonalStatsBuilders { + count: Int64Builder, + /// `sum, mean, median, mode, stddev, variance, min, max` — Sedona Spark field + /// order, matching [`zonal_stats_struct_fields`] after `count`. + floats: [Float64Builder; 8], + /// Struct-level null bitmap: `false` for a fully-NULL row (a NULL input or a + /// non-intersecting roi under `lenient`). + validity: BooleanBufferBuilder, +} + +impl ZonalStatsBuilders { + fn with_capacity(capacity: usize) -> Self { + Self { + count: Int64Builder::with_capacity(capacity), + floats: std::array::from_fn(|_| Float64Builder::with_capacity(capacity)), + validity: BooleanBufferBuilder::new(capacity), + } + } + + /// Append one computed-stats row (the struct itself is valid). The float + /// fields carry through the `Option`, so an empty roi records `count = 0` + /// with the rest NULL. + fn push_stats(&mut self, stats: &ZonalStatistics) { + self.count.append_value(stats.count); + let values = [ + stats.sum, + stats.mean, + stats.median, + stats.mode, + stats.stddev, + stats.variance, + stats.min, + stats.max, + ]; + for (builder, value) in self.floats.iter_mut().zip(values) { + builder.append_option(value); + } + self.validity.append(true); + } + + /// Append a fully-NULL struct row (a NULL input, or a non-intersecting roi + /// when `lenient`). Every field is null and the struct itself is null. + fn push_null(&mut self) { + self.count.append_null(); + for builder in &mut self.floats { + builder.append_null(); + } + self.validity.append(false); + } + + /// Assemble the accumulated columns into the struct array. + fn finish(mut self) -> StructArray { + let mut arrays: Vec<ArrayRef> = Vec::with_capacity(9); + arrays.push(Arc::new(self.count.finish())); + for builder in &mut self.floats { + arrays.push(Arc::new(builder.finish())); + } + let nulls = NullBuffer::new(self.validity.finish()); + StructArray::new(zonal_stats_struct_fields(), arrays, Some(nulls)) + } +} + +// ============================================================================= +// Core computation +// ============================================================================= + +/// Collect the pixel values a roi geometry selects on one band into `scratch`. +/// +/// Returns [`RoiCoverage::NoIntersection`] when the roi geometry does not +/// intersect the raster's footprint — a true geometry intersection (matching +/// Sedona Spark's `rsIntersects` gate), not a bounding-box overlap: a roi whose +/// envelope overlaps the raster but whose geometry is disjoint is a +/// no-intersection case. The caller turns that into NULL when `lenient`, an +/// error otherwise. A roi that intersects the footprint but whose selected +/// pixels are all outside the geometry or all nodata returns +/// [`RoiCoverage::Collected`] with `scratch` left empty (a `count = 0` result). +/// +/// The caller computes the statistic(s) it needs from `scratch` — every one for +/// RS_ZonalStatsAll, only the requested one for RS_ZonalStats — so this shared +/// collection never computes statistics the caller would discard. +/// +/// `scratch` is a reused buffer for the selected pixel values and `mask_scratch` +/// for the rasterized roi mask, both reused so the per-row computation does not +/// allocate a fresh `Vec` each call. +fn collect_zonal_values( + gdal: &Gdal, + raster: &RasterRefImpl<'_>, + geom_wkb: &[u8], + params: &ZonalStatsParams, + scratch: &mut Vec<f64>, + mask_scratch: &mut Vec<u8>, +) -> Result<RoiCoverage> { + let num_bands = raster.num_bands(); + let band_num = resolve_band(params.band, num_bands)?; + + let band = raster + .bands() + .band(band_num) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to read band {band_num}: {e}"))?; + if !band.is_spatial_2d() { + return exec_err!( + "RS_ZonalStats supports 2-D rasters only; band {band_num} is not a 2-D (y, x) grid" + ); + } + let data_type = band.data_type(); + let byte_size = data_type.byte_size(); + + let metadata = raster.metadata(); + let transform = raster_geo_transform(raster)?; + let width = usize::try_from(metadata.width()) + .map_err(|_| exec_datafusion_err!("RS_ZonalStats: negative raster width"))?; + let height = usize::try_from(metadata.height()) + .map_err(|_| exec_datafusion_err!("RS_ZonalStats: negative raster height"))?; + + // No-intersection gate: a true geometry intersection between the roi and + // the raster footprint (matching Sedona Spark's rsIntersects gate), not a + // bounding-box overlap. A roi whose envelope overlaps the raster but whose + // geometry is disjoint is a no-intersection case, not a count-0 case. The + // roi is already in the raster's CRS here, so no transform is needed. + if !raster_intersects_geom_wkb(raster, geom_wkb)? { + return Ok(RoiCoverage::NoIntersection); + } + + // Parse the roi and clamp its envelope to the raster grid for the pixel + // window to rasterize. The gate above already established overlap; a + // degenerate window (the roi only touches the raster boundary) selects no + // pixels, so it is count 0 rather than no-intersection. + let geometry = gdal + .geometry_from_wkb(geom_wkb) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to parse geometry: {e}"))?; + let Some(window) = envelope_window(&geometry, &transform, width, height)? else { + scratch.clear(); + return Ok(RoiCoverage::Collected); + }; + + // Rasterize the roi into a window-sized 0/1 mask (moves `geometry`, whose + // only remaining use is the burn). The mask reuses `mask_scratch` across + // rows rather than allocating a fresh buffer each call. + rasterize_geometry_mask( + gdal, + geometry, + &transform, + &window, + params.all_touched, + mask_scratch, + )?; + + // Read the band once (zero-copy borrow) and collect the selected values. + let nd_buffer = band + .nd_buffer() + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to read band {band_num}: {e}"))?; + let band_bytes = nd_buffer.as_contiguous().map_err(|e| { + exec_datafusion_err!("RS_ZonalStats: band {band_num} is not contiguous: {e}") + })?; + let expected = width + .checked_mul(height) + .and_then(|n| n.checked_mul(byte_size)) + .ok_or_else(|| exec_datafusion_err!("RS_ZonalStats: raster dimensions overflow"))?; + if band_bytes.len() != expected { + return sedona_internal_err!( + "RS_ZonalStats: band {band_num} byte length {} does not match {width}x{height} of {data_type:?}", + band_bytes.len() + ); + } + + // Nodata is compared in the band's own byte representation, never through + // f64 — an Int64/UInt64 nodata beyond 2^53 must not alias a nearby pixel. + let nodata = if params.exclude_no_data { + band.nodata() + } else { + None + }; + if let Some(nd) = nodata { + if nd.len() != byte_size { + return sedona_internal_err!( + "RS_ZonalStats: band {band_num} nodata is {} bytes, expected {byte_size} for {data_type:?}", + nd.len() + ); + } + } + + scratch.clear(); + collect_masked_values( + band_bytes, + data_type, + width, + &window, + mask_scratch, + nodata, + scratch, + ); + + Ok(RoiCoverage::Collected) +} + +/// Resolve the 1-based band to use. `Some(b)` must be a valid 1-based index; +/// `None` defaults to band 1 for a single-band raster and errors for a +/// multiband raster (matching the codebase's `default_band` convention, which +/// refuses to silently pick band 1 when the choice is ambiguous). +fn resolve_band(band: Option<i64>, num_bands: usize) -> Result<usize> { + match band { + Some(b) => { + if b < 1 { + return exec_err!("RS_ZonalStats: band must be >= 1, got {b}"); + } + let b = b as usize; + if b > num_bands { + return exec_err!("RS_ZonalStats: band {b} is out of range (1-{num_bands})"); + } + Ok(b) + } + None => { + if num_bands == 1 { + Ok(1) + } else { + exec_err!( + "RS_ZonalStats: raster has {num_bands} bands; pass the band argument to \ + choose one (only a single-band raster may omit it)" + ) + } + } + } +} + +/// Append every selected pixel value (masked in, and — when `nodata` is set — +/// not byte-equal to the nodata sentinel) to `out` as `f64`. +/// +/// The data type is dispatched once, outside the loop, so the per-pixel body is +/// a fixed-width little-endian read plus the mask/nodata comparisons rather +/// than a per-pixel type match. +fn collect_masked_values( + band_bytes: &[u8], + data_type: BandDataType, + width: usize, + window: &PixelWindow, + mask: &[u8], + nodata: Option<&[u8]>, + out: &mut Vec<f64>, +) { + macro_rules! collect { + ($t:ty, $n:literal) => {{ + for row in 0..window.height { + let src_row = window.row_off + row; + let mask_row = row * window.width; + for col in 0..window.width { + if mask[mask_row + col] == 0 { + continue; + } + let idx = (src_row * width + window.col_off + col) * $n; + let px = &band_bytes[idx..idx + $n]; + if let Some(nd) = nodata { + if px == nd { + continue; + } + } + let mut arr = [0u8; $n]; + arr.copy_from_slice(px); + out.push(<$t>::from_le_bytes(arr) as f64); + } + } + }}; + } + + match data_type { + BandDataType::UInt8 => collect!(u8, 1), + BandDataType::Int8 => collect!(i8, 1), + BandDataType::UInt16 => collect!(u16, 2), + BandDataType::Int16 => collect!(i16, 2), + BandDataType::UInt32 => collect!(u32, 4), + BandDataType::Int32 => collect!(i32, 4), + BandDataType::UInt64 => collect!(u64, 8), + BandDataType::Int64 => collect!(i64, 8), + BandDataType::Float32 => collect!(f32, 4), + BandDataType::Float64 => collect!(f64, 8), + } +} + +/// Compute every statistic from the selected pixel values (for +/// RS_ZonalStatsAll, which returns all of them). +/// +/// An empty slice yields `count = 0` and NULL for the rest (Sedona Spark's +/// empty-roi shortcut). Variance is the sample (n-1) variance, matching Spark; +/// for a single pixel it is 0. Median is the linear-interpolated 50th +/// percentile, which reduces to the middle element (odd n) or the mean of the +/// two central elements (even n). Mode is the most frequent value, breaking ties +/// toward the larger value. +/// +/// `values` is sorted in place (for the median and mode); the caller owns it as +/// reusable scratch. +fn compute_statistics(values: &mut [f64]) -> ZonalStatistics { + let count = values.len() as i64; + if values.is_empty() { + return ZonalStatistics { + count: 0, + sum: None, + mean: None, + median: None, + mode: None, + stddev: None, + variance: None, + min: None, + max: None, + }; + } + + // A NaN pixel (e.g. a float band whose NaN nodata was not excluded) poisons + // every statistic under numpy semantics. Return NaN for all of them rather + // than letting f64::min / f64::max silently skip NaN while sum and mean + // propagate it — an internally inconsistent, reference-diverging result. + if values.iter().any(|v| v.is_nan()) { + return ZonalStatistics { + count, + sum: Some(f64::NAN), + mean: Some(f64::NAN), + median: Some(f64::NAN), + mode: Some(f64::NAN), + stddev: Some(f64::NAN), + variance: Some(f64::NAN), + min: Some(f64::NAN), + max: Some(f64::NAN), + }; + } + + let sum: f64 = values.iter().sum(); + let mean = sum / count as f64; + let min = values.iter().copied().fold(f64::INFINITY, f64::min); + let max = values.iter().copied().fold(f64::NEG_INFINITY, f64::max); + let variance = sample_variance(values); + let stddev = variance.sqrt(); + + // Median and mode both read the values in sorted order; sort once in place. + sort_values(values); + let median = median_of_sorted(values); + let mode = mode_of_sorted(values); Review Comment: Since you're computing all of these, you can use your cumulative calculations to avoid a few passes over the data: min is the first of the sorted values; max is the last of the sorted values variance can use the cached sum with one additional pass. If you wire your sort so that `NaN`s are first, you can also avoid the pass over all the values to check for NaN. You can then eliminate the single statistic path (just return early if there is any work that can be avoided, like returning before the sort or returning before the variance calculation) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
