james-willis commented on code in PR #1066: URL: https://github.com/apache/sedona-db/pull/1066#discussion_r3670486680
########## rust/sedona-raster-gdal/src/rs_zonal_stats.rs: ########## @@ -0,0 +1,1609 @@ +// 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::collections::HashMap; +use std::sync::Arc; + +use arrow_array::builder::{Float64Builder, Int64Builder, StructBuilder}; +use arrow_array::{ArrayRef, BooleanArray, Int64Array, StringArray}; +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_gdal::geo_transform::GeoTransform; +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::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>, +} + +impl ZonalStatistics { + /// The value RS_ZonalStats returns for a single statistic. `count` is never + /// NULL (it is 0 for an empty roi); the others are NULL for an empty roi. + fn get(&self, stat_type: StatType) -> Option<f64> { + match stat_type { + StatType::Count => Some(self.count as f64), + StatType::Sum => self.sum, + StatType::Mean => self.mean, + StatType::Median => self.median, + StatType::Mode => self.mode, + StatType::StdDev => self.stddev, + StatType::Variance => self.variance, + StatType::Min => self.min, + StatType::Max => self.max, + } + } +} + +// ============================================================================= +// 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(); + + // 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 compute_zonal_stats(gdal, raster, &geom_wkb, ¶ms, &mut scratch)? { + Some(stats) => match stats.get(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. + None if params.lenient => builder.append_null(), + None => 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 builder = StructBuilder::from_fields(zonal_stats_struct_fields(), num_iterations); + let mut scratch: Vec<f64> = 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 { + append_struct_null(&mut builder)?; + return Ok(()); + }; + + let (Some(raster), Some(wkb)) = (raster_opt, wkb_opt) else { + append_struct_null(&mut builder)?; + 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 compute_zonal_stats(gdal, raster, &geom_wkb, ¶ms, &mut scratch)? { + Some(stats) => append_struct_stats(&mut builder, &stats)?, + None if params.lenient => append_struct_null(&mut builder)?, + None => return no_intersection_err(), + } + Ok(()) + }) + })?; + + let out: ArrayRef = Arc::new(builder.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), + ]) +} + +/// Append a fully-NULL struct row (the roi does not intersect the raster and +/// `lenient` is set). +fn append_struct_null(builder: &mut StructBuilder) -> Result<()> { + let Some(count) = builder.field_builder::<Int64Builder>(0) else { + return sedona_internal_err!("RS_ZonalStats: count field is not an Int64 builder"); + }; + count.append_null(); + for i in 1..=8 { + let Some(field) = builder.field_builder::<Float64Builder>(i) else { + return sedona_internal_err!("RS_ZonalStats: stat field {i} is not a Float64 builder"); + }; + field.append_null(); + } + builder.append(false); + Ok(()) +} Review Comment: fixed -- 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]
