From aa77a2e392ce602c6ddf0da76d166859e72f3780 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 20 Jul 2026 13:08:11 -0700 Subject: [PATCH 01/16] feat(rust/sedona-raster-gdal): add RS_ZonalStats and RS_ZonalStatsAll UDFs Compute summary statistics of the raster pixels covered by a zone geometry. RS_ZonalStats returns one statistic (count, sum, mean, median, mode, stddev, variance, min, max) as a Float64; RS_ZonalStatsAll returns them all as a struct. Optional settings (band, all_touched, exclude_nodata, lenient) travel as a single JSON options argument rather than a positional overload ladder. Nodata is compared in the band's own byte representation to avoid a lossy f64 round-trip, and the band value type is dispatched once outside the pixel loop. Includes a criterion benchmark over raster resolution and zone complexity. --- Cargo.lock | 2 + rust/sedona-raster-gdal/Cargo.toml | 7 + .../benches/rs_zonalstats.rs | 177 +++ rust/sedona-raster-gdal/src/lib.rs | 1 + rust/sedona-raster-gdal/src/register.rs | 2 + rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 1396 +++++++++++++++++ 6 files changed, 1585 insertions(+) create mode 100644 rust/sedona-raster-gdal/benches/rs_zonalstats.rs create mode 100644 rust/sedona-raster-gdal/src/rs_zonal_stats.rs diff --git a/Cargo.lock b/Cargo.lock index 4818fa196a..a2418e97ab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6259,6 +6259,8 @@ dependencies = [ "sedona-raster-functions", "sedona-schema", "sedona-testing", + "serde", + "serde_json", "tempfile", "tokio", ] diff --git a/rust/sedona-raster-gdal/Cargo.toml b/rust/sedona-raster-gdal/Cargo.toml index de0cf3ce89..2c3dc6eca6 100644 --- a/rust/sedona-raster-gdal/Cargo.toml +++ b/rust/sedona-raster-gdal/Cargo.toml @@ -46,6 +46,8 @@ sedona-proj = { workspace = true } sedona-raster = { workspace = true } sedona-raster-functions = { workspace = true } sedona-schema = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } tokio = { workspace = true } [features] @@ -89,3 +91,8 @@ path = "benches/rs_polygonize.rs" harness = false name = "rs_as_raster" path = "benches/rs_as_raster.rs" + +[[bench]] +harness = false +name = "rs_zonalstats" +path = "benches/rs_zonalstats.rs" diff --git a/rust/sedona-raster-gdal/benches/rs_zonalstats.rs b/rust/sedona-raster-gdal/benches/rs_zonalstats.rs new file mode 100644 index 0000000000..ec46a69ae5 --- /dev/null +++ b/rust/sedona-raster-gdal/benches/rs_zonalstats.rs @@ -0,0 +1,177 @@ +// 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_ZonalStats / RS_ZonalStatsAll UDFs. +//! +//! Both functions rasterize the zone geometry into a mask, walk the masked +//! window collecting the selected pixel values, and reduce them to statistics. +//! +//! Each case builds a raster whose world extent is exactly the zone-polygon +//! generator's `[-10, 10]²` bounds at the requested resolution, so every +//! generated polygon lands on the raster and the full mask/collect/reduce path +//! runs. `all_touched = true` (passed via the JSON options) guarantees a +//! polygon smaller than a cell still burns at least one pixel rather than +//! hitting the empty-zone early return. +//! +//! Axes: +//! - **Raster resolution** (`64²`, `256²`, `1024²`) with a small polygon: +//! rasterization + window scan dominate. +//! - **Zone polygon complexity** (vertex count) at a fixed resolution, driving +//! the GDAL rasterization cost. +//! - **Large zone**: a polygon covering most of the raster, so collecting the +//! masked values and the reduction (sort for median, frequency map for mode) +//! dominate — the `RS_ZonalStatsAll` case is the heaviest since it computes +//! every statistic. +//! +//! Numerical correctness against a reference (rasterio / numpy) is pinned by +//! the Python parity tests, not here; this bench only measures throughput. + +use std::sync::Arc; + +use arrow_array::{ArrayRef, BinaryArray, StringArray}; +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 stats_udf: ScalarUDF = f + .scalar_udf("rs_zonalstats") + .expect("rs_zonalstats is registered") + .clone() + .into(); + let stats_all_udf: ScalarUDF = f + .scalar_udf("rs_zonalstatsall") + .expect("rs_zonalstatsall is registered") + .clone() + .into(); + + // RS_ZonalStats(raster, zone, stat, options) and + // RS_ZonalStatsAll(raster, zone, options). + let stats_tester = ScalarUdfTester::new( + stats_udf, + vec![ + RASTER, + WKB_GEOMETRY, + SedonaType::Arrow(arrow_schema::DataType::Utf8), + SedonaType::Arrow(arrow_schema::DataType::Utf8), + ], + ); + let stats_all_tester = ScalarUdfTester::new( + stats_all_udf, + vec![ + RASTER, + WKB_GEOMETRY, + SedonaType::Arrow(arrow_schema::DataType::Utf8), + ], + ); + + let mean_stat: ArrayRef = Arc::new(StringArray::from(vec!["mean"])); + let all_touched_opts: ArrayRef = Arc::new(StringArray::from(vec![ + r#"{"band": 1, "all_touched": 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 = (0..(w * h)).map(|v| v as f64).collect(); + Arc::new( + RasterSpec::d2(w, h) + .band_values(&values) + .crs(None) + .transform(transform) + .build(), + ) + }; + + let gen_polygon = |vertices: usize| -> ArrayRef { + BenchmarkArgSpec::Polygon(vertices) + .build_arrays(0, 1, 1) + .expect("build zone polygon") + .remove(0) + }; + + let run_single = |c: &mut Criterion, label: &str, raster: ArrayRef, geom: ArrayRef| { + c.bench_function(label, |b| { + b.iter(|| { + stats_tester + .invoke_arrays(vec![ + raster.clone(), + geom.clone(), + mean_stat.clone(), + all_touched_opts.clone(), + ]) + .unwrap() + }) + }); + }; + + let run_all = |c: &mut Criterion, label: &str, raster: ArrayRef, geom: ArrayRef| { + c.bench_function(label, |b| { + b.iter(|| { + stats_all_tester + .invoke_arrays(vec![raster.clone(), geom.clone(), all_touched_opts.clone()]) + .unwrap() + }) + }); + }; + + // Resolution sweep (simple 8-vertex polygon), single-stat mean. + for (w, h) in [(64i64, 64i64), (256, 256), (1024, 1024)] { + let label = + format!("raster-gdal rs_zonalstats ZonalStats(Raster({w}x{h}), Polygon(8), mean)"); + run_single(c, &label, build_raster(w, h), gen_polygon(8)); + } + + // Zone-complexity axis at a fixed 64×64 resolution. + run_single( + c, + "raster-gdal rs_zonalstats ZonalStats(Raster(64x64), Polygon(50), mean)", + build_raster(64, 64), + gen_polygon(50), + ); + + // Large zone: the polygon covers nearly the whole raster, so collecting the + // masked values and the reduction dominate. RS_ZonalStatsAll is the heaviest + // (median sort + mode frequency map over every selected pixel). + 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_single( + c, + "raster-gdal rs_zonalstats ZonalStats(Raster(1024x1024), Polygon(large), mean)", + build_raster(1024, 1024), + big_geom(), + ); + run_all( + c, + "raster-gdal rs_zonalstats ZonalStatsAll(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 74065b7055..f087a2d581 100644 --- a/rust/sedona-raster-gdal/src/lib.rs +++ b/rust/sedona-raster-gdal/src/lib.rs @@ -39,6 +39,7 @@ mod rs_clip; mod rs_frompath; mod rs_metadata; mod rs_polygonize; +mod rs_zonal_stats; mod source_uri; mod utils; diff --git a/rust/sedona-raster-gdal/src/register.rs b/rust/sedona-raster-gdal/src/register.rs index a30a13a05a..2ecc242ce0 100644 --- a/rust/sedona-raster-gdal/src/register.rs +++ b/rust/sedona-raster-gdal/src/register.rs @@ -26,5 +26,7 @@ pub fn default_function_set() -> FunctionSet { function_set.insert_scalar_udf(crate::rs_frompath::rs_frompath_udf()); function_set.insert_scalar_udf(crate::rs_metadata::rs_metadata_udf()); function_set.insert_scalar_udf(crate::rs_polygonize::rs_polygonize_udf()); + function_set.insert_scalar_udf(crate::rs_zonal_stats::rs_zonal_stats_udf()); + function_set.insert_scalar_udf(crate::rs_zonal_stats::rs_zonal_stats_all_udf()); function_set } diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs new file mode 100644 index 0000000000..2e599c5414 --- /dev/null +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -0,0 +1,1396 @@ +// 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 zone geometry. +//! +//! `RS_ZonalStats(raster, zone, stat[, options])` returns one statistic as a +//! `Float64`; `RS_ZonalStatsAll(raster, zone[, options])` returns every +//! statistic as a struct. Both compute over the pixels of a single band whose +//! centre falls inside the zone (or that the zone merely touches, with +//! `all_touched`), optionally excluding the band's nodata value. +//! +//! The optional trailing `options` argument is a JSON object rather than the +//! positional `band, all_touched, exclude_nodata, lenient` overload ladder of +//! Sedona Spark, so the surface stays a single `(input, options)` shape: +//! +//! ```json +//! {"band": 1, "all_touched": false, "exclude_nodata": true, "lenient": true} +//! ``` +//! +//! 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, StringArray}; +use arrow_schema::{DataType, Field, Fields}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::error::Result; +use datafusion_common::{exec_datafusion_err, exec_err}; +use datafusion_expr::{ColumnarValue, Volatility}; +use serde::Deserialize; + +use sedona_common::sedona_internal_err; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +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_raster::array::RasterRefImpl; +use sedona_raster::traits::RasterRef; +use sedona_raster_functions::crs_utils::{crs_transform_wkb, resolve_crs, with_crs_engine}; +use sedona_raster_functions::rs_ensure_loaded::NEEDS_PIXELS_METADATA_KEY; +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; + +/// 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 { + 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, + } + } +} + +/// Optional `(input, options)` JSON payload. Missing fields take their default; +/// unknown fields are rejected so a typo surfaces rather than being ignored. +#[derive(Debug, Clone, Deserialize)] +#[serde(default, deny_unknown_fields, rename_all = "snake_case")] +struct ZonalStatsOptions { + /// 1-based band to compute over. `None` means "resolve the default": band 1 + /// for a single-band raster, an error for a multiband raster (naming the + /// band is required rather than silently getting band 1). + band: Option, + /// Include every pixel the zone touches, not only those whose centre it + /// covers. + all_touched: bool, + /// Skip pixels equal to the band's nodata value. + exclude_nodata: bool, + /// Return NULL when the zone 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, +} + +impl Default for ZonalStatsOptions { + fn default() -> Self { + Self { + band: None, + all_touched: false, + exclude_nodata: true, + lenient: true, + } + } +} + +impl ZonalStatsOptions { + /// Parse the JSON options string, or return the defaults when no options + /// argument was supplied (`None`) or it is SQL NULL. + fn parse(json: Option<&str>) -> Result { + match json { + None => Ok(Self::default()), + Some(s) => serde_json::from_str(s) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: invalid options JSON: {e}")), + } + } +} + +/// Every statistic for a zone. `count` is always present (0 when the zone +/// 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, + mean: Option, + median: Option, + mode: Option, + stddev: Option, + variance: Option, + min: Option, + max: Option, +} + +impl ZonalStatistics { + /// The value RS_ZonalStats returns for a single statistic. `count` is never + /// NULL (it is 0 for an empty zone); the others are NULL for an empty zone. + fn get(&self, stat_type: StatType) -> Option { + 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(raster, zone, stat[, options])` — one statistic as a +/// `Float64`. `stat` is a statistic name (`count`, `sum`, `mean`, `median`, +/// `mode`, `stddev`, `variance`, `min`, `max`). +pub fn rs_zonal_stats_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_zonalstats", + vec![ + Arc::new(RsZonalStats { + with_options: false, + }), + Arc::new(RsZonalStats { with_options: true }), + ], + 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 { + /// Whether this kernel matches the trailing JSON `options` argument. + with_options: bool, +} + +impl SedonaScalarKernel for RsZonalStats { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let mut matchers = vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_string(), + ]; + if self.with_options { + matchers.push(ArgMatcher::is_string()); + } + let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(DataType::Float64)); + 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 = RasterExecutor::num_iterations_over(args); + + // stat is at index 2, options (when present) at index 3. + let stat_array = expand_string_arg(&args[2], num_iterations)?; + let mut stat_iter = stat_array.iter(); + let options_array = self + .with_options + .then(|| expand_string_arg(&args[3], num_iterations)) + .transpose()?; + let mut options_iter = options_array.as_ref().map(|a| a.iter()); + + let mut builder = Float64Builder::with_capacity(num_iterations); + let mut scratch: Vec = Vec::new(); + + // The executor only sees (raster, zone); the stat/options 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 options_str = options_iter.as_mut().and_then(|i| i.next().flatten()); + let options = ZonalStatsOptions::parse(options_str)?; + + // A NULL stat, raster, or zone 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 zone into the raster's CRS; a known/unknown + // CRS mismatch on either side would mislocate it. + 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(wkb, geom_crs, raster_crs, engine)? + } + (None, None) => 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" + ), + }; + match compute_zonal_stats(gdal, raster, &geom_wkb, &options, &mut scratch)? { + Some(stats) => match stats.get(stat_type) { + Some(value) => builder.append_value(value), + None => builder.append_null(), + }, + // The zone does not intersect the raster: NULL when + // lenient (the default), an error otherwise. + None if options.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(raster, zone[, options])` — every statistic as a struct +/// with fields `count, sum, mean, median, mode, stddev, variance, min, max`. +pub fn rs_zonal_stats_all_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_zonalstatsall", + vec![ + Arc::new(RsZonalStatsAll { + with_options: false, + }), + Arc::new(RsZonalStatsAll { with_options: true }), + ], + Volatility::Immutable, + ) + .with_metadata(NEEDS_PIXELS_METADATA_KEY, "true") +} + +#[derive(Debug)] +struct RsZonalStatsAll { + with_options: bool, +} + +impl SedonaScalarKernel for RsZonalStatsAll { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let mut matchers = vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ]; + if self.with_options { + matchers.push(ArgMatcher::is_string()); + } + 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 { + 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 = RasterExecutor::num_iterations_over(args); + + // options (when present) is at index 2. + let options_array = self + .with_options + .then(|| expand_string_arg(&args[2], num_iterations)) + .transpose()?; + let mut options_iter = options_array.as_ref().map(|a| a.iter()); + + let mut builder = StructBuilder::from_fields(zonal_stats_struct_fields(), num_iterations); + let mut scratch: Vec = 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 options_str = options_iter.as_mut().and_then(|i| i.next().flatten()); + let options = ZonalStatsOptions::parse(options_str)?; + + 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 = match (geom_crs, raster_crs.as_deref()) { + (Some(geom_crs), Some(raster_crs)) => { + crs_transform_wkb(wkb, geom_crs, raster_crs, engine)? + } + (None, None) => 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" + ), + }; + match compute_zonal_stats(gdal, raster, &geom_wkb, &options, &mut scratch)? { + Some(stats) => append_struct_stats(&mut builder, &stats), + None if options.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 zone does not intersect the raster and +/// `lenient` is set). +fn append_struct_null(builder: &mut StructBuilder) { + builder + .field_builder::(0) + .unwrap() + .append_null(); + for i in 1..=8 { + builder + .field_builder::(i) + .unwrap() + .append_null(); + } + builder.append(false); +} + +/// Append one computed-stats row. The float fields carry through the `Option` +/// so an empty zone records `count = 0` with the rest NULL. +fn append_struct_stats(builder: &mut StructBuilder, stats: &ZonalStatistics) { + builder + .field_builder::(0) + .unwrap() + .append_value(stats.count); + for (i, value) in [ + stats.sum, + stats.mean, + stats.median, + stats.mode, + stats.stddev, + stats.variance, + stats.min, + stats.max, + ] + .into_iter() + .enumerate() + { + builder + .field_builder::(i + 1) + .unwrap() + .append_option(value); + } + builder.append(true); +} + +// ============================================================================= +// Core computation +// ============================================================================= + +/// A rectangular pixel window (offset + size) into the raster grid. +#[derive(Clone, Copy, Debug)] +struct PixelWindow { + col_off: usize, + row_off: usize, + width: usize, + height: usize, +} + +/// Compute the statistics of the pixels a zone geometry selects on one band. +/// +/// Returns `Ok(None)` when the zone's envelope does not intersect the raster +/// extent (the caller turns this into NULL when `lenient`, an error otherwise). +/// A zone that intersects the extent but whose selected pixels are all outside +/// the geometry or all nodata yields `Ok(Some(..))` with `count = 0`. +/// +/// `scratch` is a reused buffer for the selected pixel values so the per-row +/// collection does not allocate a fresh `Vec` each call. +fn compute_zonal_stats( + gdal: &Gdal, + raster: &RasterRefImpl<'_>, + geom_wkb: &[u8], + options: &ZonalStatsOptions, + scratch: &mut Vec, +) -> Result> { + let num_bands = raster.num_bands(); + let band_num = resolve_band(options.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_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"))?; + + // Parse the zone and clamp its envelope to the raster grid. A disjoint + // envelope is the no-intersection case. + 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 { + return Ok(None); + }; + + // Rasterize the zone into a window-sized 0/1 mask (moves `geometry`, whose + // only remaining use is the burn). + let mask = rasterize_zone_mask(gdal, geometry, &transform, &window, options.all_touched)?; + + // 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 options.exclude_nodata { + band.nodata() + } else { + None + }; + + scratch.clear(); + collect_masked_values( + band_bytes, data_type, width, &window, &mask, nodata, scratch, + ); + + Ok(Some(compute_statistics(scratch))) +} + +/// 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, num_bands: usize) -> Result { + 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; set the \"band\" option to \ + choose one (only a single-band raster may omit it)" + ) + } + } + } +} + +/// The raster's 6-coefficient GDAL geotransform as a fixed array. +fn raster_transform(raster: &RasterRefImpl<'_>) -> Result { + let t = raster.transform(); + <[f64; 6]>::try_from(t) + .map_err(|_| exec_datafusion_err!("RS_ZonalStats: expected a 6-element geotransform")) +} + +/// The zone's envelope intersected with the raster extent, snapped outward to +/// whole pixels. `None` when the envelope is disjoint from the raster. Mirrors +/// the window RS_Clip / PostGIS ST_Clip use: all four corners are mapped +/// through the inverse geotransform so a skewed raster still gets a correct +/// superset window. +fn envelope_window( + geometry: &Geometry, + transform: &GeoTransform, + width: usize, + height: usize, +) -> Result> { + let env = geometry.envelope(); + let inverse = transform + .invert() + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: 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(PixelWindow { + col_off: col0 as usize, + row_off: row0 as usize, + width: (col1 - col0) as usize, + height: (row1 - row0) as usize, + })) +} + +/// Rasterize the zone into a window-sized `u8` mask: 1 where the zone covers a +/// pixel, 0 elsewhere. The MEM band is zero-filled on creation, so only the +/// burn (value 1, inside the geometry) has to be written. +fn rasterize_zone_mask( + gdal: &Gdal, + geometry: Geometry, + transform: &GeoTransform, + window: &PixelWindow, + all_touched: bool, +) -> Result> { + let mask_dataset = + MemDatasetBuilder::create(gdal, window.width, window.height, 1, GdalDataType::UInt8) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to create mask: {e}"))?; + let (window_ulx, window_uly) = transform.apply(window.col_off as f64, window.row_off as f64); + let mask_transform = [ + window_ulx, + transform[1], + transform[2], + window_uly, + transform[4], + transform[5], + ]; + mask_dataset + .set_geo_transform(&mask_transform) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to set mask geotransform: {e}"))?; + + gdal.rasterize_affine(&mask_dataset, &[1], &[geometry], &[1.0], all_touched) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to rasterize zone: {e}"))?; + + let mask_band = mask_dataset + .rasterband(1) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to read 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!("RS_ZonalStats: failed to read mask: {e}"))?; + Ok(mask_buffer.data().to_vec()) +} + +/// 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, +) { + 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. +/// +/// An empty slice yields `count = 0` and NULL for the rest (Sedona Spark's +/// empty-zone 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 for the median 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); 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, + }; + } + + let n = values.len(); + let sum: f64 = values.iter().sum(); + let mean = sum / n 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 = if n > 1 { + let sum_sq: f64 = values.iter().map(|&v| (v - mean).powi(2)).sum(); + sum_sq / (n as f64 - 1.0) + } else { + 0.0 + }; + let stddev = variance.sqrt(); + + let mode = compute_mode(values); + + // Median needs the values sorted; do it in place on the scratch buffer. + values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + let mid = n / 2; + let median = if n.is_multiple_of(2) { + (values[mid - 1] + values[mid]) / 2.0 + } else { + values[mid] + }; + + 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), + } +} + +/// The most frequent value, breaking ties toward the larger value (matching +/// Sedona Spark's `StatUtils.mode`, which returns the largest of the tied +/// modes). Values are keyed by their exact bit pattern, so integer-valued +/// pixels compare exactly. +fn compute_mode(values: &[f64]) -> f64 { + let mut counts: HashMap = HashMap::new(); + for &v in values { + *counts.entry(v.to_bits()).or_insert(0) += 1; + } + let best_count = counts.values().copied().max().unwrap_or(0); + counts + .into_iter() + .filter(|&(_, c)| c == best_count) + .map(|(bits, _)| f64::from_bits(bits)) + .fold(f64::NEG_INFINITY, f64::max) +} + +// ============================================================================= +// Argument helpers +// ============================================================================= + +/// The error returned for a non-intersecting zone when `lenient` is off. +fn no_intersection_err() -> Result { + exec_err!( + "RS_ZonalStats: the zone geometry does not intersect the raster; \ + set the \"lenient\" option to return NULL instead" + ) +} + +/// Cast a column to `Utf8` and materialize it to a `StringArray` so its values +/// can be iterated in lockstep with the raster/zone rows. +fn expand_string_arg(arg: &ColumnarValue, num_iterations: usize) -> Result { + let array = arg + .clone() + .cast_to(&DataType::Utf8, None)? + .into_array(num_iterations)?; + Ok(datafusion_common::cast::as_string_array(&array)?.clone()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn stat_type_from_str_matches_spark_aliases() { + assert_eq!(StatType::from_str("count"), Some(StatType::Count)); + assert_eq!(StatType::from_str("COUNT"), Some(StatType::Count)); + assert_eq!(StatType::from_str("mean"), Some(StatType::Mean)); + assert_eq!(StatType::from_str("avg"), Some(StatType::Mean)); + assert_eq!(StatType::from_str("average"), Some(StatType::Mean)); + assert_eq!(StatType::from_str("stddev"), Some(StatType::StdDev)); + assert_eq!(StatType::from_str("sd"), Some(StatType::StdDev)); + assert_eq!(StatType::from_str("variance"), Some(StatType::Variance)); + assert_eq!(StatType::from_str("min"), Some(StatType::Min)); + assert_eq!(StatType::from_str("max"), Some(StatType::Max)); + assert_eq!(StatType::from_str("nonsense"), None); + } + + #[test] + fn options_parse_defaults_and_overrides() { + let d = ZonalStatsOptions::parse(None).unwrap(); + assert_eq!(d.band, None); + assert!(!d.all_touched); + assert!(d.exclude_nodata); + assert!(d.lenient); + + let o = ZonalStatsOptions::parse(Some( + r#"{"band": 2, "all_touched": true, "exclude_nodata": false, "lenient": false}"#, + )) + .unwrap(); + assert_eq!(o.band, Some(2)); + assert!(o.all_touched); + assert!(!o.exclude_nodata); + assert!(!o.lenient); + + // A partial object keeps the defaults for the omitted fields. + let p = ZonalStatsOptions::parse(Some(r#"{"band": 3}"#)).unwrap(); + assert_eq!(p.band, Some(3)); + assert!(p.exclude_nodata); + } + + #[test] + fn options_parse_rejects_unknown_field() { + let err = ZonalStatsOptions::parse(Some(r#"{"bnad": 1}"#)).unwrap_err(); + assert!(err.to_string().contains("invalid options JSON"), "{err}"); + } + + #[test] + fn resolve_band_defaults_and_bounds() { + // Single-band raster may omit the band. + assert_eq!(resolve_band(None, 1).unwrap(), 1); + // Multiband raster must name the band. + let err = resolve_band(None, 3).unwrap_err().to_string(); + assert!(err.contains("has 3 bands"), "{err}"); + // Explicit band is range-checked (1-based). + assert_eq!(resolve_band(Some(2), 3).unwrap(), 2); + assert!(resolve_band(Some(0), 3) + .unwrap_err() + .to_string() + .contains(">= 1")); + assert!(resolve_band(Some(4), 3) + .unwrap_err() + .to_string() + .contains("out of range")); + } + + #[test] + fn statistics_of_one_to_five() { + // count, sum, mean, min, max, median are exact; variance/stddev are the + // sample (n-1) values: ((1-3)^2+..+(5-3)^2)/4 = 10/4 = 2.5. + let mut values = vec![1.0, 2.0, 3.0, 4.0, 5.0]; + let s = compute_statistics(&mut values); + assert_eq!(s.count, 5); + assert_eq!(s.sum, Some(15.0)); + assert_eq!(s.mean, Some(3.0)); + assert_eq!(s.min, Some(1.0)); + assert_eq!(s.max, Some(5.0)); + assert_eq!(s.median, Some(3.0)); + assert_eq!(s.variance, Some(2.5)); + assert_eq!(s.stddev, Some(2.5_f64.sqrt())); + } + + #[test] + fn statistics_empty_is_zero_count_and_nulls() { + let mut values: Vec = vec![]; + let s = compute_statistics(&mut values); + assert_eq!(s.count, 0); + assert_eq!(s.sum, None); + assert_eq!(s.mean, None); + assert_eq!(s.median, None); + assert_eq!(s.min, None); + assert_eq!(s.max, None); + assert_eq!(s.variance, None); + // A single-stat lookup returns 0 for count, NULL for the rest. + assert_eq!(s.get(StatType::Count), Some(0.0)); + assert_eq!(s.get(StatType::Sum), None); + assert_eq!(s.get(StatType::Mean), None); + } + + #[test] + fn statistics_single_value_has_zero_variance() { + let mut values = vec![42.0]; + let s = compute_statistics(&mut values); + assert_eq!(s.count, 1); + assert_eq!(s.mean, Some(42.0)); + assert_eq!(s.median, Some(42.0)); + assert_eq!(s.variance, Some(0.0)); + assert_eq!(s.stddev, Some(0.0)); + assert_eq!(s.mode, Some(42.0)); + } + + #[test] + fn median_even_count_averages_the_middle_pair() { + let mut values = vec![4.0, 1.0, 3.0, 2.0]; + let s = compute_statistics(&mut values); + assert_eq!(s.median, Some(2.5)); + } + + #[test] + fn mode_breaks_ties_toward_the_larger_value() { + // 1 and 3 each appear twice; the tie resolves to the larger, 3. + let mut values = vec![1.0, 1.0, 3.0, 3.0, 2.0]; + assert_eq!(compute_statistics(&mut values).mode, Some(3.0)); + // A clear winner is returned as-is. + let mut values = vec![7.0, 7.0, 7.0, 1.0, 2.0]; + assert_eq!(compute_statistics(&mut values).mode, Some(7.0)); + } +} + +/// UDF-level tests: exercise the kernels end to end and pin the numbers against +/// values computed by hand (which agree with numpy — see the Python parity +/// tests for the rasterio/numpy cross-check). +#[cfg(test)] +mod udf_tests { + use super::*; + + use std::sync::Arc; + + use arrow_array::cast::AsArray; + use arrow_array::types::{Float64Type, Int64Type}; + use arrow_array::{Array, StructArray}; + use datafusion_common::ScalarValue; + use datafusion_expr::ScalarUDF; + use sedona_proj::transform::{with_global_proj_engine, LazyProjEngine}; + use sedona_schema::crs::deserialize_crs; + use sedona_schema::datatypes::{Edges, RASTER}; + use sedona_testing::create::make_wkb; + use sedona_testing::raster_spec::RasterSpec; + use sedona_testing::testers::ScalarUdfTester; + + // Struct field positions (Sedona Spark order). + const COUNT: usize = 0; + const SUM: usize = 1; + const MEAN: usize = 2; + const MEDIAN: usize = 3; + const MODE: usize = 4; + const STDDEV: usize = 5; + const VARIANCE: usize = 6; + const MIN: usize = 7; + const MAX: usize = 8; + + /// A 4×2 UInt8 raster with pixel values 1..=8 (row-major), world extent + /// x ∈ [0, 4], y ∈ [0, 2] with 1×1 north-up pixels. Pixel centres: + /// row y=1.5 → 1,2,3,4 at x=0.5,1.5,2.5,3.5; row y=0.5 → 5,6,7,8. + fn small_raster() -> RasterSpec { + 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]) + } + + /// The left half of `small_raster` (x ∈ [0, 2]) selects the four pixels + /// {1, 2, 5, 6}. + const LEFT_HALF: &str = "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))"; + + fn stats_arg_types(with_options: bool) -> Vec { + let mut t = vec![ + RASTER, + SedonaType::Wkb(Edges::Planar, None), + SedonaType::Arrow(DataType::Utf8), + ]; + if with_options { + t.push(SedonaType::Arrow(DataType::Utf8)); + } + t + } + + /// Invoke RS_ZonalStats on a scalar raster + zone, returning the raw result + /// so both the value and error paths are testable. + fn call_stats( + spec: &RasterSpec, + wkt: &str, + stat: &str, + options: Option<&str>, + ) -> Result { + let kernel = RsZonalStats { + with_options: options.is_some(), + }; + let mut args = vec![ + ColumnarValue::Scalar(spec.scalar()), + ColumnarValue::Scalar(ScalarValue::Binary(Some(make_wkb(wkt)))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(stat.to_string()))), + ]; + if let Some(o) = options { + args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some( + o.to_string(), + )))); + } + kernel.invoke_batch(&stats_arg_types(options.is_some()), &args) + } + + /// Invoke RS_ZonalStatsAll on a scalar raster + zone. + fn call_all(spec: &RasterSpec, wkt: &str, options: Option<&str>) -> Result { + let kernel = RsZonalStatsAll { + with_options: options.is_some(), + }; + let mut arg_types = vec![RASTER, SedonaType::Wkb(Edges::Planar, None)]; + let mut args = vec![ + ColumnarValue::Scalar(spec.scalar()), + ColumnarValue::Scalar(ScalarValue::Binary(Some(make_wkb(wkt)))), + ]; + if let Some(o) = options { + arg_types.push(SedonaType::Arrow(DataType::Utf8)); + args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some( + o.to_string(), + )))); + } + kernel.invoke_batch(&arg_types, &args) + } + + fn cv_f64(cv: ColumnarValue) -> Option { + match cv { + ColumnarValue::Scalar(ScalarValue::Float64(v)) => v, + other => panic!("expected a Float64 scalar, got {other:?}"), + } + } + + fn cv_struct(cv: ColumnarValue) -> Arc { + match cv { + ColumnarValue::Scalar(ScalarValue::Struct(s)) => s, + other => panic!("expected a struct scalar, got {other:?}"), + } + } + + fn f64_field(s: &StructArray, col: usize) -> Option { + let c = s.column(col); + (!c.is_null(0)).then(|| c.as_primitive::().value(0)) + } + + fn i64_field(s: &StructArray, col: usize) -> Option { + let c = s.column(col); + (!c.is_null(0)).then(|| c.as_primitive::().value(0)) + } + + #[test] + fn single_stats_match_hand_computed_values() { + let spec = small_raster(); + // Selected pixels {1, 2, 5, 6}: exact for the integer-selection stats. + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, "count", None).unwrap()), + Some(4.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, "sum", None).unwrap()), + Some(14.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, "mean", None).unwrap()), + Some(3.5) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, "min", None).unwrap()), + Some(1.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, "max", None).unwrap()), + Some(6.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, "median", None).unwrap()), + Some(3.5) + ); + // All four values are unique, so every one is a mode; the tie resolves + // to the largest (6), matching Sedona Spark. + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, "mode", None).unwrap()), + Some(6.0) + ); + + // Sample (n-1) variance / stddev: float accumulation, so approximate. + let var = cv_f64(call_stats(&spec, LEFT_HALF, "variance", None).unwrap()).unwrap(); + assert!((var - 17.0 / 3.0).abs() < 1e-9, "variance was {var}"); + let sd = cv_f64(call_stats(&spec, LEFT_HALF, "stddev", None).unwrap()).unwrap(); + assert!( + (sd - (17.0_f64 / 3.0).sqrt()).abs() < 1e-9, + "stddev was {sd}" + ); + } + + #[test] + fn all_returns_full_struct() { + let s = cv_struct(call_all(&small_raster(), LEFT_HALF, None).unwrap()); + assert!(!s.is_null(0), "the struct itself is valid"); + assert_eq!(i64_field(&s, COUNT), Some(4)); + assert_eq!(f64_field(&s, SUM), Some(14.0)); + assert_eq!(f64_field(&s, MEAN), Some(3.5)); + assert_eq!(f64_field(&s, MEDIAN), Some(3.5)); + assert_eq!(f64_field(&s, MODE), Some(6.0)); + assert_eq!(f64_field(&s, MIN), Some(1.0)); + assert_eq!(f64_field(&s, MAX), Some(6.0)); + assert!((f64_field(&s, VARIANCE).unwrap() - 17.0 / 3.0).abs() < 1e-9); + assert!((f64_field(&s, STDDEV).unwrap() - (17.0_f64 / 3.0).sqrt()).abs() < 1e-9); + } + + #[test] + fn zone_that_selects_no_pixel_centre_is_count_zero_not_null() { + // A tiny zone inside the top-left pixel (centre 0.5, 1.5) but not + // covering that centre: with all_touched off, no pixel is selected. The + // zone still overlaps the raster extent, so count is 0 (not NULL). + let tiny = "POLYGON((0.1 1.6, 0.4 1.6, 0.4 1.9, 0.1 1.9, 0.1 1.6))"; + assert_eq!( + cv_f64(call_stats(&small_raster(), tiny, "count", None).unwrap()), + Some(0.0) + ); + assert_eq!( + cv_f64(call_stats(&small_raster(), tiny, "sum", None).unwrap()), + None + ); + assert_eq!( + cv_f64(call_stats(&small_raster(), tiny, "mean", None).unwrap()), + None + ); + + let s = cv_struct(call_all(&small_raster(), tiny, None).unwrap()); + assert!( + !s.is_null(0), + "an intersecting-but-empty zone is a valid row" + ); + assert_eq!(i64_field(&s, COUNT), Some(0)); + assert_eq!(f64_field(&s, SUM), None); + assert_eq!(f64_field(&s, MEAN), None); + } + + #[test] + fn all_touched_selects_the_touched_pixel() { + // The same tiny zone, with all_touched, burns the pixel it lies inside + // (value 1) even though it misses the centre. + let tiny = "POLYGON((0.1 1.6, 0.4 1.6, 0.4 1.9, 0.1 1.9, 0.1 1.6))"; + let opts = r#"{"all_touched": true}"#; + assert_eq!( + cv_f64(call_stats(&small_raster(), tiny, "count", Some(opts)).unwrap()), + Some(1.0) + ); + assert_eq!( + cv_f64(call_stats(&small_raster(), tiny, "sum", Some(opts)).unwrap()), + Some(1.0) + ); + } + + #[test] + fn no_intersection_is_null_when_lenient_and_errors_when_strict() { + let far = "POLYGON((100 100, 101 100, 101 101, 100 101, 100 100))"; + // Lenient (default): the whole value is NULL, including count. + assert_eq!( + cv_f64(call_stats(&small_raster(), far, "count", None).unwrap()), + None + ); + assert!(cv_struct(call_all(&small_raster(), far, None).unwrap()).is_null(0)); + + // Strict: both functions error. + let err = call_stats(&small_raster(), far, "count", Some(r#"{"lenient": false}"#)) + .unwrap_err() + .to_string(); + assert!( + err.contains("does not intersect"), + "unexpected error: {err}" + ); + let err = call_all(&small_raster(), far, Some(r#"{"lenient": false}"#)) + .unwrap_err() + .to_string(); + assert!( + err.contains("does not intersect"), + "unexpected error: {err}" + ); + } + + #[test] + fn nodata_pixels_are_excluded_by_default_and_kept_when_asked() { + // A 2×2 UInt8 raster [10, 255, 20, 30] with nodata 255, world extent + // x ∈ [0, 2], y ∈ [0, 2]; the zone covers all four pixels. + let spec = RasterSpec::d2(2, 2) + .band_values(&[10u8, 255, 20, 30]) + .nodata(255u8) + .crs(None) + .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]); + // Default excludes the nodata pixel: {10, 20, 30}. + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, "count", None).unwrap()), + Some(3.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, "sum", None).unwrap()), + Some(60.0) + ); + // exclude_nodata=false keeps it: {10, 255, 20, 30}. + let keep = r#"{"exclude_nodata": false}"#; + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, "count", Some(keep)).unwrap()), + Some(4.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, "sum", Some(keep)).unwrap()), + Some(315.0) + ); + } + + /// A zone covering the whole 2×2 nodata raster above. + const LEFT_HALF_FULL: &str = "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))"; + + #[test] + fn multiband_raster_requires_the_band_option() { + let spec = RasterSpec::d2(2, 2) + .band_values(&[1u8, 2, 3, 4]) + .band_values(&[10u8, 20, 30, 40]) + .crs(None) + .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]); + // Omitting the band on a multiband raster errors rather than defaulting. + let err = call_stats(&spec, LEFT_HALF_FULL, "sum", None) + .unwrap_err() + .to_string(); + assert!(err.contains("2 bands"), "unexpected error: {err}"); + // Naming the band selects it (band 1 sums to 10, band 2 to 100). + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, "sum", Some(r#"{"band": 1}"#)).unwrap()), + Some(10.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, "sum", Some(r#"{"band": 2}"#)).unwrap()), + Some(100.0) + ); + // An out-of-range band errors. + let err = call_stats(&spec, LEFT_HALF_FULL, "sum", Some(r#"{"band": 3}"#)) + .unwrap_err() + .to_string(); + assert!(err.contains("out of range"), "unexpected error: {err}"); + } + + #[test] + fn unknown_statistic_errors() { + let err = call_stats(&small_raster(), LEFT_HALF, "bogus", None) + .unwrap_err() + .to_string(); + assert!(err.contains("unknown statistic"), "unexpected error: {err}"); + } + + #[test] + fn null_raster_or_zone_yields_null() { + // A NULL zone geometry propagates to a NULL result. + let kernel = RsZonalStats { + with_options: false, + }; + let result = kernel + .invoke_batch( + &stats_arg_types(false), + &[ + ColumnarValue::Scalar(small_raster().scalar()), + ColumnarValue::Scalar(ScalarValue::Binary(None)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("count".to_string()))), + ], + ) + .unwrap(); + assert_eq!(cv_f64(result), None); + + // A NULL statistic name also yields NULL. + let result = kernel + .invoke_batch( + &stats_arg_types(false), + &[ + ColumnarValue::Scalar(small_raster().scalar()), + ColumnarValue::Scalar(ScalarValue::Binary(Some(make_wkb(LEFT_HALF)))), + ColumnarValue::Scalar(ScalarValue::Utf8(None)), + ], + ) + .unwrap(); + assert_eq!(cv_f64(result), None); + } + + #[test] + fn reprojects_the_zone_into_the_raster_crs() { + // The raster is EPSG:4326; the zone is supplied in EPSG:3857 (the + // reprojected LEFT_HALF polygon). Reprojecting it back to the raster CRS + // must recover the same four-pixel selection. + let spec = small_raster().crs(Some("EPSG:4326")); + let crs_4326 = deserialize_crs("EPSG:4326").unwrap().unwrap(); + let crs_3857 = deserialize_crs("EPSG:3857").unwrap().unwrap(); + let wkb_4326 = make_wkb(LEFT_HALF); + let wkb_3857 = with_global_proj_engine(|engine| { + crs_transform_wkb(&wkb_4326, crs_4326.as_ref(), crs_3857.as_ref(), engine) + }) + .unwrap(); + + let udf: ScalarUDF = rs_zonal_stats_udf().into(); + let arg_types = vec![ + RASTER, + SedonaType::Wkb(Edges::Planar, Some(crs_3857)), + SedonaType::Arrow(DataType::Utf8), + ]; + let tester = ScalarUdfTester::new(udf, arg_types).with_crs_engine(Arc::new(LazyProjEngine)); + let result = tester + .invoke_scalar_scalar_scalar(&spec, ScalarValue::Binary(Some(wkb_3857)), "count") + .unwrap(); + assert_eq!(result, ScalarValue::Float64(Some(4.0))); + } +} From cb3e3a06deb194a795716afa20376910cb77ba48 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 20 Jul 2026 13:08:11 -0700 Subject: [PATCH 02/16] docs(sql): document RS_ZonalStats and RS_ZonalStatsAll Mark both functions experimental and give runnable examples over RS_Example(). --- docs/reference/sql/rs_zonalstats.qmd | 102 ++++++++++++++++++++++++ docs/reference/sql/rs_zonalstatsall.qmd | 85 ++++++++++++++++++++ 2 files changed, 187 insertions(+) create mode 100644 docs/reference/sql/rs_zonalstats.qmd create mode 100644 docs/reference/sql/rs_zonalstatsall.qmd diff --git a/docs/reference/sql/rs_zonalstats.qmd b/docs/reference/sql/rs_zonalstats.qmd new file mode 100644 index 0000000000..6929642c60 --- /dev/null +++ b/docs/reference/sql/rs_zonalstats.qmd @@ -0,0 +1,102 @@ +--- +# 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_ZonalStats +description: > + Computes a single summary statistic of the raster pixels covered by a zone + geometry. +kernels: + - returns: double + args: + - raster + - {name: zone, type: geometry} + - {name: stat, type: string} + - returns: double + args: + - raster + - name: zone + type: geometry + description: > + Zone geometry defining the region of interest. Reprojected into the + raster's CRS when both carry one; it is an error for exactly one side to + have a CRS. + - name: stat + type: string + description: > + Statistic to return (case-insensitive): count, sum, mean, median, mode, + stddev, variance, min, or max. `avg`/`average` alias mean and `sd` + aliases stddev. + - name: options + type: string + description: > + JSON object of options. `band` is the 1-based band index (required for a + multiband raster; defaults to 1 for a single-band raster). `all_touched` + (default false) keeps every pixel the zone touches rather than only those + whose center it covers. `exclude_nodata` (default true) skips pixels equal + to the band's nodata value. `lenient` (default true) returns NULL when the + zone does not intersect the raster instead of raising an error. +--- + +::: callout-warning +**Experimental.** This function is experimental; its behavior may change without notice. +::: + +## Description + +`RS_ZonalStats` returns one summary statistic of the pixels of a single band +that a zone geometry covers, following Apache Sedona's `RS_ZonalStats`. A pixel +is included when its center falls inside the zone (or, with `all_touched`, when +the zone touches it at all). By default the band's nodata pixels are excluded. + +The statistic is one of `count`, `sum`, `mean`, `median`, `mode`, `stddev`, +`variance`, `min`, or `max`. `count` is returned as a whole number; all other +statistics are floating point. Variance and standard deviation are the sample +(n-1) values, and `mode` breaks ties toward the larger value. + +When the zone overlaps the raster but selects no pixel, `count` is 0 and every +other statistic is NULL. When the zone does not intersect the raster at all, the +result is NULL under the default `lenient` behavior, or an error when +`lenient` is set to false. + +The options are supplied as a single JSON object rather than as a positional +argument list, so the argument surface stays `(raster, zone, stat, options)`. +This function operates on 2-D `(y, x)` bands; computing a statistic per +non-spatial plane of an N-D band is not supported. + +Use [`RS_ZonalStatsAll`](rs_zonalstatsall.qmd) to compute every statistic at +once. + +## Examples + +```sql +SELECT RS_ZonalStats( + RS_Example(), + ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), + 'mean', + '{"band": 1}' +); +``` + +```sql +SELECT RS_ZonalStats( + RS_Example(), + ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), + 'count', + '{"band": 1, "exclude_nodata": false}' +); +``` diff --git a/docs/reference/sql/rs_zonalstatsall.qmd b/docs/reference/sql/rs_zonalstatsall.qmd new file mode 100644 index 0000000000..7b2bc69743 --- /dev/null +++ b/docs/reference/sql/rs_zonalstatsall.qmd @@ -0,0 +1,85 @@ +--- +# 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_ZonalStatsAll +description: > + Computes every summary statistic of the raster pixels covered by a zone + geometry and returns them as a struct. +kernels: + - returns: struct + args: + - raster + - {name: zone, type: geometry} + - returns: struct + args: + - raster + - name: zone + type: geometry + description: > + Zone geometry defining the region of interest. Reprojected into the + raster's CRS when both carry one; it is an error for exactly one side to + have a CRS. + - name: options + type: string + description: > + JSON object of options. `band` is the 1-based band index (required for a + multiband raster; defaults to 1 for a single-band raster). `all_touched` + (default false) keeps every pixel the zone touches rather than only those + whose center it covers. `exclude_nodata` (default true) skips pixels equal + to the band's nodata value. `lenient` (default true) returns NULL when the + zone does not intersect the raster instead of raising an error. +--- + +::: callout-warning +**Experimental.** This function is experimental; its behavior may change without notice. +::: + +## Description + +`RS_ZonalStatsAll` returns every summary statistic of the pixels of a single +band that a zone geometry covers, as a struct with fields `count`, `sum`, +`mean`, `median`, `mode`, `stddev`, `variance`, `min`, and `max` (the field +order of Apache Sedona's `RS_ZonalStatsAll`). A pixel is included when its +center falls inside the zone (or, with `all_touched`, when the zone touches it +at all), and the band's nodata pixels are excluded by default. + +`count` is a whole number; every other field is floating point. Variance and +standard deviation are the sample (n-1) values, and `mode` breaks ties toward +the larger value. + +When the zone overlaps the raster but selects no pixel, `count` is 0 and every +other field is NULL. When the zone does not intersect the raster at all, the +whole struct is NULL under the default `lenient` behavior, or the call raises an +error when `lenient` is set to false. + +The options are supplied as a single JSON object rather than as a positional +argument list, so the argument surface stays `(raster, zone, options)`. This +function operates on 2-D `(y, x)` bands; computing statistics per non-spatial +plane of an N-D band is not supported. + +Use [`RS_ZonalStats`](rs_zonalstats.qmd) to compute a single statistic. + +## Examples + +```sql +SELECT RS_ZonalStatsAll( + RS_Example(), + ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), + '{"band": 1}' +); +``` From 38f4333a8bc054371571676e460f3b216a76b265 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 20 Jul 2026 13:08:11 -0700 Subject: [PATCH 03/16] test(python/sedonadb): RS_ZonalStats / RS_ZonalStatsAll parity tests Cross-check every statistic against a rasterio + numpy reference through the dataframe API, with SQL-text smokes and error-path coverage. --- .../tests/functions/test_rs_zonalstats.py | 355 ++++++++++-------- 1 file changed, 200 insertions(+), 155 deletions(-) diff --git a/python/sedonadb/tests/functions/test_rs_zonalstats.py b/python/sedonadb/tests/functions/test_rs_zonalstats.py index cc9ae8daad..7fa93b5f05 100644 --- a/python/sedonadb/tests/functions/test_rs_zonalstats.py +++ b/python/sedonadb/tests/functions/test_rs_zonalstats.py @@ -15,195 +15,240 @@ # specific language governing permissions and limitations # under the License. -"""RS_ZonalStats parity against geometry_mask + numpy reductions. - -The rasterio comparator selects pixels with -`rasterio.features.geometry_mask`, drops pixels valued at the band nodata, -and reduces in float64. stddev/variance are the sample (ddof=1) statistics — -that is what Sedona computes. The diagonal-edged zone under the centroid -rule is on the Sedona Spark deviation ledger (its scanline rasterizer -mis-places x-intercepts on non-square pixels and drops some center-inside -pixels there, apache/sedona#3111). Zones that select no pixels are not -compared here. +"""RS_ZonalStats / RS_ZonalStatsAll cross-checked against a numpy reference. + +The fixture raster is CRS-less (so nothing reprojects and pixel selection is +bit-comparable). The reference rasterizes the zone with `rasterio.features` +(the same GDAL rasterizer the kernel uses) and reduces the selected pixels with +numpy; exact-selection statistics (count, sum, min, max, median, mode) are +compared exactly and the float-accumulation ones (mean, variance, stddev) with +a tolerance. + +rasterio is required to write the fixture GeoTIFF, so the whole module skips +when it is unavailable rather than importing it at module scope. """ +import json + +import numpy as np import pyarrow as pa import pytest -from sedonadb.raster_testing import ( - Deviation, - SedonaDB, - SedonaSpark, - expect_deviations, - random_raster_data, - write_geotiff, -) - pytest.importorskip("rasterio") -pytest.importorskip("shapely") -pytestmark = pytest.mark.skipif( - not SedonaDB.implements("zonal_stats"), - reason="RS_ZonalStats is not implemented in SedonaDB (the parity subject)", -) +from sedonadb.raster_testing import random_raster_data, write_geotiff # noqa: E402 -# GDAL-order geotransform: origin (100, 500), 2-wide by 3-tall north-up -# pixels; with a 7x6 raster the extent is x in [100, 114], y in [482, 500]. +# GDAL-order geotransform: origin (100, 500), 2-wide by 3-tall north-up pixels; +# a 6x7 raster then spans x in [100, 114], y in [482, 500]. GDAL_TRANSFORM = (100.0, 2.0, 0.0, 500.0, 0.0, -3.0) -HEIGHT, WIDTH = 6, 7 +BANDS, HEIGHT, WIDTH = 1, 6, 7 +NODATA = -9999 + +# A rectangle well inside the raster that selects a block of pixels. GEOM_RECT = ( "POLYGON ((102.6 495.8, 109.3 495.8, 109.3 485.9, 102.6 485.9, 102.6 495.8))" ) -# Diagonal edges make all_touched matter, while staying clear of the corner -# pixels where the fixture plants the dtype extremes (a float64 extreme in -# the zone would push the squared-deviation statistics to infinity). -GEOM_TRIANGLE = "POLYGON ((102.7 497.4, 112.4 496.9, 104.2 483.7, 102.7 497.4))" - -STATS = ["count", "sum", "mean", "min", "max", "stddev", "variance", "median"] - -DEVIATIONS = [ - Deviation( - SedonaSpark, - "zonal_stats", - matches=lambda p: p.get("wkt") == GEOM_TRIANGLE and not p.get("all_touched"), - reason="Sedona's scanline rasterizer mis-places x-intercepts on " - "non-square pixels and drops some center-inside pixels along " - "diagonal edges; GDAL selects every center-inside pixel " - "(https://github.com/apache/sedona/issues/3111)", - ), -] +# Entirely outside the raster extent. +GEOM_DISJOINT = "POLYGON ((900 900, 910 900, 910 890, 900 890, 900 900))" +# A thin strip crossing the x = 104 pixel boundary but covering no pixel center +# (centers sit at odd x): selects nothing unless all_touched. +GEOM_SLIVER = "POLYGON ((103.6 499, 104.4 499, 104.4 483, 103.6 483, 103.6 499))" +STATS = ["count", "sum", "mean", "median", "mode", "stddev", "variance", "min", "max"] +EXACT_STATS = {"count", "sum", "min", "max", "median", "mode"} -@pytest.mark.parametrize("stat", STATS) -@pytest.mark.parametrize( - ("wkt", "all_touched"), - [ - (GEOM_RECT, False), - (GEOM_RECT, True), - (GEOM_TRIANGLE, False), - (GEOM_TRIANGLE, True), - ], - ids=["rect-centroid", "rect-touched", "triangle-centroid", "triangle-touched"], -) -def test_rs_zonalstats_matches_comparators( - subject, comparator, request, tmp_path, wkt, all_touched, stat -): - """Every statistic over the float64 fixture, on both selection rules. - The zone stays clear of the corners so the planted dtype extremes don't - collapse sums to infinity.""" - expect_deviations(request, comparator, "zonal_stats", DEVIATIONS) - tiff = tmp_path / "zonal.tif" - write_geotiff( - tiff, - random_raster_data("float64", bands=2, height=HEIGHT, width=WIDTH), - gdal_transform=GDAL_TRANSFORM, - ) - got = subject.zonal_stats(tiff, wkt, band=2, stat=stat, all_touched=all_touched) - expected = comparator.zonal_stats( - tiff, wkt, band=2, stat=stat, all_touched=all_touched +def fixture_raster(tmp_path): + """A single-band int32 raster with planted nodata and a repeated value. + + Returns `(path, band)` where `band` is the `(HEIGHT, WIDTH)` numpy array. + Two interior pixels hold the nodata value and three hold a repeated value + (66) so the mode is unambiguous and nodata exclusion is observable. + """ + data = random_raster_data( + "int32", + bands=BANDS, + height=HEIGHT, + width=WIDTH, + seed=7, + plants={(1, 1): NODATA, (2, 2): NODATA, (1, 2): 66, (2, 3): 66, (3, 1): 66}, ) - # Engines reduce in different orders, so exact float equality is not - # attainable; 1e-9 passes summation noise and still fails any semantic - # mismatch (selection, nodata handling, ddof). - assert got == pytest.approx(expected, rel=1e-9), (wkt, all_touched, stat) - - -@pytest.mark.parametrize("stat", ["count", "sum"]) -def test_rs_zonalstats_excludes_nodata(subject, comparator, tmp_path, stat): - """A pixel valued at the band nodata inside the zone is excluded from - the reduction by every engine.""" - tiff = tmp_path / "zonal_nodata.tif" - write_geotiff( - tiff, - random_raster_data( - "uint8", bands=1, height=HEIGHT, width=WIDTH, plants={(2, 3): 200} - ), - gdal_transform=GDAL_TRANSFORM, - nodata=200.0, + path = tmp_path / "zonal.tif" + write_geotiff(path, data, gdal_transform=GDAL_TRANSFORM, nodata=NODATA) + return path, data[0] + + +def numpy_reference(band, wkt, *, all_touched, exclude_nodata): + """Reference statistics over the pixels the zone selects, via rasterio+numpy. + + Returns a dict of every statistic, or the sentinel string ``"empty"`` when + the selection is empty (the caller maps that to count 0 / NULLs). + """ + import rasterio.features + import shapely + from rasterio.transform import Affine + + geom = shapely.from_wkt(wkt) + mask = rasterio.features.rasterize( + [(geom, 1)], + out_shape=band.shape, + transform=Affine.from_gdal(*GDAL_TRANSFORM), + all_touched=all_touched, + fill=0, + dtype="uint8", ) + sel = band[mask == 1].astype(np.float64) + if exclude_nodata: + sel = sel[sel != NODATA] + if sel.size == 0: + return "empty" + + values, counts = np.unique(sel, return_counts=True) + mode = float(values[counts == counts.max()].max()) # ties -> largest + n = sel.size + return { + "count": float(n), + "sum": float(sel.sum()), + "mean": float(sel.mean()), + "median": float(np.median(sel)), + "mode": mode, + "stddev": float(sel.std(ddof=1)) if n > 1 else 0.0, + "variance": float(sel.var(ddof=1)) if n > 1 else 0.0, + "min": float(sel.min()), + "max": float(sel.max()), + } - got = subject.zonal_stats(tiff, GEOM_RECT, stat=stat) - expected = comparator.zonal_stats(tiff, GEOM_RECT, stat=stat) - assert got == pytest.approx(expected, rel=1e-9), stat +def zonal_stat(con, path, wkt, stat, options=None): + """RS_ZonalStats over a one-row table (values as columns, not literals).""" + columns = { + "path": pa.array([str(path)], pa.utf8()), + "wkt": pa.array([wkt], pa.utf8()), + "stat": pa.array([stat], pa.utf8()), + } + if options is not None: + columns["options"] = pa.array([options], pa.utf8()) + df = con.create_data_frame(pa.table(columns)) + raster = df.path.funcs.rs_frompath() + geom = con.funcs.st_geomfromtext(df.wkt) + args = [geom, df.stat] + ([df.options] if options is not None else []) + table = df.select(r=raster.funcs.rs_zonalstats(*args)).to_arrow_table() + return table["r"][0].as_py() -def _invoke_zonalstats(con, tiff, wkt, *, all_stats, options=None): - """Invoke RS_ZonalStats ('sum') or RS_ZonalStatsAll over a one-row table - and return the Arrow result. - Arguments travel as table columns so the kernel runs its real array path - (literals constant-fold). The argument surface — `(raster, zone, stat, - options)` for RS_ZonalStats and `(raster, zone, options)` for - RS_ZonalStatsAll, with `band` inside the JSON `options` — mirrors the - RS_ZonalStats function tests. These invoke the subject (SedonaDB) directly - because the harness `zonal_stats` op exposes only a single scalar statistic, - not the all-stats struct or the band-unspecified path. - """ +def zonal_stats_all(con, path, wkt, options=None): + """RS_ZonalStatsAll over a one-row table; returns the struct as a dict.""" columns = { - "path": pa.array([str(tiff)], pa.utf8()), + "path": pa.array([str(path)], pa.utf8()), "wkt": pa.array([wkt], pa.utf8()), } - if not all_stats: - columns["stat"] = pa.array(["sum"], pa.utf8()) if options is not None: columns["options"] = pa.array([options], pa.utf8()) df = con.create_data_frame(pa.table(columns)) raster = df.path.funcs.rs_frompath() geom = con.funcs.st_geomfromtext(df.wkt) - tail = [df.options] if options is not None else [] - if all_stats: - expr = raster.funcs.rs_zonalstatsall(geom, *tail) - else: - expr = raster.funcs.rs_zonalstats(geom, df.stat, *tail) - return df.select(r=expr).to_arrow_table() + args = [geom] + ([df.options] if options is not None else []) + table = df.select(r=raster.funcs.rs_zonalstatsall(*args)).to_arrow_table() + return table["r"][0].as_py() -@pytest.mark.parametrize( - "all_stats", [False, True], ids=["RS_ZonalStats", "RS_ZonalStatsAll"] -) -def test_multiband_raster_requires_band_option(con, tmp_path, all_stats): - """On a multiband raster with no band chosen, SedonaDB raises rather than - reducing an arbitrary band. - - Sedona Spark defaults to band 1 here (the documented divergence); asserting - the raise pins SedonaDB's stricter contract — an ambiguous multiband - selection is an error, not a silent band-1 pick. - - This is a subject-error case (the parity subject itself raises), so a plain - `pytest.raises` on the subject is the right shape; it does not go through the - comparator/deviation ledger. A ledger-integrated "subject_error" Deviation - kind that also captured the Spark band-1 default declaratively would be a - possible future enhancement. - """ - tiff = tmp_path / "multiband.tif" - write_geotiff( - tiff, - random_raster_data("float64", bands=2, height=HEIGHT, width=WIDTH), - gdal_transform=GDAL_TRANSFORM, +@pytest.mark.parametrize("stat", STATS) +@pytest.mark.parametrize("all_touched", [False, True]) +def test_single_stat_matches_numpy(con, tmp_path, stat, all_touched): + path, band = fixture_raster(tmp_path) + options = json.dumps({"band": 1, "all_touched": all_touched}) + expected = numpy_reference( + band, GEOM_RECT, all_touched=all_touched, exclude_nodata=True ) - # GEOM_RECT intersects the raster, so band resolution — not a - # no-intersection short-circuit — is what fails. - with pytest.raises(Exception, match="option to choose one"): - _invoke_zonalstats(con, tiff, GEOM_RECT, all_stats=all_stats) + assert expected != "empty", "GEOM_RECT should select pixels" + got = zonal_stat(con, path, GEOM_RECT, stat, options) + if stat in EXACT_STATS: + assert got == expected[stat] + else: + assert got == pytest.approx(expected[stat]) -def test_zonalstatsall_count_field_is_int64(con, tmp_path): - """RS_ZonalStatsAll returns `count` as an Int64 pixel count. - Sedona Spark returns a uniform `Double[]` for every statistic, so its count - is a floating-point value; SedonaDB keeps count as an integer. Pinning the - Arrow struct field type guards that contract (the rest of the struct is - Float64). - """ - tiff = tmp_path / "singleband.tif" - write_geotiff( - tiff, - random_raster_data("float64", bands=1, height=HEIGHT, width=WIDTH), - gdal_transform=GDAL_TRANSFORM, +def test_all_struct_matches_numpy(con, tmp_path): + path, band = fixture_raster(tmp_path) + expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=True) + got = zonal_stats_all(con, path, GEOM_RECT, json.dumps({"band": 1})) + + assert got["count"] == expected["count"] + for stat in EXACT_STATS - {"count"}: + assert got[stat] == expected[stat] + for stat in ("mean", "variance", "stddev"): + assert got[stat] == pytest.approx(expected[stat]) + + +def test_exclude_nodata_default_and_disabled(con, tmp_path): + path, band = fixture_raster(tmp_path) + # Default excludes nodata; disabling it keeps those pixels, raising count. + excluded = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=True) + included = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=False) + assert included["count"] > excluded["count"] + + assert ( + zonal_stat(con, path, GEOM_RECT, "count", json.dumps({"band": 1})) + == excluded["count"] + ) + assert ( + zonal_stat( + con, + path, + GEOM_RECT, + "count", + json.dumps({"band": 1, "exclude_nodata": False}), + ) + == included["count"] ) - table = _invoke_zonalstats( - con, tiff, GEOM_RECT, all_stats=True, options='{"band": 1}' + + +def test_sliver_selects_nothing_unless_all_touched(con, tmp_path): + path, _ = fixture_raster(tmp_path) + # The zone overlaps the raster but covers no pixel center: count 0, rest NULL. + assert zonal_stat(con, path, GEOM_SLIVER, "count", json.dumps({"band": 1})) == 0.0 + assert zonal_stat(con, path, GEOM_SLIVER, "sum", json.dumps({"band": 1})) is None + # all_touched picks up the pixels it crosses. + touched = zonal_stat( + con, path, GEOM_SLIVER, "count", json.dumps({"band": 1, "all_touched": True}) + ) + assert touched > 0.0 + + +def test_no_intersection_is_null_when_lenient_and_errors_when_strict(con, tmp_path): + path, _ = fixture_raster(tmp_path) + # Lenient (default): NULL, including count. + assert ( + zonal_stat(con, path, GEOM_DISJOINT, "count", json.dumps({"band": 1})) is None ) - struct_type = table.schema.field("r").type - assert struct_type.field("count").type == pa.int64() + assert zonal_stats_all(con, path, GEOM_DISJOINT, json.dumps({"band": 1})) is None + # Strict: errors. + with pytest.raises(Exception, match="does not intersect"): + zonal_stat( + con, path, GEOM_DISJOINT, "count", json.dumps({"band": 1, "lenient": False}) + ) + + +def test_unknown_statistic_errors(con, tmp_path): + path, _ = fixture_raster(tmp_path) + with pytest.raises(Exception, match="unknown statistic"): + zonal_stat(con, path, GEOM_RECT, "nonsense", json.dumps({"band": 1})) + + +def test_sql_text_smoke(con, tmp_path): + """One raw-SQL invocation per function keeps the parser path covered.""" + path, band = fixture_raster(tmp_path) + expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=True) + + single = con.sql( + "SELECT RS_ZonalStats(RS_FromPath($1), ST_GeomFromText($2), 'sum', '{\"band\": 1}') AS r", + params=(str(path), GEOM_RECT), + ).to_arrow_table() + assert single["r"][0].as_py() == expected["sum"] + + everything = con.sql( + "SELECT RS_ZonalStatsAll(RS_FromPath($1), ST_GeomFromText($2), '{\"band\": 1}') AS r", + params=(str(path), GEOM_RECT), + ).to_arrow_table() + assert everything["r"][0].as_py()["count"] == expected["count"] From 72ef6138a6c11df3d0e55fed5ba9a2b09805279d Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 20 Jul 2026 16:53:19 -0700 Subject: [PATCH 04/16] fix(rust/sedona-raster-gdal): gate RS_ZonalStats on true geometry intersection The no-intersection gate now tests a true geometry intersection between the zone and the raster footprint (its convex hull) rather than a zone-envelope intersected with the raster extent (a bounding-box overlap), matching Sedona Spark's rsIntersects gate. A zone whose bounding box overlaps the raster but whose geometry is disjoint is now a no-intersection case (NULL under lenient, an error under strict) instead of count 0. Reuses the RS_Intersects convex-hull predicate through a shared raster_intersects_geom_wkb helper in sedona-raster-functions. --- .../tests/functions/test_rs_zonalstats.py | 38 +++++++++ .../src/rs_spatial_predicates.rs | 14 ++++ rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 77 ++++++++++++++++--- 3 files changed, 118 insertions(+), 11 deletions(-) diff --git a/python/sedonadb/tests/functions/test_rs_zonalstats.py b/python/sedonadb/tests/functions/test_rs_zonalstats.py index 7fa93b5f05..8b0d2653f9 100644 --- a/python/sedonadb/tests/functions/test_rs_zonalstats.py +++ b/python/sedonadb/tests/functions/test_rs_zonalstats.py @@ -50,6 +50,11 @@ ) # Entirely outside the raster extent. GEOM_DISJOINT = "POLYGON ((900 900, 910 900, 910 890, 900 890, 900 900))" +# Bounding box overlaps the raster, but the geometry itself is disjoint: the +# triangle sits in the far corner of its bounding box, clear of the raster. A +# bounding-box gate would burn no pixels and report count 0; a true-geometry gate +# (matching Sedona Spark's rsIntersects) treats it as a no-intersection case. +GEOM_DISJOINT_BBOX = "POLYGON ((124 490, 124 510, 108 510, 124 490))" # A thin strip crossing the x = 104 pixel boundary but covering no pixel center # (centers sit at odd x): selects nothing unless all_touched. GEOM_SLIVER = "POLYGON ((103.6 499, 104.4 499, 104.4 483, 103.6 483, 103.6 499))" @@ -230,6 +235,39 @@ def test_no_intersection_is_null_when_lenient_and_errors_when_strict(con, tmp_pa ) +def test_bbox_overlapping_but_geometry_disjoint_is_no_intersection(con, tmp_path): + import shapely + from shapely.geometry import box + + path, _ = fixture_raster(tmp_path) + + # Premise: the zone's bounding box overlaps the raster extent, but the + # geometry is disjoint from it (unlike GEOM_DISJOINT, whose bbox misses too). + ox, px, _, oy, _, py = GDAL_TRANSFORM + raster_extent = box(ox, oy + py * HEIGHT, ox + px * WIDTH, oy) + geom = shapely.from_wkt(GEOM_DISJOINT_BBOX) + assert box(*geom.bounds).intersects(raster_extent), "bbox must overlap the raster" + assert geom.disjoint(raster_extent), "geometry must be disjoint from the raster" + + # Lenient (default): NULL, not count 0 — a true no-intersection case. + assert ( + zonal_stat(con, path, GEOM_DISJOINT_BBOX, "count", json.dumps({"band": 1})) + is None + ) + assert ( + zonal_stats_all(con, path, GEOM_DISJOINT_BBOX, json.dumps({"band": 1})) is None + ) + # Strict: errors, exactly like the fully-disjoint zone. + with pytest.raises(Exception, match="does not intersect"): + zonal_stat( + con, + path, + GEOM_DISJOINT_BBOX, + "count", + json.dumps({"band": 1, "lenient": False}), + ) + + def test_unknown_statistic_errors(con, tmp_path): path, _ = fixture_raster(tmp_path) with pytest.raises(Exception, match="unknown statistic"): diff --git a/rust/sedona-raster-functions/src/rs_spatial_predicates.rs b/rust/sedona-raster-functions/src/rs_spatial_predicates.rs index 7c5650f9ce..8c95ca58d5 100644 --- a/rust/sedona-raster-functions/src/rs_spatial_predicates.rs +++ b/rust/sedona-raster-functions/src/rs_spatial_predicates.rs @@ -375,6 +375,20 @@ fn evaluate_predicate(wkb_a: &[u8], wkb_b: &[u8]) -> Re /// = 9 + 4 + 80 = 93 const CONVEXHULL_WKB_SIZE: usize = 93; +/// Test whether a geometry intersects a raster's footprint (its convex-hull +/// polygon), using the same true geometry intersection as RS_Intersects rather +/// than a bounding-box overlap. This is the gate Sedona Spark's zonal statistics +/// use (`RasterPredicates.rsIntersects`): a zone whose bounding box overlaps the +/// raster but whose geometry is disjoint is treated as not intersecting. +/// +/// `geom_wkb` must already be in the raster's CRS; this performs no CRS +/// transformation (the raster footprint is built in the raster's own CRS). +pub fn raster_intersects_geom_wkb(raster: &dyn RasterRef, geom_wkb: &[u8]) -> Result { + let mut raster_wkb = Vec::with_capacity(CONVEXHULL_WKB_SIZE); + write_convexhull_wkb(raster, &mut raster_wkb)?; + evaluate_predicate::(&raster_wkb, geom_wkb) +} + /// Create WKB for a convex hull polygon for the raster fn write_convexhull_wkb(raster: &dyn RasterRef, out: &mut impl std::io::Write) -> Result<()> { let width = raster.metadata().width(); diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs index 2e599c5414..4e620a60a1 100644 --- a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -59,6 +59,7 @@ use sedona_raster::array::RasterRefImpl; use sedona_raster::traits::RasterRef; use sedona_raster_functions::crs_utils::{crs_transform_wkb, 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; @@ -503,10 +504,13 @@ struct PixelWindow { /// Compute the statistics of the pixels a zone geometry selects on one band. /// -/// Returns `Ok(None)` when the zone's envelope does not intersect the raster -/// extent (the caller turns this into NULL when `lenient`, an error otherwise). -/// A zone that intersects the extent but whose selected pixels are all outside -/// the geometry or all nodata yields `Ok(Some(..))` with `count = 0`. +/// Returns `Ok(None)` when the zone geometry does not intersect the raster's +/// footprint. This is a true geometry intersection (matching Sedona Spark's +/// `rsIntersects` gate), not a bounding-box overlap: a zone whose envelope +/// overlaps the raster but whose geometry is disjoint is a no-intersection case. +/// The caller turns `None` into NULL when `lenient`, an error otherwise. A zone +/// that intersects the footprint but whose selected pixels are all outside the +/// geometry or all nodata yields `Ok(Some(..))` with `count = 0`. /// /// `scratch` is a reused buffer for the selected pixel values so the per-row /// collection does not allocate a fresh `Vec` each call. @@ -539,13 +543,25 @@ fn compute_zonal_stats( let height = usize::try_from(metadata.height()) .map_err(|_| exec_datafusion_err!("RS_ZonalStats: negative raster height"))?; - // Parse the zone and clamp its envelope to the raster grid. A disjoint - // envelope is the no-intersection case. + // No-intersection gate: a true geometry intersection between the zone and + // the raster footprint (matching Sedona Spark's rsIntersects gate), not a + // bounding-box overlap. A zone whose envelope overlaps the raster but whose + // geometry is disjoint is a no-intersection case, not a count-0 case. The + // zone is already in the raster's CRS here, so no transform is needed. + if !raster_intersects_geom_wkb(raster, geom_wkb)? { + return Ok(None); + } + + // Parse the zone and clamp its envelope to the raster grid for the pixel + // window to rasterize. The gate above already established overlap; a + // degenerate window (the zone 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 { - return Ok(None); + scratch.clear(); + return Ok(Some(compute_statistics(scratch))); }; // Rasterize the zone into a window-sized 0/1 mask (moves `geometry`, whose @@ -623,10 +639,10 @@ fn raster_transform(raster: &RasterRefImpl<'_>) -> Result { } /// The zone's envelope intersected with the raster extent, snapped outward to -/// whole pixels. `None` when the envelope is disjoint from the raster. Mirrors -/// the window RS_Clip / PostGIS ST_Clip use: all four corners are mapped -/// through the inverse geotransform so a skewed raster still gets a correct -/// superset window. +/// whole pixels. `None` when the clamped window has no area (the envelope only +/// touches the raster boundary). Mirrors the window RS_Clip / PostGIS ST_Clip +/// use: all four corners are mapped through the inverse geotransform so a skewed +/// raster still gets a correct superset window. fn envelope_window( geometry: &Geometry, transform: &GeoTransform, @@ -1266,6 +1282,45 @@ mod udf_tests { ); } + #[test] + fn bbox_overlapping_but_geometry_disjoint_zone_is_no_intersection() { + // small_raster covers x ∈ [0, 4], y ∈ [0, 2]. This triangle lives on the + // far side of the line x + y = 7, so its geometry is disjoint from the + // raster (every raster point has x + y ≤ 6), yet its bounding box + // [0, 7] × [0, 7] contains the whole raster. A bounding-box gate would + // burn zero pixels and report count 0; the true-geometry gate (matching + // Sedona Spark's rsIntersects) treats it as a no-intersection case. + let disjoint = "POLYGON((7 0, 0 7, 7 7, 7 0))"; + + // Lenient (default): NULL, not count 0. + assert_eq!( + cv_f64(call_stats(&small_raster(), disjoint, "count", None).unwrap()), + None + ); + assert!(cv_struct(call_all(&small_raster(), disjoint, None).unwrap()).is_null(0)); + + // Strict: both functions error. + let err = call_stats( + &small_raster(), + disjoint, + "count", + Some(r#"{"lenient": false}"#), + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("does not intersect"), + "unexpected error: {err}" + ); + let err = call_all(&small_raster(), disjoint, Some(r#"{"lenient": false}"#)) + .unwrap_err() + .to_string(); + assert!( + err.contains("does not intersect"), + "unexpected error: {err}" + ); + } + #[test] fn nodata_pixels_are_excluded_by_default_and_kept_when_asked() { // A 2×2 UInt8 raster [10, 255, 20, 30] with nodata 255, world extent From 1e0fbab77d8c852ef33e8637e268201fd06da80e Mon Sep 17 00:00:00 2001 From: jameswillis Date: Tue, 21 Jul 2026 13:15:35 -0700 Subject: [PATCH 05/16] refactor(rust/sedona-raster-gdal): RS_ZonalStats(All) positional overloads Replace the trailing JSON-options argument with Apache Sedona Spark's positional overload ladder (band, stat_type, all_touched, exclude_nodata, lenient) so Spark SQL tends to run unchanged. The band-less overloads still error on a multiband raster rather than defaulting to band 1. Drop the now-unused serde/serde_json dependencies. --- Cargo.lock | 2 - docs/reference/sql/rs_zonalstats.qmd | 70 +- docs/reference/sql/rs_zonalstatsall.qmd | 64 +- .../tests/functions/test_rs_zonalstats.py | 160 ++-- rust/sedona-raster-gdal/Cargo.toml | 2 - .../benches/rs_zonalstats.rs | 29 +- rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 732 ++++++++++++------ 7 files changed, 683 insertions(+), 376 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a2418e97ab..4818fa196a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6259,8 +6259,6 @@ dependencies = [ "sedona-raster-functions", "sedona-schema", "sedona-testing", - "serde", - "serde_json", "tempfile", "tokio", ] diff --git a/docs/reference/sql/rs_zonalstats.qmd b/docs/reference/sql/rs_zonalstats.qmd index 6929642c60..efafdd3470 100644 --- a/docs/reference/sql/rs_zonalstats.qmd +++ b/docs/reference/sql/rs_zonalstats.qmd @@ -25,7 +25,28 @@ kernels: args: - raster - {name: zone, type: geometry} - - {name: stat, type: string} + - {name: stat_type, type: string} + - returns: double + args: + - raster + - {name: zone, type: geometry} + - {name: band, type: integer} + - {name: stat_type, type: string} + - returns: double + args: + - raster + - {name: zone, type: geometry} + - {name: band, type: integer} + - {name: stat_type, type: string} + - {name: all_touched, type: boolean} + - returns: double + args: + - raster + - {name: zone, type: geometry} + - {name: band, type: integer} + - {name: stat_type, type: string} + - {name: all_touched, type: boolean} + - {name: exclude_nodata, type: boolean} - returns: double args: - raster @@ -35,21 +56,31 @@ kernels: Zone geometry defining the region of interest. Reprojected into the raster's CRS when both carry one; it is an error for exactly one side to have a CRS. - - name: stat + - name: band + type: integer + description: > + 1-based band index to compute over. Required for a multiband raster; a + single-band raster may omit it via the band-less overload. + - name: stat_type type: string description: > Statistic to return (case-insensitive): count, sum, mean, median, mode, stddev, variance, min, or max. `avg`/`average` alias mean and `sd` aliases stddev. - - name: options - type: string + - name: all_touched + type: boolean + description: > + If true, include every pixel the zone touches; otherwise only pixels + whose center falls inside it. Defaults to false. + - name: exclude_nodata + type: boolean + description: > + If true (the default), skip pixels equal to the band's nodata value. + - name: lenient + type: boolean description: > - JSON object of options. `band` is the 1-based band index (required for a - multiband raster; defaults to 1 for a single-band raster). `all_touched` - (default false) keeps every pixel the zone touches rather than only those - whose center it covers. `exclude_nodata` (default true) skips pixels equal - to the band's nodata value. `lenient` (default true) returns NULL when the - zone does not intersect the raster instead of raising an error. + If true (the default), return NULL when the zone does not intersect the + raster; if false, raise an error. --- ::: callout-warning @@ -73,10 +104,13 @@ other statistic is NULL. When the zone does not intersect the raster at all, the result is NULL under the default `lenient` behavior, or an error when `lenient` is set to false. -The options are supplied as a single JSON object rather than as a positional -argument list, so the argument surface stays `(raster, zone, stat, options)`. -This function operates on 2-D `(y, x)` bands; computing a statistic per -non-spatial plane of an N-D band is not supported. +The positional overloads mirror Apache Sedona Spark's `RS_ZonalStats` verbatim. +The `band`, `all_touched`, `exclude_nodata`, and `lenient` arguments are added +one at a time by the wider overloads; `all_touched` defaults to false, +`exclude_nodata` to true, and `lenient` to true. Unlike Sedona Spark, the +band-less overload does not default to band 1 on a multiband raster: naming the +band is required there. This function operates on 2-D `(y, x)` bands; computing a +statistic per non-spatial plane of an N-D band is not supported. Use [`RS_ZonalStatsAll`](rs_zonalstatsall.qmd) to compute every statistic at once. @@ -87,8 +121,8 @@ once. SELECT RS_ZonalStats( RS_Example(), ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), - 'mean', - '{"band": 1}' + 1, + 'mean' ); ``` @@ -96,7 +130,9 @@ SELECT RS_ZonalStats( SELECT RS_ZonalStats( RS_Example(), ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), + 1, 'count', - '{"band": 1, "exclude_nodata": false}' + false, + false ); ``` diff --git a/docs/reference/sql/rs_zonalstatsall.qmd b/docs/reference/sql/rs_zonalstatsall.qmd index 7b2bc69743..c252f324a2 100644 --- a/docs/reference/sql/rs_zonalstatsall.qmd +++ b/docs/reference/sql/rs_zonalstatsall.qmd @@ -25,6 +25,24 @@ kernels: args: - raster - {name: zone, type: geometry} + - returns: struct + args: + - raster + - {name: zone, type: geometry} + - {name: band, type: integer} + - returns: struct + args: + - raster + - {name: zone, type: geometry} + - {name: band, type: integer} + - {name: all_touched, type: boolean} + - returns: struct + args: + - raster + - {name: zone, type: geometry} + - {name: band, type: integer} + - {name: all_touched, type: boolean} + - {name: exclude_nodata, type: boolean} - returns: struct args: - raster @@ -34,15 +52,25 @@ kernels: Zone geometry defining the region of interest. Reprojected into the raster's CRS when both carry one; it is an error for exactly one side to have a CRS. - - name: options - type: string + - name: band + type: integer + description: > + 1-based band index to compute over. Required for a multiband raster; a + single-band raster may omit it via the band-less overload. + - name: all_touched + type: boolean + description: > + If true, include every pixel the zone touches; otherwise only pixels + whose center falls inside it. Defaults to false. + - name: exclude_nodata + type: boolean + description: > + If true (the default), skip pixels equal to the band's nodata value. + - name: lenient + type: boolean description: > - JSON object of options. `band` is the 1-based band index (required for a - multiband raster; defaults to 1 for a single-band raster). `all_touched` - (default false) keeps every pixel the zone touches rather than only those - whose center it covers. `exclude_nodata` (default true) skips pixels equal - to the band's nodata value. `lenient` (default true) returns NULL when the - zone does not intersect the raster instead of raising an error. + If true (the default), return NULL when the zone does not intersect the + raster; if false, raise an error. --- ::: callout-warning @@ -58,19 +86,23 @@ order of Apache Sedona's `RS_ZonalStatsAll`). A pixel is included when its center falls inside the zone (or, with `all_touched`, when the zone touches it at all), and the band's nodata pixels are excluded by default. -`count` is a whole number; every other field is floating point. Variance and -standard deviation are the sample (n-1) values, and `mode` breaks ties toward -the larger value. +`count` is a whole number (a 64-bit integer); every other field is floating +point. Variance and standard deviation are the sample (n-1) values, and `mode` +breaks ties toward the larger value. When the zone overlaps the raster but selects no pixel, `count` is 0 and every other field is NULL. When the zone does not intersect the raster at all, the whole struct is NULL under the default `lenient` behavior, or the call raises an error when `lenient` is set to false. -The options are supplied as a single JSON object rather than as a positional -argument list, so the argument surface stays `(raster, zone, options)`. This -function operates on 2-D `(y, x)` bands; computing statistics per non-spatial -plane of an N-D band is not supported. +The positional overloads mirror Apache Sedona Spark's `RS_ZonalStatsAll` +verbatim (the same ladder as `RS_ZonalStats` without `stat_type`). The `band`, +`all_touched`, `exclude_nodata`, and `lenient` arguments are added one at a time +by the wider overloads; `all_touched` defaults to false, `exclude_nodata` to +true, and `lenient` to true. Unlike Sedona Spark, the band-less overload does +not default to band 1 on a multiband raster: naming the band is required there. +This function operates on 2-D `(y, x)` bands; computing statistics per +non-spatial plane of an N-D band is not supported. Use [`RS_ZonalStats`](rs_zonalstats.qmd) to compute a single statistic. @@ -80,6 +112,6 @@ Use [`RS_ZonalStats`](rs_zonalstats.qmd) to compute a single statistic. SELECT RS_ZonalStatsAll( RS_Example(), ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), - '{"band": 1}' + 1 ); ``` diff --git a/python/sedonadb/tests/functions/test_rs_zonalstats.py b/python/sedonadb/tests/functions/test_rs_zonalstats.py index 8b0d2653f9..0ad888fb86 100644 --- a/python/sedonadb/tests/functions/test_rs_zonalstats.py +++ b/python/sedonadb/tests/functions/test_rs_zonalstats.py @@ -17,6 +17,11 @@ """RS_ZonalStats / RS_ZonalStatsAll cross-checked against a numpy reference. +Both functions mirror Apache Sedona Spark's positional overloads, so the tests +call them positionally: `(raster, zone, stat_type)` / +`(raster, zone, band, stat_type[, all_touched[, exclude_nodata[, lenient]]])` +for RS_ZonalStats and the same ladder without `stat_type` for RS_ZonalStatsAll. + The fixture raster is CRS-less (so nothing reprojects and pixel selection is bit-comparable). The reference rasterizes the zone with `rasterio.features` (the same GDAL rasterizer the kernel uses) and reduces the selected pixels with @@ -28,15 +33,16 @@ when it is unavailable rather than importing it at module scope. """ -import json - import numpy as np import pyarrow as pa import pytest pytest.importorskip("rasterio") -from sedonadb.raster_testing import random_raster_data, write_geotiff # noqa: E402 +from sedonadb.raster_testing import ( # noqa: E402 + random_raster_data, + write_geotiff, +) # GDAL-order geotransform: origin (100, 500), 2-wide by 3-tall north-up pixels; # a 6x7 raster then spans x in [100, 114], y in [482, 500]. @@ -124,36 +130,41 @@ def numpy_reference(band, wkt, *, all_touched, exclude_nodata): } -def zonal_stat(con, path, wkt, stat, options=None): - """RS_ZonalStats over a one-row table (values as columns, not literals).""" - columns = { - "path": pa.array([str(path)], pa.utf8()), - "wkt": pa.array([wkt], pa.utf8()), - "stat": pa.array([stat], pa.utf8()), - } - if options is not None: - columns["options"] = pa.array([options], pa.utf8()) - df = con.create_data_frame(pa.table(columns)) - raster = df.path.funcs.rs_frompath() - geom = con.funcs.st_geomfromtext(df.wkt) - args = [geom, df.stat] + ([df.options] if options is not None else []) - table = df.select(r=raster.funcs.rs_zonalstats(*args)).to_arrow_table() +def _one_row(con, path, wkt): + """A one-row frame with the raster and zone as columns, so the kernel runs + its real per-row array path rather than constant-folding a scalar.""" + df = con.create_data_frame( + pa.table( + { + "path": pa.array([str(path)], pa.utf8()), + "wkt": pa.array([wkt], pa.utf8()), + } + ) + ) + return df, df.path.funcs.rs_frompath(), con.funcs.st_geomfromtext(df.wkt) + + +def zonal_stat(con, path, wkt, trailing): + """RS_ZonalStats over a one-row table. + + `trailing` is the positional argument list after `(raster, zone)` — e.g. + `["mean"]` (band-less overload) or `[1, "mean", all_touched, exclude_nodata, + lenient]`. Raster and zone travel as columns; the trailing scalars are + literals, matching how the SQL reads. + """ + df, raster, geom = _one_row(con, path, wkt) + table = df.select(r=raster.funcs.rs_zonalstats(geom, *trailing)).to_arrow_table() return table["r"][0].as_py() -def zonal_stats_all(con, path, wkt, options=None): - """RS_ZonalStatsAll over a one-row table; returns the struct as a dict.""" - columns = { - "path": pa.array([str(path)], pa.utf8()), - "wkt": pa.array([wkt], pa.utf8()), - } - if options is not None: - columns["options"] = pa.array([options], pa.utf8()) - df = con.create_data_frame(pa.table(columns)) - raster = df.path.funcs.rs_frompath() - geom = con.funcs.st_geomfromtext(df.wkt) - args = [geom] + ([df.options] if options is not None else []) - table = df.select(r=raster.funcs.rs_zonalstatsall(*args)).to_arrow_table() +def zonal_stats_all(con, path, wkt, trailing): + """RS_ZonalStatsAll over a one-row table; returns the struct as a dict. + + `trailing` is the positional argument list after `(raster, zone)` — e.g. + `[]` (band-less overload) or `[1, all_touched, exclude_nodata, lenient]`. + """ + df, raster, geom = _one_row(con, path, wkt) + table = df.select(r=raster.funcs.rs_zonalstatsall(geom, *trailing)).to_arrow_table() return table["r"][0].as_py() @@ -161,13 +172,13 @@ def zonal_stats_all(con, path, wkt, options=None): @pytest.mark.parametrize("all_touched", [False, True]) def test_single_stat_matches_numpy(con, tmp_path, stat, all_touched): path, band = fixture_raster(tmp_path) - options = json.dumps({"band": 1, "all_touched": all_touched}) expected = numpy_reference( band, GEOM_RECT, all_touched=all_touched, exclude_nodata=True ) assert expected != "empty", "GEOM_RECT should select pixels" - got = zonal_stat(con, path, GEOM_RECT, stat, options) + # (raster, zone, band, stat_type, all_touched) — the 5-arg overload. + got = zonal_stat(con, path, GEOM_RECT, [1, stat, all_touched]) if stat in EXACT_STATS: assert got == expected[stat] else: @@ -177,8 +188,12 @@ def test_single_stat_matches_numpy(con, tmp_path, stat, all_touched): def test_all_struct_matches_numpy(con, tmp_path): path, band = fixture_raster(tmp_path) expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=True) - got = zonal_stats_all(con, path, GEOM_RECT, json.dumps({"band": 1})) + # (raster, zone, band) — all_touched defaults to false. + got = zonal_stats_all(con, path, GEOM_RECT, [1]) + # count is an integer field (Int64); every other field is floating point. + assert isinstance(got["count"], int) + assert isinstance(got["sum"], float) assert got["count"] == expected["count"] for stat in EXACT_STATS - {"count"}: assert got[stat] == expected[stat] @@ -193,18 +208,12 @@ def test_exclude_nodata_default_and_disabled(con, tmp_path): included = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=False) assert included["count"] > excluded["count"] + # Default (4-arg (raster, zone, band, stat_type)) excludes nodata. + assert zonal_stat(con, path, GEOM_RECT, [1, "count"]) == excluded["count"] + # exclude_nodata => false keeps it: the 6-arg overload trails (all_touched, + # exclude_nodata). assert ( - zonal_stat(con, path, GEOM_RECT, "count", json.dumps({"band": 1})) - == excluded["count"] - ) - assert ( - zonal_stat( - con, - path, - GEOM_RECT, - "count", - json.dumps({"band": 1, "exclude_nodata": False}), - ) + zonal_stat(con, path, GEOM_RECT, [1, "count", False, False]) == included["count"] ) @@ -212,27 +221,22 @@ def test_exclude_nodata_default_and_disabled(con, tmp_path): def test_sliver_selects_nothing_unless_all_touched(con, tmp_path): path, _ = fixture_raster(tmp_path) # The zone overlaps the raster but covers no pixel center: count 0, rest NULL. - assert zonal_stat(con, path, GEOM_SLIVER, "count", json.dumps({"band": 1})) == 0.0 - assert zonal_stat(con, path, GEOM_SLIVER, "sum", json.dumps({"band": 1})) is None - # all_touched picks up the pixels it crosses. - touched = zonal_stat( - con, path, GEOM_SLIVER, "count", json.dumps({"band": 1, "all_touched": True}) - ) + assert zonal_stat(con, path, GEOM_SLIVER, [1, "count"]) == 0.0 + assert zonal_stat(con, path, GEOM_SLIVER, [1, "sum"]) is None + # all_touched (5-arg overload) picks up the pixels it crosses. + touched = zonal_stat(con, path, GEOM_SLIVER, [1, "count", True]) assert touched > 0.0 def test_no_intersection_is_null_when_lenient_and_errors_when_strict(con, tmp_path): path, _ = fixture_raster(tmp_path) # Lenient (default): NULL, including count. - assert ( - zonal_stat(con, path, GEOM_DISJOINT, "count", json.dumps({"band": 1})) is None - ) - assert zonal_stats_all(con, path, GEOM_DISJOINT, json.dumps({"band": 1})) is None - # Strict: errors. + assert zonal_stat(con, path, GEOM_DISJOINT, [1, "count"]) is None + assert zonal_stats_all(con, path, GEOM_DISJOINT, [1]) is None + # Strict (lenient => false): the 7-arg overload trails (all_touched, + # exclude_nodata, lenient). with pytest.raises(Exception, match="does not intersect"): - zonal_stat( - con, path, GEOM_DISJOINT, "count", json.dumps({"band": 1, "lenient": False}) - ) + zonal_stat(con, path, GEOM_DISJOINT, [1, "count", False, True, False]) def test_bbox_overlapping_but_geometry_disjoint_is_no_intersection(con, tmp_path): @@ -250,28 +254,34 @@ def test_bbox_overlapping_but_geometry_disjoint_is_no_intersection(con, tmp_path assert geom.disjoint(raster_extent), "geometry must be disjoint from the raster" # Lenient (default): NULL, not count 0 — a true no-intersection case. - assert ( - zonal_stat(con, path, GEOM_DISJOINT_BBOX, "count", json.dumps({"band": 1})) - is None - ) - assert ( - zonal_stats_all(con, path, GEOM_DISJOINT_BBOX, json.dumps({"band": 1})) is None - ) + assert zonal_stat(con, path, GEOM_DISJOINT_BBOX, [1, "count"]) is None + assert zonal_stats_all(con, path, GEOM_DISJOINT_BBOX, [1]) is None # Strict: errors, exactly like the fully-disjoint zone. with pytest.raises(Exception, match="does not intersect"): - zonal_stat( - con, - path, - GEOM_DISJOINT_BBOX, - "count", - json.dumps({"band": 1, "lenient": False}), - ) + zonal_stat(con, path, GEOM_DISJOINT_BBOX, [1, "count", False, True, False]) def test_unknown_statistic_errors(con, tmp_path): path, _ = fixture_raster(tmp_path) with pytest.raises(Exception, match="unknown statistic"): - zonal_stat(con, path, GEOM_RECT, "nonsense", json.dumps({"band": 1})) + zonal_stat(con, path, GEOM_RECT, [1, "nonsense"]) + + +def test_implicit_band_on_multiband_raster_errors(con, tmp_path): + # A 2-band raster: the band-less overloads must error rather than default to + # band 1 (this deliberately diverges from Sedona Spark). + data = random_raster_data("int32", bands=2, height=HEIGHT, width=WIDTH, seed=3) + path = tmp_path / "multiband.tif" + write_geotiff(path, data, gdal_transform=GDAL_TRANSFORM) + + # RS_ZonalStats 3-arg (raster, zone, stat_type): band implicit. + with pytest.raises(Exception, match="2 bands"): + zonal_stat(con, path, GEOM_RECT, ["count"]) + # RS_ZonalStatsAll 2-arg (raster, zone): band implicit. + with pytest.raises(Exception, match="2 bands"): + zonal_stats_all(con, path, GEOM_RECT, []) + # Naming the band resolves the ambiguity. + assert zonal_stat(con, path, GEOM_RECT, [1, "count"]) > 0.0 def test_sql_text_smoke(con, tmp_path): @@ -280,13 +290,13 @@ def test_sql_text_smoke(con, tmp_path): expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=True) single = con.sql( - "SELECT RS_ZonalStats(RS_FromPath($1), ST_GeomFromText($2), 'sum', '{\"band\": 1}') AS r", + "SELECT RS_ZonalStats(RS_FromPath($1), ST_GeomFromText($2), 1, 'sum') AS r", params=(str(path), GEOM_RECT), ).to_arrow_table() assert single["r"][0].as_py() == expected["sum"] everything = con.sql( - "SELECT RS_ZonalStatsAll(RS_FromPath($1), ST_GeomFromText($2), '{\"band\": 1}') AS r", + "SELECT RS_ZonalStatsAll(RS_FromPath($1), ST_GeomFromText($2), 1) AS r", params=(str(path), GEOM_RECT), ).to_arrow_table() assert everything["r"][0].as_py()["count"] == expected["count"] diff --git a/rust/sedona-raster-gdal/Cargo.toml b/rust/sedona-raster-gdal/Cargo.toml index 2c3dc6eca6..6aa9f034c5 100644 --- a/rust/sedona-raster-gdal/Cargo.toml +++ b/rust/sedona-raster-gdal/Cargo.toml @@ -46,8 +46,6 @@ sedona-proj = { workspace = true } sedona-raster = { workspace = true } sedona-raster-functions = { workspace = true } sedona-schema = { workspace = true } -serde = { workspace = true } -serde_json = { workspace = true } tokio = { workspace = true } [features] diff --git a/rust/sedona-raster-gdal/benches/rs_zonalstats.rs b/rust/sedona-raster-gdal/benches/rs_zonalstats.rs index ec46a69ae5..7f5122682a 100644 --- a/rust/sedona-raster-gdal/benches/rs_zonalstats.rs +++ b/rust/sedona-raster-gdal/benches/rs_zonalstats.rs @@ -23,7 +23,7 @@ //! Each case builds a raster whose world extent is exactly the zone-polygon //! generator's `[-10, 10]²` bounds at the requested resolution, so every //! generated polygon lands on the raster and the full mask/collect/reduce path -//! runs. `all_touched = true` (passed via the JSON options) guarantees a +//! runs. `all_touched = true` (the trailing boolean argument) guarantees a //! polygon smaller than a cell still burns at least one pixel rather than //! hitting the empty-zone early return. //! @@ -42,7 +42,7 @@ use std::sync::Arc; -use arrow_array::{ArrayRef, BinaryArray, StringArray}; +use arrow_array::{ArrayRef, BinaryArray, BooleanArray, Int64Array, StringArray}; use criterion::{criterion_group, criterion_main, Criterion}; use datafusion_expr::ScalarUDF; use sedona_schema::datatypes::{SedonaType, RASTER, WKB_GEOMETRY}; @@ -64,15 +64,16 @@ fn criterion_benchmark(c: &mut Criterion) { .clone() .into(); - // RS_ZonalStats(raster, zone, stat, options) and - // RS_ZonalStatsAll(raster, zone, options). + // RS_ZonalStats(raster, zone, band, stat, all_touched) and + // RS_ZonalStatsAll(raster, zone, band, all_touched). let stats_tester = ScalarUdfTester::new( stats_udf, vec![ RASTER, WKB_GEOMETRY, + SedonaType::Arrow(arrow_schema::DataType::Int64), SedonaType::Arrow(arrow_schema::DataType::Utf8), - SedonaType::Arrow(arrow_schema::DataType::Utf8), + SedonaType::Arrow(arrow_schema::DataType::Boolean), ], ); let stats_all_tester = ScalarUdfTester::new( @@ -80,14 +81,14 @@ fn criterion_benchmark(c: &mut Criterion) { vec![ RASTER, WKB_GEOMETRY, - SedonaType::Arrow(arrow_schema::DataType::Utf8), + SedonaType::Arrow(arrow_schema::DataType::Int64), + SedonaType::Arrow(arrow_schema::DataType::Boolean), ], ); + let band: ArrayRef = Arc::new(Int64Array::from(vec![1])); let mean_stat: ArrayRef = Arc::new(StringArray::from(vec!["mean"])); - let all_touched_opts: ArrayRef = Arc::new(StringArray::from(vec![ - r#"{"band": 1, "all_touched": true}"#, - ])); + 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. @@ -117,8 +118,9 @@ fn criterion_benchmark(c: &mut Criterion) { .invoke_arrays(vec![ raster.clone(), geom.clone(), + band.clone(), mean_stat.clone(), - all_touched_opts.clone(), + all_touched.clone(), ]) .unwrap() }) @@ -129,7 +131,12 @@ fn criterion_benchmark(c: &mut Criterion) { c.bench_function(label, |b| { b.iter(|| { stats_all_tester - .invoke_arrays(vec![raster.clone(), geom.clone(), all_touched_opts.clone()]) + .invoke_arrays(vec![ + raster.clone(), + geom.clone(), + band.clone(), + all_touched.clone(), + ]) .unwrap() }) }); diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs index 4e620a60a1..28c72f7a08 100644 --- a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -18,19 +18,31 @@ //! RS_ZonalStats / RS_ZonalStatsAll UDFs — summary statistics of the raster //! pixels covered by a zone geometry. //! -//! `RS_ZonalStats(raster, zone, stat[, options])` returns one statistic as a -//! `Float64`; `RS_ZonalStatsAll(raster, zone[, options])` returns every -//! statistic as a struct. Both compute over the pixels of a single band whose -//! centre falls inside the zone (or that the zone merely touches, with -//! `all_touched`), optionally excluding the band's nodata value. +//! 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`: //! -//! The optional trailing `options` argument is a JSON object rather than the -//! positional `band, all_touched, exclude_nodata, lenient` overload ladder of -//! Sedona Spark, so the surface stays a single `(input, options)` shape: +//! - `RS_ZonalStats(raster, zone, stat)` +//! - `RS_ZonalStats(raster, zone, band, stat)` +//! - `RS_ZonalStats(raster, zone, band, stat, all_touched)` +//! - `RS_ZonalStats(raster, zone, band, stat, all_touched, exclude_nodata)` +//! - `RS_ZonalStats(raster, zone, band, stat, all_touched, exclude_nodata, lenient)` //! -//! ```json -//! {"band": 1, "all_touched": false, "exclude_nodata": true, "lenient": true} -//! ``` +//! `RS_ZonalStatsAll` returns every statistic as a struct, with the same ladder +//! minus `stat`: +//! +//! - `RS_ZonalStatsAll(raster, zone)` +//! - `RS_ZonalStatsAll(raster, zone, band)` +//! - `RS_ZonalStatsAll(raster, zone, band, all_touched)` +//! - `RS_ZonalStatsAll(raster, zone, band, all_touched, exclude_nodata)` +//! - `RS_ZonalStatsAll(raster, zone, band, all_touched, exclude_nodata, lenient)` +//! +//! A pixel is included when its centre falls inside the zone (or that the zone +//! merely touches, with `all_touched`), optionally excluding the band's nodata +//! value. `all_touched` defaults to false, `exclude_nodata` 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 @@ -40,13 +52,13 @@ use std::collections::HashMap; use std::sync::Arc; use arrow_array::builder::{Float64Builder, Int64Builder, StructBuilder}; -use arrow_array::{ArrayRef, StringArray}; +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}; +use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue}; use datafusion_expr::{ColumnarValue, Volatility}; -use serde::Deserialize; use sedona_common::sedona_internal_err; use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; @@ -102,14 +114,20 @@ impl StatType { } } -/// Optional `(input, options)` JSON payload. Missing fields take their default; -/// unknown fields are rejected so a typo surfaces rather than being ignored. -#[derive(Debug, Clone, Deserialize)] -#[serde(default, deny_unknown_fields, rename_all = "snake_case")] -struct ZonalStatsOptions { - /// 1-based band to compute over. `None` means "resolve the default": band 1 - /// for a single-band raster, an error for a multiband raster (naming the - /// band is required rather than silently getting band 1). +/// 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, /// Include every pixel the zone touches, not only those whose centre it /// covers. @@ -122,29 +140,6 @@ struct ZonalStatsOptions { lenient: bool, } -impl Default for ZonalStatsOptions { - fn default() -> Self { - Self { - band: None, - all_touched: false, - exclude_nodata: true, - lenient: true, - } - } -} - -impl ZonalStatsOptions { - /// Parse the JSON options string, or return the defaults when no options - /// argument was supplied (`None`) or it is SQL NULL. - fn parse(json: Option<&str>) -> Result { - match json { - None => Ok(Self::default()), - Some(s) => serde_json::from_str(s) - .map_err(|e| exec_datafusion_err!("RS_ZonalStats: invalid options JSON: {e}")), - } - } -} - /// Every statistic for a zone. `count` is always present (0 when the zone /// selects no pixels); the remaining fields are `None` in exactly that /// no-pixel case and `Some` otherwise, mirroring Sedona Spark (which returns @@ -184,17 +179,18 @@ impl ZonalStatistics { // RS_ZonalStats // ============================================================================= -/// `RS_ZonalStats(raster, zone, stat[, options])` — one statistic as a -/// `Float64`. `stat` is a statistic name (`count`, `sum`, `mean`, `median`, -/// `mode`, `stddev`, `variance`, `min`, `max`). +/// `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 { - with_options: false, - }), - Arc::new(RsZonalStats { with_options: true }), + Arc::new(RsZonalStats { arg_count: 3 }), // (raster, zone, stat) + Arc::new(RsZonalStats { arg_count: 4 }), // (raster, zone, band, stat) + Arc::new(RsZonalStats { arg_count: 5 }), // + all_touched + Arc::new(RsZonalStats { arg_count: 6 }), // + exclude_nodata + Arc::new(RsZonalStats { arg_count: 7 }), // + lenient ], Volatility::Immutable, ) @@ -205,20 +201,59 @@ pub fn rs_zonal_stats_udf() -> SedonaScalarUDF { #[derive(Debug)] struct RsZonalStats { - /// Whether this kernel matches the trailing JSON `options` argument. - with_options: bool, + /// Number of arguments in the matched signature (3..=7). + arg_count: usize, } impl SedonaScalarKernel for RsZonalStats { fn return_type(&self, args: &[SedonaType]) -> Result> { - let mut matchers = vec![ - ArgMatcher::is_raster(), - ArgMatcher::is_geometry_or_geography(), - ArgMatcher::is_string(), - ]; - if self.with_options { - matchers.push(ArgMatcher::is_string()); - } + // Argument order mirrors Sedona Spark: (raster, zone, [band,] stat, + // [all_touched, [exclude_nodata, [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) } @@ -241,20 +276,50 @@ impl SedonaScalarKernel for RsZonalStats { ) -> Result { let num_iterations = RasterExecutor::num_iterations_over(args); - // stat is at index 2, options (when present) at index 3. - let stat_array = expand_string_arg(&args[2], num_iterations)?; + // 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 options_array = self - .with_options - .then(|| expand_string_arg(&args[3], num_iterations)) + + let band_array = has_band + .then(|| expand_int64_arg(&args[2], num_iterations)) .transpose()?; - let mut options_iter = options_array.as_ref().map(|a| a.iter()); + let mut band_iter = band_array.as_ref().map(|a| a.iter()); + + // all_touched (index 4), exclude_nodata (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_nodata_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_nodata_iter = exclude_nodata_array.iter(); + let mut lenient_iter = lenient_array.iter(); let mut builder = Float64Builder::with_capacity(num_iterations); let mut scratch: Vec = Vec::new(); - // The executor only sees (raster, zone); the stat/options columns are - // advanced in lockstep below. + // The executor only sees (raster, zone); 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 = @@ -265,8 +330,15 @@ impl SedonaScalarKernel for RsZonalStats { 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 options_str = options_iter.as_mut().and_then(|i| i.next().flatten()); - let options = ZonalStatsOptions::parse(options_str)?; + let Some(params) = next_params( + &mut band_iter, + &mut all_touched_iter, + &mut exclude_nodata_iter, + &mut lenient_iter, + ) else { + builder.append_null(); + return Ok(()); + }; // A NULL stat, raster, or zone propagates to a NULL row. let (Some(stat_str), Some(raster), Some(wkb)) = (stat_str, raster_opt, wkb_opt) @@ -293,14 +365,14 @@ impl SedonaScalarKernel for RsZonalStats { "Cannot operate on geometry and raster: geometry has no CRS but raster does" ), }; - match compute_zonal_stats(gdal, raster, &geom_wkb, &options, &mut scratch)? { + 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 zone does not intersect the raster: NULL when // lenient (the default), an error otherwise. - None if options.lenient => builder.append_null(), + None if params.lenient => builder.append_null(), None => return no_intersection_err(), } Ok(()) @@ -317,16 +389,18 @@ impl SedonaScalarKernel for RsZonalStats { // RS_ZonalStatsAll // ============================================================================= -/// `RS_ZonalStatsAll(raster, zone[, options])` — every statistic as a struct -/// with fields `count, sum, mean, median, mode, stddev, variance, min, max`. +/// `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 { - with_options: false, - }), - Arc::new(RsZonalStatsAll { with_options: true }), + Arc::new(RsZonalStatsAll { arg_count: 2 }), // (raster, zone) + Arc::new(RsZonalStatsAll { arg_count: 3 }), // (raster, zone, band) + Arc::new(RsZonalStatsAll { arg_count: 4 }), // + all_touched + Arc::new(RsZonalStatsAll { arg_count: 5 }), // + exclude_nodata + Arc::new(RsZonalStatsAll { arg_count: 6 }), // + lenient ], Volatility::Immutable, ) @@ -335,17 +409,30 @@ pub fn rs_zonal_stats_all_udf() -> SedonaScalarUDF { #[derive(Debug)] struct RsZonalStatsAll { - with_options: bool, + /// Number of arguments in the matched signature (2..=6). + arg_count: usize, } impl SedonaScalarKernel for RsZonalStatsAll { fn return_type(&self, args: &[SedonaType]) -> Result> { + // Argument order mirrors Sedona Spark: (raster, zone, [band, + // [all_touched, [exclude_nodata, [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.with_options { - matchers.push(ArgMatcher::is_string()); + 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_nodata, 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) @@ -369,12 +456,38 @@ impl SedonaScalarKernel for RsZonalStatsAll { ) -> Result { let num_iterations = RasterExecutor::num_iterations_over(args); - // options (when present) is at index 2. - let options_array = self - .with_options - .then(|| expand_string_arg(&args[2], num_iterations)) + // band (index 2) only exists in the 3+-arg overloads; the 2-arg overload + // leaves it implicit. all_touched (index 3), exclude_nodata (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 options_iter = options_array.as_ref().map(|a| a.iter()); + 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_nodata_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_nodata_iter = exclude_nodata_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 = Vec::new(); @@ -388,8 +501,15 @@ impl SedonaScalarKernel for RsZonalStatsAll { 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 options_str = options_iter.as_mut().and_then(|i| i.next().flatten()); - let options = ZonalStatsOptions::parse(options_str)?; + let Some(params) = next_params( + &mut band_iter, + &mut all_touched_iter, + &mut exclude_nodata_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); @@ -409,9 +529,9 @@ impl SedonaScalarKernel for RsZonalStatsAll { "Cannot operate on geometry and raster: geometry has no CRS but raster does" ), }; - match compute_zonal_stats(gdal, raster, &geom_wkb, &options, &mut scratch)? { + match compute_zonal_stats(gdal, raster, &geom_wkb, ¶ms, &mut scratch)? { Some(stats) => append_struct_stats(&mut builder, &stats), - None if options.lenient => append_struct_null(&mut builder), + None if params.lenient => append_struct_null(&mut builder), None => return no_intersection_err(), } Ok(()) @@ -518,11 +638,11 @@ fn compute_zonal_stats( gdal: &Gdal, raster: &RasterRefImpl<'_>, geom_wkb: &[u8], - options: &ZonalStatsOptions, + params: &ZonalStatsParams, scratch: &mut Vec, ) -> Result> { let num_bands = raster.num_bands(); - let band_num = resolve_band(options.band, num_bands)?; + let band_num = resolve_band(params.band, num_bands)?; let band = raster .bands() @@ -566,7 +686,7 @@ fn compute_zonal_stats( // Rasterize the zone into a window-sized 0/1 mask (moves `geometry`, whose // only remaining use is the burn). - let mask = rasterize_zone_mask(gdal, geometry, &transform, &window, options.all_touched)?; + let mask = rasterize_zone_mask(gdal, geometry, &transform, &window, params.all_touched)?; // Read the band once (zero-copy borrow) and collect the selected values. let nd_buffer = band @@ -588,7 +708,7 @@ fn compute_zonal_stats( // 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 options.exclude_nodata { + let nodata = if params.exclude_nodata { band.nodata() } else { None @@ -623,7 +743,7 @@ fn resolve_band(band: Option, num_bands: usize) -> Result { Ok(1) } else { exec_err!( - "RS_ZonalStats: raster has {num_bands} bands; set the \"band\" option to \ + "RS_ZonalStats: raster has {num_bands} bands; pass the band argument to \ choose one (only a single-band raster may omit it)" ) } @@ -881,10 +1001,78 @@ fn compute_mode(values: &[f64]) -> f64 { fn no_intersection_err() -> Result { exec_err!( "RS_ZonalStats: the zone geometry does not intersect the raster; \ - set the \"lenient\" option to return NULL instead" + pass lenient => true to return NULL instead" ) } +/// Advance the option iterators one row (keeping every column in lockstep) and +/// assemble the resolved params, or return `None` when an explicit-but-NULL +/// value makes this a NULL output row. +/// +/// A `band_iter` of `None` is the band-less overload, whose implicit band stays +/// `None` (resolved to band 1 for a single-band raster, an error otherwise). A +/// `Some` band iterator carrying a NULL, or any NULL flag, is a NULL row. +fn next_params( + band_iter: &mut Option, + all_touched_iter: &mut F, + exclude_nodata_iter: &mut F, + lenient_iter: &mut F, +) -> Option +where + B: Iterator>, + F: Iterator>, +{ + // Advance every iterator first so a NULL-driven early return does not desync + // the columns on the next row. + let band_cell = band_iter.as_mut().map(|iter| iter.next().flatten()); + let all_touched = all_touched_iter.next().flatten(); + let exclude_nodata = exclude_nodata_iter.next().flatten(); + let lenient = lenient_iter.next().flatten(); + + let band = match band_cell { + None => None, // band-less overload: implicit band + Some(Some(b)) => Some(b), // explicit band + Some(None) => return None, // explicit NULL band -> NULL row + }; + Some(ZonalStatsParams { + band, + all_touched: all_touched?, + exclude_nodata: exclude_nodata?, + lenient: lenient?, + }) +} + +/// The boolean flag column at `args[index]` when the overload carries it +/// (`present`), otherwise a constant array of `default`. +fn expand_flag( + args: &[ColumnarValue], + index: usize, + present: bool, + default: bool, + num_iterations: usize, +) -> Result { + if present { + let array = args[index] + .clone() + .cast_to(&DataType::Boolean, None)? + .into_array(num_iterations)?; + Ok(as_boolean_array(&array)?.clone()) + } else { + let array = ScalarValue::Boolean(Some(default)).to_array_of_size(num_iterations)?; + Ok(as_boolean_array(&array)?.clone()) + } +} + +/// Cast a column to `Int64` and materialize it so its values can be iterated in +/// lockstep with the raster/zone rows. +fn expand_int64_arg(arg: &ColumnarValue, num_iterations: usize) -> Result { + let array = arg + .clone() + .cast_to(&DataType::Int64, None)? + .into_array(num_iterations)?; + Ok(as_int64_array(&array)?.clone()) +} + /// Cast a column to `Utf8` and materialize it to a `StringArray` so its values /// can be iterated in lockstep with the raster/zone rows. fn expand_string_arg(arg: &ColumnarValue, num_iterations: usize) -> Result { @@ -892,7 +1080,7 @@ fn expand_string_arg(arg: &ColumnarValue, num_iterations: usize) -> Result Vec { - let mut t = vec![ - RASTER, - SedonaType::Wkb(Edges::Planar, None), - SedonaType::Arrow(DataType::Utf8), - ]; - if with_options { - t.push(SedonaType::Arrow(DataType::Utf8)); - } - t + // ScalarValue constructors for the positional trailing arguments, so a call + // reads close to its Sedona Spark SQL form: `[band(1), stat("sum"), + // flag(true), flag(false)]` is `(raster, zone, 1, 'sum', true, false)`. + fn band(b: i64) -> ScalarValue { + ScalarValue::Int64(Some(b)) + } + fn stat(s: &str) -> ScalarValue { + ScalarValue::Utf8(Some(s.to_string())) + } + fn flag(b: bool) -> ScalarValue { + ScalarValue::Boolean(Some(b)) } - /// Invoke RS_ZonalStats on a scalar raster + zone, returning the raw result - /// so both the value and error paths are testable. - fn call_stats( + /// Invoke a zonal-stats UDF on a scalar raster + zone with the given + /// positional trailing arguments. Routing through the UDF (rather than a + /// hand-picked kernel) exercises overload selection by argument count and + /// type; the raw `ScalarValue` is returned so both value and error paths are + /// testable. + fn invoke_udf( + udf: SedonaScalarUDF, spec: &RasterSpec, - wkt: &str, - stat: &str, - options: Option<&str>, - ) -> Result { - let kernel = RsZonalStats { - with_options: options.is_some(), - }; + geom: ScalarValue, + trailing: Vec, + ) -> Result { + let mut arg_types = vec![RASTER, SedonaType::Wkb(Edges::Planar, None)]; let mut args = vec![ ColumnarValue::Scalar(spec.scalar()), - ColumnarValue::Scalar(ScalarValue::Binary(Some(make_wkb(wkt)))), - ColumnarValue::Scalar(ScalarValue::Utf8(Some(stat.to_string()))), + ColumnarValue::Scalar(geom), ]; - if let Some(o) = options { - args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some( - o.to_string(), - )))); + for value in trailing { + arg_types.push(SedonaType::Arrow(value.data_type())); + args.push(ColumnarValue::Scalar(value)); + } + match ScalarUdfTester::new(udf.into(), arg_types).invoke(args)? { + ColumnarValue::Scalar(s) => Ok(s), + other => panic!("expected a scalar result, got {other:?}"), } - kernel.invoke_batch(&stats_arg_types(options.is_some()), &args) } - /// Invoke RS_ZonalStatsAll on a scalar raster + zone. - fn call_all(spec: &RasterSpec, wkt: &str, options: Option<&str>) -> Result { - let kernel = RsZonalStatsAll { - with_options: options.is_some(), - }; - let mut arg_types = vec![RASTER, SedonaType::Wkb(Edges::Planar, None)]; - let mut args = vec![ - ColumnarValue::Scalar(spec.scalar()), - ColumnarValue::Scalar(ScalarValue::Binary(Some(make_wkb(wkt)))), - ]; - if let Some(o) = options { - arg_types.push(SedonaType::Arrow(DataType::Utf8)); - args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some( - o.to_string(), - )))); - } - kernel.invoke_batch(&arg_types, &args) + /// RS_ZonalStats over a scalar raster + zone with the given trailing args. + fn call_stats(spec: &RasterSpec, wkt: &str, trailing: Vec) -> Result { + let geom = ScalarValue::Binary(Some(make_wkb(wkt))); + invoke_udf(rs_zonal_stats_udf(), spec, geom, trailing) + } + + /// RS_ZonalStatsAll over a scalar raster + zone with the given trailing args. + fn call_all(spec: &RasterSpec, wkt: &str, trailing: Vec) -> Result { + let geom = ScalarValue::Binary(Some(make_wkb(wkt))); + invoke_udf(rs_zonal_stats_all_udf(), spec, geom, trailing) } - fn cv_f64(cv: ColumnarValue) -> Option { - match cv { - ColumnarValue::Scalar(ScalarValue::Float64(v)) => v, + fn cv_f64(s: ScalarValue) -> Option { + match s { + ScalarValue::Float64(v) => v, other => panic!("expected a Float64 scalar, got {other:?}"), } } - fn cv_struct(cv: ColumnarValue) -> Arc { - match cv { - ColumnarValue::Scalar(ScalarValue::Struct(s)) => s, + fn cv_struct(s: ScalarValue) -> Arc { + match s { + ScalarValue::Struct(s) => s, other => panic!("expected a struct scalar, got {other:?}"), } } @@ -1154,41 +1308,42 @@ mod udf_tests { fn single_stats_match_hand_computed_values() { let spec = small_raster(); // Selected pixels {1, 2, 5, 6}: exact for the integer-selection stats. + // The 3-arg overload leaves the band implicit (unambiguous single band). assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF, "count", None).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("count")]).unwrap()), Some(4.0) ); assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF, "sum", None).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("sum")]).unwrap()), Some(14.0) ); assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF, "mean", None).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("mean")]).unwrap()), Some(3.5) ); assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF, "min", None).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("min")]).unwrap()), Some(1.0) ); assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF, "max", None).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("max")]).unwrap()), Some(6.0) ); assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF, "median", None).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("median")]).unwrap()), Some(3.5) ); // All four values are unique, so every one is a mode; the tie resolves // to the largest (6), matching Sedona Spark. assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF, "mode", None).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("mode")]).unwrap()), Some(6.0) ); // Sample (n-1) variance / stddev: float accumulation, so approximate. - let var = cv_f64(call_stats(&spec, LEFT_HALF, "variance", None).unwrap()).unwrap(); + let var = cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("variance")]).unwrap()).unwrap(); assert!((var - 17.0 / 3.0).abs() < 1e-9, "variance was {var}"); - let sd = cv_f64(call_stats(&spec, LEFT_HALF, "stddev", None).unwrap()).unwrap(); + let sd = cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("stddev")]).unwrap()).unwrap(); assert!( (sd - (17.0_f64 / 3.0).sqrt()).abs() < 1e-9, "stddev was {sd}" @@ -1197,7 +1352,8 @@ mod udf_tests { #[test] fn all_returns_full_struct() { - let s = cv_struct(call_all(&small_raster(), LEFT_HALF, None).unwrap()); + // The 2-arg overload leaves the band implicit. + let s = cv_struct(call_all(&small_raster(), LEFT_HALF, vec![]).unwrap()); assert!(!s.is_null(0), "the struct itself is valid"); assert_eq!(i64_field(&s, COUNT), Some(4)); assert_eq!(f64_field(&s, SUM), Some(14.0)); @@ -1210,6 +1366,38 @@ mod udf_tests { assert!((f64_field(&s, STDDEV).unwrap() - (17.0_f64 / 3.0).sqrt()).abs() < 1e-9); } + #[test] + fn overloads_dispatch_by_arg_count() { + // The 3-arg (raster, zone, stat) and 4-arg (raster, zone, band, stat) + // overloads resolve by argument count and by the type at position 2 (a + // stat string vs. a band integer). On a single-band raster both compute + // the same mean over {1, 2, 5, 6}. + let spec = small_raster(); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("mean")]).unwrap()), + Some(3.5) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![band(1), stat("mean")]).unwrap()), + Some(3.5) + ); + // RS_ZonalStatsAll: the 2-arg and 3-arg (band) overloads agree too. + assert_eq!( + i64_field( + &cv_struct(call_all(&spec, LEFT_HALF, vec![]).unwrap()), + COUNT + ), + Some(4) + ); + assert_eq!( + i64_field( + &cv_struct(call_all(&spec, LEFT_HALF, vec![band(1)]).unwrap()), + COUNT + ), + Some(4) + ); + } + #[test] fn zone_that_selects_no_pixel_centre_is_count_zero_not_null() { // A tiny zone inside the top-left pixel (centre 0.5, 1.5) but not @@ -1217,19 +1405,19 @@ mod udf_tests { // zone still overlaps the raster extent, so count is 0 (not NULL). let tiny = "POLYGON((0.1 1.6, 0.4 1.6, 0.4 1.9, 0.1 1.9, 0.1 1.6))"; assert_eq!( - cv_f64(call_stats(&small_raster(), tiny, "count", None).unwrap()), + cv_f64(call_stats(&small_raster(), tiny, vec![stat("count")]).unwrap()), Some(0.0) ); assert_eq!( - cv_f64(call_stats(&small_raster(), tiny, "sum", None).unwrap()), + cv_f64(call_stats(&small_raster(), tiny, vec![stat("sum")]).unwrap()), None ); assert_eq!( - cv_f64(call_stats(&small_raster(), tiny, "mean", None).unwrap()), + cv_f64(call_stats(&small_raster(), tiny, vec![stat("mean")]).unwrap()), None ); - let s = cv_struct(call_all(&small_raster(), tiny, None).unwrap()); + let s = cv_struct(call_all(&small_raster(), tiny, vec![]).unwrap()); assert!( !s.is_null(0), "an intersecting-but-empty zone is a valid row" @@ -1242,15 +1430,30 @@ mod udf_tests { #[test] fn all_touched_selects_the_touched_pixel() { // The same tiny zone, with all_touched, burns the pixel it lies inside - // (value 1) even though it misses the centre. + // (value 1) even though it misses the centre. all_touched first appears + // in the 5-arg overload (raster, zone, band, stat, all_touched), so the + // band must be named to reach it. let tiny = "POLYGON((0.1 1.6, 0.4 1.6, 0.4 1.9, 0.1 1.9, 0.1 1.6))"; - let opts = r#"{"all_touched": true}"#; assert_eq!( - cv_f64(call_stats(&small_raster(), tiny, "count", Some(opts)).unwrap()), + cv_f64( + call_stats( + &small_raster(), + tiny, + vec![band(1), stat("count"), flag(true)] + ) + .unwrap() + ), Some(1.0) ); assert_eq!( - cv_f64(call_stats(&small_raster(), tiny, "sum", Some(opts)).unwrap()), + cv_f64( + call_stats( + &small_raster(), + tiny, + vec![band(1), stat("sum"), flag(true)] + ) + .unwrap() + ), Some(1.0) ); } @@ -1260,22 +1463,34 @@ mod udf_tests { let far = "POLYGON((100 100, 101 100, 101 101, 100 101, 100 100))"; // Lenient (default): the whole value is NULL, including count. assert_eq!( - cv_f64(call_stats(&small_raster(), far, "count", None).unwrap()), + cv_f64(call_stats(&small_raster(), far, vec![stat("count")]).unwrap()), None ); - assert!(cv_struct(call_all(&small_raster(), far, None).unwrap()).is_null(0)); + assert!(cv_struct(call_all(&small_raster(), far, vec![]).unwrap()).is_null(0)); - // Strict: both functions error. - let err = call_stats(&small_raster(), far, "count", Some(r#"{"lenient": false}"#)) - .unwrap_err() - .to_string(); + // Strict (lenient => false): both functions error. RS_ZonalStats reaches + // lenient only in its 7-arg overload, whose trailing flags are + // (all_touched, exclude_nodata, lenient). + let err = call_stats( + &small_raster(), + far, + vec![band(1), stat("count"), flag(false), flag(true), flag(false)], + ) + .unwrap_err() + .to_string(); assert!( err.contains("does not intersect"), "unexpected error: {err}" ); - let err = call_all(&small_raster(), far, Some(r#"{"lenient": false}"#)) - .unwrap_err() - .to_string(); + // RS_ZonalStatsAll's 6-arg overload trails (all_touched, exclude_nodata, + // lenient) after the band. + let err = call_all( + &small_raster(), + far, + vec![band(1), flag(false), flag(true), flag(false)], + ) + .unwrap_err() + .to_string(); assert!( err.contains("does not intersect"), "unexpected error: {err}" @@ -1294,17 +1509,16 @@ mod udf_tests { // Lenient (default): NULL, not count 0. assert_eq!( - cv_f64(call_stats(&small_raster(), disjoint, "count", None).unwrap()), + cv_f64(call_stats(&small_raster(), disjoint, vec![stat("count")]).unwrap()), None ); - assert!(cv_struct(call_all(&small_raster(), disjoint, None).unwrap()).is_null(0)); + assert!(cv_struct(call_all(&small_raster(), disjoint, vec![]).unwrap()).is_null(0)); - // Strict: both functions error. + // Strict (lenient => false): both functions error. let err = call_stats( &small_raster(), disjoint, - "count", - Some(r#"{"lenient": false}"#), + vec![band(1), stat("count"), flag(false), flag(true), flag(false)], ) .unwrap_err() .to_string(); @@ -1312,9 +1526,13 @@ mod udf_tests { err.contains("does not intersect"), "unexpected error: {err}" ); - let err = call_all(&small_raster(), disjoint, Some(r#"{"lenient": false}"#)) - .unwrap_err() - .to_string(); + let err = call_all( + &small_raster(), + disjoint, + vec![band(1), flag(false), flag(true), flag(false)], + ) + .unwrap_err() + .to_string(); assert!( err.contains("does not intersect"), "unexpected error: {err}" @@ -1332,21 +1550,36 @@ mod udf_tests { .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]); // Default excludes the nodata pixel: {10, 20, 30}. assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF_FULL, "count", None).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF_FULL, vec![stat("count")]).unwrap()), Some(3.0) ); assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF_FULL, "sum", None).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF_FULL, vec![stat("sum")]).unwrap()), Some(60.0) ); - // exclude_nodata=false keeps it: {10, 255, 20, 30}. - let keep = r#"{"exclude_nodata": false}"#; + // exclude_nodata => false keeps it: {10, 255, 20, 30}. It first appears + // in the 6-arg overload, whose trailing flags are (all_touched, + // exclude_nodata). assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF_FULL, "count", Some(keep)).unwrap()), + cv_f64( + call_stats( + &spec, + LEFT_HALF_FULL, + vec![band(1), stat("count"), flag(false), flag(false)] + ) + .unwrap() + ), Some(4.0) ); assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF_FULL, "sum", Some(keep)).unwrap()), + cv_f64( + call_stats( + &spec, + LEFT_HALF_FULL, + vec![band(1), stat("sum"), flag(false), flag(false)] + ) + .unwrap() + ), Some(315.0) ); } @@ -1355,28 +1588,35 @@ mod udf_tests { const LEFT_HALF_FULL: &str = "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))"; #[test] - fn multiband_raster_requires_the_band_option() { + fn multiband_raster_requires_the_band_argument() { let spec = RasterSpec::d2(2, 2) .band_values(&[1u8, 2, 3, 4]) .band_values(&[10u8, 20, 30, 40]) .crs(None) .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]); - // Omitting the band on a multiband raster errors rather than defaulting. - let err = call_stats(&spec, LEFT_HALF_FULL, "sum", None) + // Omitting the band on a multiband raster errors rather than defaulting + // to band 1 (this deliberately diverges from Sedona Spark). The 3-arg + // RS_ZonalStats overload and the 2-arg RS_ZonalStatsAll overload both + // leave the band implicit. + let err = call_stats(&spec, LEFT_HALF_FULL, vec![stat("sum")]) + .unwrap_err() + .to_string(); + assert!(err.contains("2 bands"), "unexpected error: {err}"); + let err = call_all(&spec, LEFT_HALF_FULL, vec![]) .unwrap_err() .to_string(); assert!(err.contains("2 bands"), "unexpected error: {err}"); // Naming the band selects it (band 1 sums to 10, band 2 to 100). assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF_FULL, "sum", Some(r#"{"band": 1}"#)).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF_FULL, vec![band(1), stat("sum")]).unwrap()), Some(10.0) ); assert_eq!( - cv_f64(call_stats(&spec, LEFT_HALF_FULL, "sum", Some(r#"{"band": 2}"#)).unwrap()), + cv_f64(call_stats(&spec, LEFT_HALF_FULL, vec![band(2), stat("sum")]).unwrap()), Some(100.0) ); // An out-of-range band errors. - let err = call_stats(&spec, LEFT_HALF_FULL, "sum", Some(r#"{"band": 3}"#)) + let err = call_stats(&spec, LEFT_HALF_FULL, vec![band(3), stat("sum")]) .unwrap_err() .to_string(); assert!(err.contains("out of range"), "unexpected error: {err}"); @@ -1384,7 +1624,7 @@ mod udf_tests { #[test] fn unknown_statistic_errors() { - let err = call_stats(&small_raster(), LEFT_HALF, "bogus", None) + let err = call_stats(&small_raster(), LEFT_HALF, vec![stat("bogus")]) .unwrap_err() .to_string(); assert!(err.contains("unknown statistic"), "unexpected error: {err}"); @@ -1392,34 +1632,20 @@ mod udf_tests { #[test] fn null_raster_or_zone_yields_null() { - // A NULL zone geometry propagates to a NULL result. - let kernel = RsZonalStats { - with_options: false, - }; - let result = kernel - .invoke_batch( - &stats_arg_types(false), - &[ - ColumnarValue::Scalar(small_raster().scalar()), - ColumnarValue::Scalar(ScalarValue::Binary(None)), - ColumnarValue::Scalar(ScalarValue::Utf8(Some("count".to_string()))), - ], - ) - .unwrap(); - assert_eq!(cv_f64(result), None); + // A NULL zone geometry propagates to a NULL result (3-arg overload). + let null_zone = invoke_udf( + rs_zonal_stats_udf(), + &small_raster(), + ScalarValue::Binary(None), + vec![stat("count")], + ) + .unwrap(); + assert_eq!(cv_f64(null_zone), None); // A NULL statistic name also yields NULL. - let result = kernel - .invoke_batch( - &stats_arg_types(false), - &[ - ColumnarValue::Scalar(small_raster().scalar()), - ColumnarValue::Scalar(ScalarValue::Binary(Some(make_wkb(LEFT_HALF)))), - ColumnarValue::Scalar(ScalarValue::Utf8(None)), - ], - ) - .unwrap(); - assert_eq!(cv_f64(result), None); + let null_stat = + call_stats(&small_raster(), LEFT_HALF, vec![ScalarValue::Utf8(None)]).unwrap(); + assert_eq!(cv_f64(null_stat), None); } #[test] From d6bbdf7bd8ed46a07d989689e5a394b1639bdf61 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Thu, 23 Jul 2026 09:38:39 -0700 Subject: [PATCH 06/16] docs(rs_zonalstats): drop Sedona Spark comparisons from qmd prose --- docs/reference/sql/rs_zonalstats.qmd | 8 +++----- docs/reference/sql/rs_zonalstatsall.qmd | 13 ++++++------- 2 files changed, 9 insertions(+), 12 deletions(-) diff --git a/docs/reference/sql/rs_zonalstats.qmd b/docs/reference/sql/rs_zonalstats.qmd index efafdd3470..f049de6e19 100644 --- a/docs/reference/sql/rs_zonalstats.qmd +++ b/docs/reference/sql/rs_zonalstats.qmd @@ -90,7 +90,7 @@ kernels: ## Description `RS_ZonalStats` returns one summary statistic of the pixels of a single band -that a zone geometry covers, following Apache Sedona's `RS_ZonalStats`. A pixel +that a zone geometry covers. A pixel is included when its center falls inside the zone (or, with `all_touched`, when the zone touches it at all). By default the band's nodata pixels are excluded. @@ -104,12 +104,10 @@ other statistic is NULL. When the zone does not intersect the raster at all, the result is NULL under the default `lenient` behavior, or an error when `lenient` is set to false. -The positional overloads mirror Apache Sedona Spark's `RS_ZonalStats` verbatim. The `band`, `all_touched`, `exclude_nodata`, and `lenient` arguments are added one at a time by the wider overloads; `all_touched` defaults to false, -`exclude_nodata` to true, and `lenient` to true. Unlike Sedona Spark, the -band-less overload does not default to band 1 on a multiband raster: naming the -band is required there. This function operates on 2-D `(y, x)` bands; computing a +`exclude_nodata` to true, and `lenient` to true. The band-less overload does +not default to band 1 on a multiband raster: naming the band is required there. This function operates on 2-D `(y, x)` bands; computing a statistic per non-spatial plane of an N-D band is not supported. Use [`RS_ZonalStatsAll`](rs_zonalstatsall.qmd) to compute every statistic at diff --git a/docs/reference/sql/rs_zonalstatsall.qmd b/docs/reference/sql/rs_zonalstatsall.qmd index c252f324a2..9bf850eb17 100644 --- a/docs/reference/sql/rs_zonalstatsall.qmd +++ b/docs/reference/sql/rs_zonalstatsall.qmd @@ -81,8 +81,8 @@ kernels: `RS_ZonalStatsAll` returns every summary statistic of the pixels of a single band that a zone geometry covers, as a struct with fields `count`, `sum`, -`mean`, `median`, `mode`, `stddev`, `variance`, `min`, and `max` (the field -order of Apache Sedona's `RS_ZonalStatsAll`). A pixel is included when its +`mean`, `median`, `mode`, `stddev`, `variance`, `min`, and `max`. A pixel is +included when its center falls inside the zone (or, with `all_touched`, when the zone touches it at all), and the band's nodata pixels are excluded by default. @@ -95,11 +95,10 @@ other field is NULL. When the zone does not intersect the raster at all, the whole struct is NULL under the default `lenient` behavior, or the call raises an error when `lenient` is set to false. -The positional overloads mirror Apache Sedona Spark's `RS_ZonalStatsAll` -verbatim (the same ladder as `RS_ZonalStats` without `stat_type`). The `band`, -`all_touched`, `exclude_nodata`, and `lenient` arguments are added one at a time -by the wider overloads; `all_touched` defaults to false, `exclude_nodata` to -true, and `lenient` to true. Unlike Sedona Spark, the band-less overload does +The overloads are the same ladder as `RS_ZonalStats` without `stat_type`. The +`band`, `all_touched`, `exclude_nodata`, and `lenient` arguments are added one +at a time by the wider overloads; `all_touched` defaults to false, +`exclude_nodata` to true, and `lenient` to true. The band-less overload does not default to band 1 on a multiband raster: naming the band is required there. This function operates on 2-D `(y, x)` bands; computing statistics per non-spatial plane of an N-D band is not supported. From f78249934c72e8ecad899cb3512404d951ca5dca Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 27 Jul 2026 01:32:23 -0700 Subject: [PATCH 07/16] refactor(rust/sedona-raster-gdal): share geometry-mask machinery between RS_Clip and RS_ZonalStats Extract the pixel-window addressing and geometry rasterization that RS_Clip and RS_ZonalStats each carried into a shared mask module: one PixelWindow, one envelope_window, one rasterize_geometry_mask. RS_Clip now builds its geotransform from raster.transform() instead of an inline coefficient array. Behavior is unchanged; the deliberately-different per-pixel consumers (apply_mask_to_band writes nodata, collect_masked_values reads f64) stay separate. --- rust/sedona-raster-gdal/src/lib.rs | 1 + rust/sedona-raster-gdal/src/mask.rs | 155 ++++++++++++++++++ rust/sedona-raster-gdal/src/rs_clip.rs | 143 ++-------------- rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 117 +------------ 4 files changed, 172 insertions(+), 244 deletions(-) create mode 100644 rust/sedona-raster-gdal/src/mask.rs diff --git a/rust/sedona-raster-gdal/src/lib.rs b/rust/sedona-raster-gdal/src/lib.rs index f1d692eff6..2e9f72c1e1 100644 --- a/rust/sedona-raster-gdal/src/lib.rs +++ b/rust/sedona-raster-gdal/src/lib.rs @@ -32,6 +32,7 @@ mod gdal_common; #[allow(dead_code)] mod gdal_dataset_provider; +mod mask; mod raster_loader; mod rs_as_geotiff; mod rs_as_raster; diff --git a/rust/sedona-raster-gdal/src/mask.rs b/rust/sedona-raster-gdal/src/mask.rs new file mode 100644 index 0000000000..2385d51d91 --- /dev/null +++ b/rust/sedona-raster-gdal/src/mask.rs @@ -0,0 +1,155 @@ +// 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. + +//! Geometry masking machinery shared by the raster functions that select the +//! pixels a geometry covers (RS_Clip, RS_ZonalStats). +//! +//! A mask is built in two steps: [`envelope_window`] clamps the geometry's +//! envelope to a rectangular pixel window on the raster grid, and +//! [`rasterize_geometry_mask`] burns the geometry into a window-sized 0/1 `u8` +//! mask. Callers then interpret the mask however they need — RS_Clip writes +//! nodata outside it, RS_ZonalStats reads the selected pixel values — so this +//! module owns only the window addressing and rasterization, not the +//! per-pixel consumption. + +use datafusion_common::{exec_datafusion_err, Result}; +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; + +/// A rectangular pixel window (offset + size) into a raster grid. +#[derive(Debug, Clone, Copy)] +pub struct PixelWindow { + pub col_off: usize, + pub row_off: usize, + pub width: usize, + pub height: usize, +} + +/// The geometry's envelope intersected with the raster extent, snapped outward +/// to whole pixels. `None` when the clamped window has no area (the envelope +/// falls entirely outside the raster, or only touches its boundary). +/// +/// This is the window PostGIS ST_Clip, `gdalwarp -crop_to_cutline`, and Sedona +/// Spark's raster functions use. All four envelope corners are mapped through +/// the inverse geotransform (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. +pub fn envelope_window( + geometry: &Geometry, + transform: &GeoTransform, + width: usize, + height: usize, +) -> Result> { + let env = geometry.envelope(); + let inverse = transform + .invert() + .map_err(|e| exec_datafusion_err!("raster mask: 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(PixelWindow { + col_off: col0 as usize, + row_off: row0 as usize, + width: (col1 - col0) as usize, + height: (row1 - row0) as usize, + })) +} + +/// Rasterize `geometry` into a window-sized `u8` mask: 1 where the geometry +/// covers a pixel, 0 elsewhere. +/// +/// The mask is a MEM UInt8 dataset covering only `window`, with the raster +/// geotransform shifted to the window's upper-left corner so pixel indices in +/// the mask line up with the same-offset pixels of the source raster. GDAL's +/// MEM driver zero-fills the band on creation, so only the burn (value 1, +/// inside the geometry) has to be written. `geometry` is consumed, since the +/// burn is its only remaining use. +pub fn rasterize_geometry_mask( + gdal: &Gdal, + geometry: Geometry, + transform: &GeoTransform, + window: &PixelWindow, + all_touched: bool, +) -> Result> { + let mask_dataset = + MemDatasetBuilder::create(gdal, window.width, window.height, 1, GdalDataType::UInt8) + .map_err(|e| exec_datafusion_err!("raster mask: failed to create mask dataset: {e}"))?; + let (window_ulx, window_uly) = transform.apply(window.col_off as f64, window.row_off as f64); + let mask_transform = [ + window_ulx, + transform[1], + transform[2], + window_uly, + transform[4], + transform[5], + ]; + mask_dataset + .set_geo_transform(&mask_transform) + .map_err(|e| exec_datafusion_err!("raster mask: failed to set mask geotransform: {e}"))?; + + gdal.rasterize_affine(&mask_dataset, &[1], &[geometry], &[1.0], all_touched) + .map_err(|e| exec_datafusion_err!("raster mask: failed to rasterize geometry: {e}"))?; + + let mask_band = mask_dataset + .rasterband(1) + .map_err(|e| exec_datafusion_err!("raster mask: failed to read 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!("raster mask: failed to read mask: {e}"))?; + Ok(mask_buffer.data().to_vec()) +} diff --git a/rust/sedona-raster-gdal/src/rs_clip.rs b/rust/sedona-raster-gdal/src/rs_clip.rs index 8134c01097..4a6a889c4e 100644 --- a/rust/sedona-raster-gdal/src/rs_clip.rs +++ b/rust/sedona-raster-gdal/src/rs_clip.rs @@ -34,10 +34,7 @@ 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_gdal::geo_transform::GeoTransform; use arrow_schema::DataType; use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; @@ -55,6 +52,7 @@ 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}; use sedona_raster::traits::nodata_f64_to_bytes; /// RS_Clip() scalar UDF implementation @@ -364,16 +362,7 @@ struct ClippedRasterData { /// 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, + crop_window: Option, } /// Clip a raster to a geometry. @@ -399,14 +388,9 @@ fn clip_raster( .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(), - ]; + // GDAL geotransform: [upper_left_x, scale_x, skew_x, upper_left_y, skew_y, scale_y]. + let geotransform: GeoTransform = <[f64; 6]>::try_from(raster.transform()) + .map_err(|_| exec_datafusion_err!("RS_Clip: expected a 6-element geotransform"))?; // The clip window is the geometry's envelope intersected with the raster // extent, snapped outward to the pixel grid — the window PostGIS ST_Clip, @@ -417,49 +401,9 @@ fn clip_raster( 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(); + // Rasterize the geometry into a window-sized 0/1 mask: 1 inside, 0 outside + // (moves `geometry`, whose only remaining use is the burn). + let mask = rasterize_geometry_mask(gdal, geometry, &geotransform, &window, all_touched)?; // The envelope may overlap the raster while the geometry itself selects no // pixel (e.g. it falls between pixel centers); that is still the @@ -586,7 +530,7 @@ fn clip_raster( if let Some(cw) = crop_window { apply_mask_and_crop( plane_bytes, - mask, + &mask, width, &data_type, &nodata_bytes, @@ -596,7 +540,7 @@ fn clip_raster( } else { apply_mask_to_band( plane_bytes, - mask, + &mask, width, &data_type, &nodata_bytes, @@ -632,67 +576,6 @@ fn clip_raster( })) } -/// 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. The plane's bytes are appended @@ -704,7 +587,7 @@ fn apply_mask_to_band( width: usize, data_type: &BandDataType, nodata_bytes: &[u8], - window: &CropWindow, + window: &PixelWindow, out: &mut Vec, ) -> Result<()> { let byte_size = data_type.byte_size(); @@ -754,7 +637,7 @@ fn apply_mask_and_crop( full_width: usize, data_type: &BandDataType, nodata_bytes: &[u8], - cw: &CropWindow, + cw: &PixelWindow, out: &mut Vec, ) -> Result<()> { let byte_size = data_type.byte_size(); diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs index 28c72f7a08..4c7d5cbb50 100644 --- a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -63,10 +63,7 @@ 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, GeoTransformEx}; -use sedona_gdal::mem::MemDatasetBuilder; -use sedona_gdal::raster::types::GdalDataType; -use sedona_gdal::vector::geometry::Geometry; +use sedona_gdal::geo_transform::GeoTransform; use sedona_raster::array::RasterRefImpl; use sedona_raster::traits::RasterRef; use sedona_raster_functions::crs_utils::{crs_transform_wkb, resolve_crs, with_crs_engine}; @@ -79,6 +76,7 @@ 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. @@ -613,15 +611,6 @@ fn append_struct_stats(builder: &mut StructBuilder, stats: &ZonalStatistics) { // Core computation // ============================================================================= -/// A rectangular pixel window (offset + size) into the raster grid. -#[derive(Clone, Copy, Debug)] -struct PixelWindow { - col_off: usize, - row_off: usize, - width: usize, - height: usize, -} - /// Compute the statistics of the pixels a zone geometry selects on one band. /// /// Returns `Ok(None)` when the zone geometry does not intersect the raster's @@ -686,7 +675,7 @@ fn compute_zonal_stats( // Rasterize the zone into a window-sized 0/1 mask (moves `geometry`, whose // only remaining use is the burn). - let mask = rasterize_zone_mask(gdal, geometry, &transform, &window, params.all_touched)?; + let mask = rasterize_geometry_mask(gdal, geometry, &transform, &window, params.all_touched)?; // Read the band once (zero-copy borrow) and collect the selected values. let nd_buffer = band @@ -758,106 +747,6 @@ fn raster_transform(raster: &RasterRefImpl<'_>) -> Result { .map_err(|_| exec_datafusion_err!("RS_ZonalStats: expected a 6-element geotransform")) } -/// The zone's envelope intersected with the raster extent, snapped outward to -/// whole pixels. `None` when the clamped window has no area (the envelope only -/// touches the raster boundary). Mirrors the window RS_Clip / PostGIS ST_Clip -/// use: all four corners are mapped through the inverse geotransform so a skewed -/// raster still gets a correct superset window. -fn envelope_window( - geometry: &Geometry, - transform: &GeoTransform, - width: usize, - height: usize, -) -> Result> { - let env = geometry.envelope(); - let inverse = transform - .invert() - .map_err(|e| exec_datafusion_err!("RS_ZonalStats: 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(PixelWindow { - col_off: col0 as usize, - row_off: row0 as usize, - width: (col1 - col0) as usize, - height: (row1 - row0) as usize, - })) -} - -/// Rasterize the zone into a window-sized `u8` mask: 1 where the zone covers a -/// pixel, 0 elsewhere. The MEM band is zero-filled on creation, so only the -/// burn (value 1, inside the geometry) has to be written. -fn rasterize_zone_mask( - gdal: &Gdal, - geometry: Geometry, - transform: &GeoTransform, - window: &PixelWindow, - all_touched: bool, -) -> Result> { - let mask_dataset = - MemDatasetBuilder::create(gdal, window.width, window.height, 1, GdalDataType::UInt8) - .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to create mask: {e}"))?; - let (window_ulx, window_uly) = transform.apply(window.col_off as f64, window.row_off as f64); - let mask_transform = [ - window_ulx, - transform[1], - transform[2], - window_uly, - transform[4], - transform[5], - ]; - mask_dataset - .set_geo_transform(&mask_transform) - .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to set mask geotransform: {e}"))?; - - gdal.rasterize_affine(&mask_dataset, &[1], &[geometry], &[1.0], all_touched) - .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to rasterize zone: {e}"))?; - - let mask_band = mask_dataset - .rasterband(1) - .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to read 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!("RS_ZonalStats: failed to read mask: {e}"))?; - Ok(mask_buffer.data().to_vec()) -} - /// Append every selected pixel value (masked in, and — when `nodata` is set — /// not byte-equal to the nodata sentinel) to `out` as `f64`. /// From d930b09a24350c0114026fce184ec1d8c0cd3d1e Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 27 Jul 2026 10:11:22 -0700 Subject: [PATCH 08/16] fix(rust/sedona-raster-gdal): tighten RS_ZonalStats CRS reuse, nodata, and error paths - Reuse crs_utils::align_wkb_to_crs in both kernels: the equal-CRS fast path now borrows the zone WKB instead of reprojecting and copying it every row, and a CRS on exactly one side still errors (same policy, shared wording). - Guard that the band's nodata sentinel width matches the dtype byte size before the byte-equality compare, so a malformed nodata errors instead of silently disabling nodata exclusion. - Return sedona_internal_err! instead of unwrap() when appending struct fields, so a future field-layout change surfaces an error rather than aborting a Python release build. --- rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 105 ++++++++++-------- 1 file changed, 56 insertions(+), 49 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs index 4c7d5cbb50..a1abdab7b9 100644 --- a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -66,7 +66,7 @@ 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::{crs_transform_wkb, resolve_crs, with_crs_engine}; +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; @@ -348,21 +348,18 @@ impl SedonaScalarKernel for RsZonalStats { exec_datafusion_err!("RS_ZonalStats: unknown statistic {stat_str:?}") })?; - // Reproject the zone into the raster's CRS; a known/unknown - // CRS mismatch on either side would mislocate it. + // Reproject the zone 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 zone. 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(wkb, geom_crs, raster_crs, engine)? - } - (None, None) => 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" - ), - }; + 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), @@ -505,31 +502,27 @@ impl SedonaScalarKernel for RsZonalStatsAll { &mut exclude_nodata_iter, &mut lenient_iter, ) else { - append_struct_null(&mut builder); + append_struct_null(&mut builder)?; return Ok(()); }; let (Some(raster), Some(wkb)) = (raster_opt, wkb_opt) else { - append_struct_null(&mut builder); + append_struct_null(&mut builder)?; 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(wkb, geom_crs, raster_crs, engine)? - } - (None, None) => 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" - ), - }; + 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), + Some(stats) => append_struct_stats(&mut builder, &stats)?, + None if params.lenient => append_struct_null(&mut builder)?, None => return no_intersection_err(), } Ok(()) @@ -565,27 +558,28 @@ fn zonal_stats_struct_fields() -> Fields { /// Append a fully-NULL struct row (the zone does not intersect the raster and /// `lenient` is set). -fn append_struct_null(builder: &mut StructBuilder) { - builder - .field_builder::(0) - .unwrap() - .append_null(); +fn append_struct_null(builder: &mut StructBuilder) -> Result<()> { + let Some(count) = builder.field_builder::(0) else { + return sedona_internal_err!("RS_ZonalStats: count field is not an Int64 builder"); + }; + count.append_null(); for i in 1..=8 { - builder - .field_builder::(i) - .unwrap() - .append_null(); + let Some(field) = builder.field_builder::(i) else { + return sedona_internal_err!("RS_ZonalStats: stat field {i} is not a Float64 builder"); + }; + field.append_null(); } builder.append(false); + Ok(()) } /// Append one computed-stats row. The float fields carry through the `Option` /// so an empty zone records `count = 0` with the rest NULL. -fn append_struct_stats(builder: &mut StructBuilder, stats: &ZonalStatistics) { - builder - .field_builder::(0) - .unwrap() - .append_value(stats.count); +fn append_struct_stats(builder: &mut StructBuilder, stats: &ZonalStatistics) -> Result<()> { + let Some(count) = builder.field_builder::(0) else { + return sedona_internal_err!("RS_ZonalStats: count field is not an Int64 builder"); + }; + count.append_value(stats.count); for (i, value) in [ stats.sum, stats.mean, @@ -599,12 +593,16 @@ fn append_struct_stats(builder: &mut StructBuilder, stats: &ZonalStatistics) { .into_iter() .enumerate() { - builder - .field_builder::(i + 1) - .unwrap() - .append_option(value); + let Some(field) = builder.field_builder::(i + 1) else { + return sedona_internal_err!( + "RS_ZonalStats: stat field {} is not a Float64 builder", + i + 1 + ); + }; + field.append_option(value); } builder.append(true); + Ok(()) } // ============================================================================= @@ -702,6 +700,14 @@ fn compute_zonal_stats( } 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( @@ -1087,6 +1093,7 @@ mod udf_tests { use arrow_array::{Array, StructArray}; use datafusion_expr::ScalarUDF; use sedona_proj::transform::{with_global_proj_engine, LazyProjEngine}; + use sedona_raster_functions::crs_utils::crs_transform_wkb; use sedona_schema::crs::deserialize_crs; use sedona_schema::datatypes::{Edges, RASTER}; use sedona_testing::create::make_wkb; From 563c2ab0203247c4e90c44b162b19638c095545d Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 27 Jul 2026 10:25:20 -0700 Subject: [PATCH 09/16] refactor(rust/sedona-raster-gdal): match Spark arg names; NaN stat consistency - Rename the RS_ZonalStats / RS_ZonalStatsAll arguments zone -> roi and exclude_nodata -> exclude_no_data so the whole signature mirrors Sedona Spark's getZonalStats(raster, roi, band, statType, allTouched, excludeNoData, lenient) (band / stat_type / all_touched / lenient already matched). Updates the qmd frontmatter (the generated accessor API), kernels, docs, and tests. - compute_statistics: a NaN pixel now yields NaN for every statistic (numpy semantics) instead of finite min/max (f64::min/f64::max silently skip NaN) alongside NaN sum/mean. Adds a unit test. --- docs/reference/sql/rs_zonalstats.qmd | 36 +-- docs/reference/sql/rs_zonalstatsall.qmd | 34 +-- .../tests/functions/test_rs_zonalstats.py | 60 ++--- rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 210 ++++++++++-------- 4 files changed, 189 insertions(+), 151 deletions(-) diff --git a/docs/reference/sql/rs_zonalstats.qmd b/docs/reference/sql/rs_zonalstats.qmd index f049de6e19..e792ed4800 100644 --- a/docs/reference/sql/rs_zonalstats.qmd +++ b/docs/reference/sql/rs_zonalstats.qmd @@ -18,42 +18,42 @@ title: RS_ZonalStats description: > - Computes a single summary statistic of the raster pixels covered by a zone + Computes a single summary statistic of the raster pixels covered by a roi geometry. kernels: - returns: double args: - raster - - {name: zone, type: geometry} + - {name: roi, type: geometry} - {name: stat_type, type: string} - returns: double args: - raster - - {name: zone, type: geometry} + - {name: roi, type: geometry} - {name: band, type: integer} - {name: stat_type, type: string} - returns: double args: - raster - - {name: zone, type: geometry} + - {name: roi, type: geometry} - {name: band, type: integer} - {name: stat_type, type: string} - {name: all_touched, type: boolean} - returns: double args: - raster - - {name: zone, type: geometry} + - {name: roi, type: geometry} - {name: band, type: integer} - {name: stat_type, type: string} - {name: all_touched, type: boolean} - - {name: exclude_nodata, type: boolean} + - {name: exclude_no_data, type: boolean} - returns: double args: - raster - - name: zone + - name: roi type: geometry description: > - Zone geometry defining the region of interest. Reprojected into the + Region-of-interest geometry. Reprojected into the raster's CRS when both carry one; it is an error for exactly one side to have a CRS. - name: band @@ -70,16 +70,16 @@ kernels: - name: all_touched type: boolean description: > - If true, include every pixel the zone touches; otherwise only pixels + If true, include every pixel the roi touches; otherwise only pixels whose center falls inside it. Defaults to false. - - name: exclude_nodata + - name: exclude_no_data type: boolean description: > If true (the default), skip pixels equal to the band's nodata value. - name: lenient type: boolean description: > - If true (the default), return NULL when the zone does not intersect the + If true (the default), return NULL when the roi does not intersect the raster; if false, raise an error. --- @@ -90,23 +90,23 @@ kernels: ## Description `RS_ZonalStats` returns one summary statistic of the pixels of a single band -that a zone geometry covers. A pixel -is included when its center falls inside the zone (or, with `all_touched`, when -the zone touches it at all). By default the band's nodata pixels are excluded. +that a roi geometry covers. A pixel +is included when its center falls inside the roi (or, with `all_touched`, when +the roi touches it at all). By default the band's nodata pixels are excluded. The statistic is one of `count`, `sum`, `mean`, `median`, `mode`, `stddev`, `variance`, `min`, or `max`. `count` is returned as a whole number; all other statistics are floating point. Variance and standard deviation are the sample (n-1) values, and `mode` breaks ties toward the larger value. -When the zone overlaps the raster but selects no pixel, `count` is 0 and every -other statistic is NULL. When the zone does not intersect the raster at all, the +When the roi overlaps the raster but selects no pixel, `count` is 0 and every +other statistic is NULL. When the roi does not intersect the raster at all, the result is NULL under the default `lenient` behavior, or an error when `lenient` is set to false. -The `band`, `all_touched`, `exclude_nodata`, and `lenient` arguments are added +The `band`, `all_touched`, `exclude_no_data`, and `lenient` arguments are added one at a time by the wider overloads; `all_touched` defaults to false, -`exclude_nodata` to true, and `lenient` to true. The band-less overload does +`exclude_no_data` to true, and `lenient` to true. The band-less overload does not default to band 1 on a multiband raster: naming the band is required there. This function operates on 2-D `(y, x)` bands; computing a statistic per non-spatial plane of an N-D band is not supported. diff --git a/docs/reference/sql/rs_zonalstatsall.qmd b/docs/reference/sql/rs_zonalstatsall.qmd index 9bf850eb17..4bc83669e2 100644 --- a/docs/reference/sql/rs_zonalstatsall.qmd +++ b/docs/reference/sql/rs_zonalstatsall.qmd @@ -18,38 +18,38 @@ title: RS_ZonalStatsAll description: > - Computes every summary statistic of the raster pixels covered by a zone + Computes every summary statistic of the raster pixels covered by a roi geometry and returns them as a struct. kernels: - returns: struct args: - raster - - {name: zone, type: geometry} + - {name: roi, type: geometry} - returns: struct args: - raster - - {name: zone, type: geometry} + - {name: roi, type: geometry} - {name: band, type: integer} - returns: struct args: - raster - - {name: zone, type: geometry} + - {name: roi, type: geometry} - {name: band, type: integer} - {name: all_touched, type: boolean} - returns: struct args: - raster - - {name: zone, type: geometry} + - {name: roi, type: geometry} - {name: band, type: integer} - {name: all_touched, type: boolean} - - {name: exclude_nodata, type: boolean} + - {name: exclude_no_data, type: boolean} - returns: struct args: - raster - - name: zone + - name: roi type: geometry description: > - Zone geometry defining the region of interest. Reprojected into the + Region-of-interest geometry. Reprojected into the raster's CRS when both carry one; it is an error for exactly one side to have a CRS. - name: band @@ -60,16 +60,16 @@ kernels: - name: all_touched type: boolean description: > - If true, include every pixel the zone touches; otherwise only pixels + If true, include every pixel the roi touches; otherwise only pixels whose center falls inside it. Defaults to false. - - name: exclude_nodata + - name: exclude_no_data type: boolean description: > If true (the default), skip pixels equal to the band's nodata value. - name: lenient type: boolean description: > - If true (the default), return NULL when the zone does not intersect the + If true (the default), return NULL when the roi does not intersect the raster; if false, raise an error. --- @@ -80,25 +80,25 @@ kernels: ## Description `RS_ZonalStatsAll` returns every summary statistic of the pixels of a single -band that a zone geometry covers, as a struct with fields `count`, `sum`, +band that a roi geometry covers, as a struct with fields `count`, `sum`, `mean`, `median`, `mode`, `stddev`, `variance`, `min`, and `max`. A pixel is included when its -center falls inside the zone (or, with `all_touched`, when the zone touches it +center falls inside the roi (or, with `all_touched`, when the roi touches it at all), and the band's nodata pixels are excluded by default. `count` is a whole number (a 64-bit integer); every other field is floating point. Variance and standard deviation are the sample (n-1) values, and `mode` breaks ties toward the larger value. -When the zone overlaps the raster but selects no pixel, `count` is 0 and every -other field is NULL. When the zone does not intersect the raster at all, the +When the roi overlaps the raster but selects no pixel, `count` is 0 and every +other field is NULL. When the roi does not intersect the raster at all, the whole struct is NULL under the default `lenient` behavior, or the call raises an error when `lenient` is set to false. The overloads are the same ladder as `RS_ZonalStats` without `stat_type`. The -`band`, `all_touched`, `exclude_nodata`, and `lenient` arguments are added one +`band`, `all_touched`, `exclude_no_data`, and `lenient` arguments are added one at a time by the wider overloads; `all_touched` defaults to false, -`exclude_nodata` to true, and `lenient` to true. The band-less overload does +`exclude_no_data` to true, and `lenient` to true. The band-less overload does not default to band 1 on a multiband raster: naming the band is required there. This function operates on 2-D `(y, x)` bands; computing statistics per non-spatial plane of an N-D band is not supported. diff --git a/python/sedonadb/tests/functions/test_rs_zonalstats.py b/python/sedonadb/tests/functions/test_rs_zonalstats.py index 0ad888fb86..3b78b7af1e 100644 --- a/python/sedonadb/tests/functions/test_rs_zonalstats.py +++ b/python/sedonadb/tests/functions/test_rs_zonalstats.py @@ -18,12 +18,12 @@ """RS_ZonalStats / RS_ZonalStatsAll cross-checked against a numpy reference. Both functions mirror Apache Sedona Spark's positional overloads, so the tests -call them positionally: `(raster, zone, stat_type)` / -`(raster, zone, band, stat_type[, all_touched[, exclude_nodata[, lenient]]])` +call them positionally: `(raster, roi, stat_type)` / +`(raster, roi, band, stat_type[, all_touched[, exclude_no_data[, lenient]]])` for RS_ZonalStats and the same ladder without `stat_type` for RS_ZonalStatsAll. The fixture raster is CRS-less (so nothing reprojects and pixel selection is -bit-comparable). The reference rasterizes the zone with `rasterio.features` +bit-comparable). The reference rasterizes the roi with `rasterio.features` (the same GDAL rasterizer the kernel uses) and reduces the selected pixels with numpy; exact-selection statistics (count, sum, min, max, median, mode) are compared exactly and the float-accumulation ones (mean, variance, stddev) with @@ -89,8 +89,8 @@ def fixture_raster(tmp_path): return path, data[0] -def numpy_reference(band, wkt, *, all_touched, exclude_nodata): - """Reference statistics over the pixels the zone selects, via rasterio+numpy. +def numpy_reference(band, wkt, *, all_touched, exclude_no_data): + """Reference statistics over the pixels the roi selects, via rasterio+numpy. Returns a dict of every statistic, or the sentinel string ``"empty"`` when the selection is empty (the caller maps that to count 0 / NULLs). @@ -109,7 +109,7 @@ def numpy_reference(band, wkt, *, all_touched, exclude_nodata): dtype="uint8", ) sel = band[mask == 1].astype(np.float64) - if exclude_nodata: + if exclude_no_data: sel = sel[sel != NODATA] if sel.size == 0: return "empty" @@ -131,7 +131,7 @@ def numpy_reference(band, wkt, *, all_touched, exclude_nodata): def _one_row(con, path, wkt): - """A one-row frame with the raster and zone as columns, so the kernel runs + """A one-row frame with the raster and roi as columns, so the kernel runs its real per-row array path rather than constant-folding a scalar.""" df = con.create_data_frame( pa.table( @@ -147,9 +147,9 @@ def _one_row(con, path, wkt): def zonal_stat(con, path, wkt, trailing): """RS_ZonalStats over a one-row table. - `trailing` is the positional argument list after `(raster, zone)` — e.g. - `["mean"]` (band-less overload) or `[1, "mean", all_touched, exclude_nodata, - lenient]`. Raster and zone travel as columns; the trailing scalars are + `trailing` is the positional argument list after `(raster, roi)` — e.g. + `["mean"]` (band-less overload) or `[1, "mean", all_touched, exclude_no_data, + lenient]`. Raster and roi travel as columns; the trailing scalars are literals, matching how the SQL reads. """ df, raster, geom = _one_row(con, path, wkt) @@ -160,8 +160,8 @@ def zonal_stat(con, path, wkt, trailing): def zonal_stats_all(con, path, wkt, trailing): """RS_ZonalStatsAll over a one-row table; returns the struct as a dict. - `trailing` is the positional argument list after `(raster, zone)` — e.g. - `[]` (band-less overload) or `[1, all_touched, exclude_nodata, lenient]`. + `trailing` is the positional argument list after `(raster, roi)` — e.g. + `[]` (band-less overload) or `[1, all_touched, exclude_no_data, lenient]`. """ df, raster, geom = _one_row(con, path, wkt) table = df.select(r=raster.funcs.rs_zonalstatsall(geom, *trailing)).to_arrow_table() @@ -173,11 +173,11 @@ def zonal_stats_all(con, path, wkt, trailing): def test_single_stat_matches_numpy(con, tmp_path, stat, all_touched): path, band = fixture_raster(tmp_path) expected = numpy_reference( - band, GEOM_RECT, all_touched=all_touched, exclude_nodata=True + band, GEOM_RECT, all_touched=all_touched, exclude_no_data=True ) assert expected != "empty", "GEOM_RECT should select pixels" - # (raster, zone, band, stat_type, all_touched) — the 5-arg overload. + # (raster, roi, band, stat_type, all_touched) — the 5-arg overload. got = zonal_stat(con, path, GEOM_RECT, [1, stat, all_touched]) if stat in EXACT_STATS: assert got == expected[stat] @@ -187,8 +187,8 @@ def test_single_stat_matches_numpy(con, tmp_path, stat, all_touched): def test_all_struct_matches_numpy(con, tmp_path): path, band = fixture_raster(tmp_path) - expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=True) - # (raster, zone, band) — all_touched defaults to false. + expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_no_data=True) + # (raster, roi, band) — all_touched defaults to false. got = zonal_stats_all(con, path, GEOM_RECT, [1]) # count is an integer field (Int64); every other field is floating point. @@ -201,17 +201,19 @@ def test_all_struct_matches_numpy(con, tmp_path): assert got[stat] == pytest.approx(expected[stat]) -def test_exclude_nodata_default_and_disabled(con, tmp_path): +def test_exclude_no_data_default_and_disabled(con, tmp_path): path, band = fixture_raster(tmp_path) # Default excludes nodata; disabling it keeps those pixels, raising count. - excluded = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=True) - included = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=False) + excluded = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_no_data=True) + included = numpy_reference( + band, GEOM_RECT, all_touched=False, exclude_no_data=False + ) assert included["count"] > excluded["count"] - # Default (4-arg (raster, zone, band, stat_type)) excludes nodata. + # Default (4-arg (raster, roi, band, stat_type)) excludes nodata. assert zonal_stat(con, path, GEOM_RECT, [1, "count"]) == excluded["count"] - # exclude_nodata => false keeps it: the 6-arg overload trails (all_touched, - # exclude_nodata). + # exclude_no_data => false keeps it: the 6-arg overload trails (all_touched, + # exclude_no_data). assert ( zonal_stat(con, path, GEOM_RECT, [1, "count", False, False]) == included["count"] @@ -220,7 +222,7 @@ def test_exclude_nodata_default_and_disabled(con, tmp_path): def test_sliver_selects_nothing_unless_all_touched(con, tmp_path): path, _ = fixture_raster(tmp_path) - # The zone overlaps the raster but covers no pixel center: count 0, rest NULL. + # The roi overlaps the raster but covers no pixel center: count 0, rest NULL. assert zonal_stat(con, path, GEOM_SLIVER, [1, "count"]) == 0.0 assert zonal_stat(con, path, GEOM_SLIVER, [1, "sum"]) is None # all_touched (5-arg overload) picks up the pixels it crosses. @@ -234,7 +236,7 @@ def test_no_intersection_is_null_when_lenient_and_errors_when_strict(con, tmp_pa assert zonal_stat(con, path, GEOM_DISJOINT, [1, "count"]) is None assert zonal_stats_all(con, path, GEOM_DISJOINT, [1]) is None # Strict (lenient => false): the 7-arg overload trails (all_touched, - # exclude_nodata, lenient). + # exclude_no_data, lenient). with pytest.raises(Exception, match="does not intersect"): zonal_stat(con, path, GEOM_DISJOINT, [1, "count", False, True, False]) @@ -245,7 +247,7 @@ def test_bbox_overlapping_but_geometry_disjoint_is_no_intersection(con, tmp_path path, _ = fixture_raster(tmp_path) - # Premise: the zone's bounding box overlaps the raster extent, but the + # Premise: the roi's bounding box overlaps the raster extent, but the # geometry is disjoint from it (unlike GEOM_DISJOINT, whose bbox misses too). ox, px, _, oy, _, py = GDAL_TRANSFORM raster_extent = box(ox, oy + py * HEIGHT, ox + px * WIDTH, oy) @@ -256,7 +258,7 @@ def test_bbox_overlapping_but_geometry_disjoint_is_no_intersection(con, tmp_path # Lenient (default): NULL, not count 0 — a true no-intersection case. assert zonal_stat(con, path, GEOM_DISJOINT_BBOX, [1, "count"]) is None assert zonal_stats_all(con, path, GEOM_DISJOINT_BBOX, [1]) is None - # Strict: errors, exactly like the fully-disjoint zone. + # Strict: errors, exactly like the fully-disjoint roi. with pytest.raises(Exception, match="does not intersect"): zonal_stat(con, path, GEOM_DISJOINT_BBOX, [1, "count", False, True, False]) @@ -274,10 +276,10 @@ def test_implicit_band_on_multiband_raster_errors(con, tmp_path): path = tmp_path / "multiband.tif" write_geotiff(path, data, gdal_transform=GDAL_TRANSFORM) - # RS_ZonalStats 3-arg (raster, zone, stat_type): band implicit. + # RS_ZonalStats 3-arg (raster, roi, stat_type): band implicit. with pytest.raises(Exception, match="2 bands"): zonal_stat(con, path, GEOM_RECT, ["count"]) - # RS_ZonalStatsAll 2-arg (raster, zone): band implicit. + # RS_ZonalStatsAll 2-arg (raster, roi): band implicit. with pytest.raises(Exception, match="2 bands"): zonal_stats_all(con, path, GEOM_RECT, []) # Naming the band resolves the ambiguity. @@ -287,7 +289,7 @@ def test_implicit_band_on_multiband_raster_errors(con, tmp_path): def test_sql_text_smoke(con, tmp_path): """One raw-SQL invocation per function keeps the parser path covered.""" path, band = fixture_raster(tmp_path) - expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_nodata=True) + expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_no_data=True) single = con.sql( "SELECT RS_ZonalStats(RS_FromPath($1), ST_GeomFromText($2), 1, 'sum') AS r", diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs index a1abdab7b9..97ea8f4825 100644 --- a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -16,30 +16,30 @@ // under the License. //! RS_ZonalStats / RS_ZonalStatsAll UDFs — summary statistics of the raster -//! pixels covered by a zone geometry. +//! 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, zone, stat)` -//! - `RS_ZonalStats(raster, zone, band, stat)` -//! - `RS_ZonalStats(raster, zone, band, stat, all_touched)` -//! - `RS_ZonalStats(raster, zone, band, stat, all_touched, exclude_nodata)` -//! - `RS_ZonalStats(raster, zone, band, stat, all_touched, exclude_nodata, lenient)` +//! - `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, zone)` -//! - `RS_ZonalStatsAll(raster, zone, band)` -//! - `RS_ZonalStatsAll(raster, zone, band, all_touched)` -//! - `RS_ZonalStatsAll(raster, zone, band, all_touched, exclude_nodata)` -//! - `RS_ZonalStatsAll(raster, zone, band, all_touched, exclude_nodata, lenient)` +//! - `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 zone (or that the zone +//! 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_nodata` to true, and +//! 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). @@ -127,18 +127,18 @@ struct ZonalStatsParams { /// the band is required rather than silently getting band 1). Only the /// band-less overloads leave this `None`. band: Option, - /// Include every pixel the zone touches, not only those whose centre it + /// 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_nodata: bool, - /// Return NULL when the zone does not intersect the raster, rather than + 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 zone. `count` is always present (0 when the zone +/// 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). @@ -157,7 +157,7 @@ struct ZonalStatistics { impl ZonalStatistics { /// The value RS_ZonalStats returns for a single statistic. `count` is never - /// NULL (it is 0 for an empty zone); the others are NULL for an empty zone. + /// NULL (it is 0 for an empty roi); the others are NULL for an empty roi. fn get(&self, stat_type: StatType) -> Option { match stat_type { StatType::Count => Some(self.count as f64), @@ -184,10 +184,10 @@ pub fn rs_zonal_stats_udf() -> SedonaScalarUDF { SedonaScalarUDF::new( "rs_zonalstats", vec![ - Arc::new(RsZonalStats { arg_count: 3 }), // (raster, zone, stat) - Arc::new(RsZonalStats { arg_count: 4 }), // (raster, zone, band, stat) + 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_nodata + Arc::new(RsZonalStats { arg_count: 6 }), // + exclude_no_data Arc::new(RsZonalStats { arg_count: 7 }), // + lenient ], Volatility::Immutable, @@ -205,8 +205,8 @@ struct RsZonalStats { impl SedonaScalarKernel for RsZonalStats { fn return_type(&self, args: &[SedonaType]) -> Result> { - // Argument order mirrors Sedona Spark: (raster, zone, [band,] stat, - // [all_touched, [exclude_nodata, [lenient]]]). The 3-arg overload omits + // 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 { @@ -286,7 +286,7 @@ impl SedonaScalarKernel for RsZonalStats { .transpose()?; let mut band_iter = band_array.as_ref().map(|a| a.iter()); - // all_touched (index 4), exclude_nodata (index 5), lenient (index 6): + // 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, @@ -295,7 +295,7 @@ impl SedonaScalarKernel for RsZonalStats { DEFAULT_ALL_TOUCHED, num_iterations, )?; - let exclude_nodata_array = expand_flag( + let exclude_no_data_array = expand_flag( args, 5, self.arg_count >= 6, @@ -310,13 +310,13 @@ impl SedonaScalarKernel for RsZonalStats { num_iterations, )?; let mut all_touched_iter = all_touched_array.iter(); - let mut exclude_nodata_iter = exclude_nodata_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 = Vec::new(); - // The executor only sees (raster, zone); the option columns are advanced + // 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()]; @@ -331,14 +331,14 @@ impl SedonaScalarKernel for RsZonalStats { let Some(params) = next_params( &mut band_iter, &mut all_touched_iter, - &mut exclude_nodata_iter, + &mut exclude_no_data_iter, &mut lenient_iter, ) else { builder.append_null(); return Ok(()); }; - // A NULL stat, raster, or zone propagates to a NULL row. + // 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(); @@ -348,9 +348,9 @@ impl SedonaScalarKernel for RsZonalStats { exec_datafusion_err!("RS_ZonalStats: unknown statistic {stat_str:?}") })?; - // Reproject the zone into the raster's CRS, borrowing it + // 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 zone. + // 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, @@ -365,7 +365,7 @@ impl SedonaScalarKernel for RsZonalStats { Some(value) => builder.append_value(value), None => builder.append_null(), }, - // The zone does not intersect the raster: NULL when + // 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(), @@ -391,10 +391,10 @@ pub fn rs_zonal_stats_all_udf() -> SedonaScalarUDF { SedonaScalarUDF::new( "rs_zonalstatsall", vec![ - Arc::new(RsZonalStatsAll { arg_count: 2 }), // (raster, zone) - Arc::new(RsZonalStatsAll { arg_count: 3 }), // (raster, zone, band) + 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_nodata + Arc::new(RsZonalStatsAll { arg_count: 5 }), // + exclude_no_data Arc::new(RsZonalStatsAll { arg_count: 6 }), // + lenient ], Volatility::Immutable, @@ -410,8 +410,8 @@ struct RsZonalStatsAll { impl SedonaScalarKernel for RsZonalStatsAll { fn return_type(&self, args: &[SedonaType]) -> Result> { - // Argument order mirrors Sedona Spark: (raster, zone, [band, - // [all_touched, [exclude_nodata, [lenient]]]]). The 2-arg overload omits + // 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(), @@ -421,7 +421,7 @@ impl SedonaScalarKernel for RsZonalStatsAll { matchers.push(ArgMatcher::is_integer()); // band } for _ in 4..=self.arg_count { - matchers.push(ArgMatcher::is_boolean()); // all_touched, exclude_nodata, lenient + matchers.push(ArgMatcher::is_boolean()); // all_touched, exclude_no_data, lenient } if self.arg_count < 2 || self.arg_count > 6 { return sedona_internal_err!( @@ -452,7 +452,7 @@ impl SedonaScalarKernel for RsZonalStatsAll { 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_nodata (index 4), + // 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)) @@ -466,7 +466,7 @@ impl SedonaScalarKernel for RsZonalStatsAll { DEFAULT_ALL_TOUCHED, num_iterations, )?; - let exclude_nodata_array = expand_flag( + let exclude_no_data_array = expand_flag( args, 4, self.arg_count >= 5, @@ -481,7 +481,7 @@ impl SedonaScalarKernel for RsZonalStatsAll { num_iterations, )?; let mut all_touched_iter = all_touched_array.iter(); - let mut exclude_nodata_iter = exclude_nodata_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); @@ -499,7 +499,7 @@ impl SedonaScalarKernel for RsZonalStatsAll { let Some(params) = next_params( &mut band_iter, &mut all_touched_iter, - &mut exclude_nodata_iter, + &mut exclude_no_data_iter, &mut lenient_iter, ) else { append_struct_null(&mut builder)?; @@ -556,7 +556,7 @@ fn zonal_stats_struct_fields() -> Fields { ]) } -/// Append a fully-NULL struct row (the zone does not intersect the raster and +/// 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::(0) else { @@ -574,7 +574,7 @@ fn append_struct_null(builder: &mut StructBuilder) -> Result<()> { } /// Append one computed-stats row. The float fields carry through the `Option` -/// so an empty zone records `count = 0` with the rest NULL. +/// so an empty roi records `count = 0` with the rest NULL. fn append_struct_stats(builder: &mut StructBuilder, stats: &ZonalStatistics) -> Result<()> { let Some(count) = builder.field_builder::(0) else { return sedona_internal_err!("RS_ZonalStats: count field is not an Int64 builder"); @@ -609,13 +609,13 @@ fn append_struct_stats(builder: &mut StructBuilder, stats: &ZonalStatistics) -> // Core computation // ============================================================================= -/// Compute the statistics of the pixels a zone geometry selects on one band. +/// Compute the statistics of the pixels a roi geometry selects on one band. /// -/// Returns `Ok(None)` when the zone geometry does not intersect the raster's +/// Returns `Ok(None)` when the roi geometry does not intersect the raster's /// footprint. This is a true geometry intersection (matching Sedona Spark's -/// `rsIntersects` gate), not a bounding-box overlap: a zone whose envelope +/// `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 `None` into NULL when `lenient`, an error otherwise. A zone +/// The caller turns `None` 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 yields `Ok(Some(..))` with `count = 0`. /// @@ -650,18 +650,18 @@ fn compute_zonal_stats( 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 zone and + // 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 zone whose envelope overlaps the raster but whose + // 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 - // zone is already in the raster's CRS here, so no transform is needed. + // roi is already in the raster's CRS here, so no transform is needed. if !raster_intersects_geom_wkb(raster, geom_wkb)? { return Ok(None); } - // Parse the zone and clamp its envelope to the raster grid for the pixel + // 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 zone only touches the raster boundary) selects no + // 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) @@ -671,7 +671,7 @@ fn compute_zonal_stats( return Ok(Some(compute_statistics(scratch))); }; - // Rasterize the zone into a window-sized 0/1 mask (moves `geometry`, whose + // Rasterize the roi into a window-sized 0/1 mask (moves `geometry`, whose // only remaining use is the burn). let mask = rasterize_geometry_mask(gdal, geometry, &transform, &window, params.all_touched)?; @@ -695,7 +695,7 @@ fn compute_zonal_stats( // 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_nodata { + let nodata = if params.exclude_no_data { band.nodata() } else { None @@ -809,7 +809,7 @@ fn collect_masked_values( /// Compute every statistic from the selected pixel values. /// /// An empty slice yields `count = 0` and NULL for the rest (Sedona Spark's -/// empty-zone shortcut). Variance is the sample (n-1) variance, matching Spark; +/// 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 for the median reduces to the middle element (odd n) or /// the mean of the two central elements (even n). Mode is the most frequent @@ -833,6 +833,24 @@ fn compute_statistics(values: &mut [f64]) -> ZonalStatistics { }; } + // 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 n = values.len(); let sum: f64 = values.iter().sum(); let mean = sum / n as f64; @@ -892,10 +910,10 @@ fn compute_mode(values: &[f64]) -> f64 { // Argument helpers // ============================================================================= -/// The error returned for a non-intersecting zone when `lenient` is off. +/// The error returned for a non-intersecting roi when `lenient` is off. fn no_intersection_err() -> Result { exec_err!( - "RS_ZonalStats: the zone geometry does not intersect the raster; \ + "RS_ZonalStats: the roi geometry does not intersect the raster; \ pass lenient => true to return NULL instead" ) } @@ -910,7 +928,7 @@ fn no_intersection_err() -> Result { fn next_params( band_iter: &mut Option, all_touched_iter: &mut F, - exclude_nodata_iter: &mut F, + exclude_no_data_iter: &mut F, lenient_iter: &mut F, ) -> Option where @@ -921,7 +939,7 @@ where // the columns on the next row. let band_cell = band_iter.as_mut().map(|iter| iter.next().flatten()); let all_touched = all_touched_iter.next().flatten(); - let exclude_nodata = exclude_nodata_iter.next().flatten(); + let exclude_no_data = exclude_no_data_iter.next().flatten(); let lenient = lenient_iter.next().flatten(); let band = match band_cell { @@ -932,7 +950,7 @@ where Some(ZonalStatsParams { band, all_touched: all_touched?, - exclude_nodata: exclude_nodata?, + exclude_no_data: exclude_no_data?, lenient: lenient?, }) } @@ -959,7 +977,7 @@ fn expand_flag( } /// Cast a column to `Int64` and materialize it so its values can be iterated in -/// lockstep with the raster/zone rows. +/// lockstep with the raster/roi rows. fn expand_int64_arg(arg: &ColumnarValue, num_iterations: usize) -> Result { let array = arg .clone() @@ -969,7 +987,7 @@ fn expand_int64_arg(arg: &ColumnarValue, num_iterations: usize) -> Result Result { let array = arg .clone() @@ -1077,6 +1095,24 @@ mod tests { let mut values = vec![7.0, 7.0, 7.0, 1.0, 2.0]; assert_eq!(compute_statistics(&mut values).mode, Some(7.0)); } + + #[test] + fn statistics_with_nan_are_all_nan_except_count() { + // A NaN pixel poisons every statistic (numpy semantics); count still + // reflects the number of selected values. Guards against min/max + // silently skipping NaN while sum/mean propagate it. + let mut values = vec![1.0, f64::NAN, 3.0]; + let s = compute_statistics(&mut values); + assert_eq!(s.count, 3); + assert!(s.sum.unwrap().is_nan()); + assert!(s.mean.unwrap().is_nan()); + assert!(s.median.unwrap().is_nan()); + assert!(s.mode.unwrap().is_nan()); + assert!(s.min.unwrap().is_nan()); + assert!(s.max.unwrap().is_nan()); + assert!(s.variance.unwrap().is_nan()); + assert!(s.stddev.unwrap().is_nan()); + } } /// UDF-level tests: exercise the kernels end to end and pin the numbers against @@ -1127,7 +1163,7 @@ mod udf_tests { // ScalarValue constructors for the positional trailing arguments, so a call // reads close to its Sedona Spark SQL form: `[band(1), stat("sum"), - // flag(true), flag(false)]` is `(raster, zone, 1, 'sum', true, false)`. + // flag(true), flag(false)]` is `(raster, roi, 1, 'sum', true, false)`. fn band(b: i64) -> ScalarValue { ScalarValue::Int64(Some(b)) } @@ -1138,7 +1174,7 @@ mod udf_tests { ScalarValue::Boolean(Some(b)) } - /// Invoke a zonal-stats UDF on a scalar raster + zone with the given + /// Invoke a zonal-stats UDF on a scalar raster + roi with the given /// positional trailing arguments. Routing through the UDF (rather than a /// hand-picked kernel) exercises overload selection by argument count and /// type; the raw `ScalarValue` is returned so both value and error paths are @@ -1164,13 +1200,13 @@ mod udf_tests { } } - /// RS_ZonalStats over a scalar raster + zone with the given trailing args. + /// RS_ZonalStats over a scalar raster + roi with the given trailing args. fn call_stats(spec: &RasterSpec, wkt: &str, trailing: Vec) -> Result { let geom = ScalarValue::Binary(Some(make_wkb(wkt))); invoke_udf(rs_zonal_stats_udf(), spec, geom, trailing) } - /// RS_ZonalStatsAll over a scalar raster + zone with the given trailing args. + /// RS_ZonalStatsAll over a scalar raster + roi with the given trailing args. fn call_all(spec: &RasterSpec, wkt: &str, trailing: Vec) -> Result { let geom = ScalarValue::Binary(Some(make_wkb(wkt))); invoke_udf(rs_zonal_stats_all_udf(), spec, geom, trailing) @@ -1264,7 +1300,7 @@ mod udf_tests { #[test] fn overloads_dispatch_by_arg_count() { - // The 3-arg (raster, zone, stat) and 4-arg (raster, zone, band, stat) + // The 3-arg (raster, roi, stat) and 4-arg (raster, roi, band, stat) // overloads resolve by argument count and by the type at position 2 (a // stat string vs. a band integer). On a single-band raster both compute // the same mean over {1, 2, 5, 6}. @@ -1295,10 +1331,10 @@ mod udf_tests { } #[test] - fn zone_that_selects_no_pixel_centre_is_count_zero_not_null() { - // A tiny zone inside the top-left pixel (centre 0.5, 1.5) but not + fn roi_that_selects_no_pixel_centre_is_count_zero_not_null() { + // A tiny roi inside the top-left pixel (centre 0.5, 1.5) but not // covering that centre: with all_touched off, no pixel is selected. The - // zone still overlaps the raster extent, so count is 0 (not NULL). + // roi still overlaps the raster extent, so count is 0 (not NULL). let tiny = "POLYGON((0.1 1.6, 0.4 1.6, 0.4 1.9, 0.1 1.9, 0.1 1.6))"; assert_eq!( cv_f64(call_stats(&small_raster(), tiny, vec![stat("count")]).unwrap()), @@ -1316,7 +1352,7 @@ mod udf_tests { let s = cv_struct(call_all(&small_raster(), tiny, vec![]).unwrap()); assert!( !s.is_null(0), - "an intersecting-but-empty zone is a valid row" + "an intersecting-but-empty roi is a valid row" ); assert_eq!(i64_field(&s, COUNT), Some(0)); assert_eq!(f64_field(&s, SUM), None); @@ -1325,9 +1361,9 @@ mod udf_tests { #[test] fn all_touched_selects_the_touched_pixel() { - // The same tiny zone, with all_touched, burns the pixel it lies inside + // The same tiny roi, with all_touched, burns the pixel it lies inside // (value 1) even though it misses the centre. all_touched first appears - // in the 5-arg overload (raster, zone, band, stat, all_touched), so the + // in the 5-arg overload (raster, roi, band, stat, all_touched), so the // band must be named to reach it. let tiny = "POLYGON((0.1 1.6, 0.4 1.6, 0.4 1.9, 0.1 1.9, 0.1 1.6))"; assert_eq!( @@ -1366,7 +1402,7 @@ mod udf_tests { // Strict (lenient => false): both functions error. RS_ZonalStats reaches // lenient only in its 7-arg overload, whose trailing flags are - // (all_touched, exclude_nodata, lenient). + // (all_touched, exclude_no_data, lenient). let err = call_stats( &small_raster(), far, @@ -1378,7 +1414,7 @@ mod udf_tests { err.contains("does not intersect"), "unexpected error: {err}" ); - // RS_ZonalStatsAll's 6-arg overload trails (all_touched, exclude_nodata, + // RS_ZonalStatsAll's 6-arg overload trails (all_touched, exclude_no_data, // lenient) after the band. let err = call_all( &small_raster(), @@ -1394,7 +1430,7 @@ mod udf_tests { } #[test] - fn bbox_overlapping_but_geometry_disjoint_zone_is_no_intersection() { + fn bbox_overlapping_but_geometry_disjoint_roi_is_no_intersection() { // small_raster covers x ∈ [0, 4], y ∈ [0, 2]. This triangle lives on the // far side of the line x + y = 7, so its geometry is disjoint from the // raster (every raster point has x + y ≤ 6), yet its bounding box @@ -1438,7 +1474,7 @@ mod udf_tests { #[test] fn nodata_pixels_are_excluded_by_default_and_kept_when_asked() { // A 2×2 UInt8 raster [10, 255, 20, 30] with nodata 255, world extent - // x ∈ [0, 2], y ∈ [0, 2]; the zone covers all four pixels. + // x ∈ [0, 2], y ∈ [0, 2]; the roi covers all four pixels. let spec = RasterSpec::d2(2, 2) .band_values(&[10u8, 255, 20, 30]) .nodata(255u8) @@ -1453,9 +1489,9 @@ mod udf_tests { cv_f64(call_stats(&spec, LEFT_HALF_FULL, vec![stat("sum")]).unwrap()), Some(60.0) ); - // exclude_nodata => false keeps it: {10, 255, 20, 30}. It first appears + // exclude_no_data => false keeps it: {10, 255, 20, 30}. It first appears // in the 6-arg overload, whose trailing flags are (all_touched, - // exclude_nodata). + // exclude_no_data). assert_eq!( cv_f64( call_stats( @@ -1480,7 +1516,7 @@ mod udf_tests { ); } - /// A zone covering the whole 2×2 nodata raster above. + /// A roi covering the whole 2×2 nodata raster above. const LEFT_HALF_FULL: &str = "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))"; #[test] @@ -1527,16 +1563,16 @@ mod udf_tests { } #[test] - fn null_raster_or_zone_yields_null() { - // A NULL zone geometry propagates to a NULL result (3-arg overload). - let null_zone = invoke_udf( + fn null_raster_or_roi_yields_null() { + // A NULL roi geometry propagates to a NULL result (3-arg overload). + let null_roi = invoke_udf( rs_zonal_stats_udf(), &small_raster(), ScalarValue::Binary(None), vec![stat("count")], ) .unwrap(); - assert_eq!(cv_f64(null_zone), None); + assert_eq!(cv_f64(null_roi), None); // A NULL statistic name also yields NULL. let null_stat = @@ -1545,8 +1581,8 @@ mod udf_tests { } #[test] - fn reprojects_the_zone_into_the_raster_crs() { - // The raster is EPSG:4326; the zone is supplied in EPSG:3857 (the + fn reprojects_the_roi_into_the_raster_crs() { + // The raster is EPSG:4326; the roi is supplied in EPSG:3857 (the // reprojected LEFT_HALF polygon). Reprojecting it back to the raster CRS // must recover the same four-pixel selection. let spec = small_raster().crs(Some("EPSG:4326")); From da9a46e05fc0739548c6edaded8a90dbecd6a5bd Mon Sep 17 00:00:00 2001 From: James Willis Date: Tue, 28 Jul 2026 14:21:23 -0700 Subject: [PATCH 10/16] ROI clean up Co-authored-by: Dewey Dunnington --- docs/reference/sql/rs_zonalstats.qmd | 8 ++++---- docs/reference/sql/rs_zonalstatsall.qmd | 2 +- python/sedonadb/tests/functions/test_rs_zonalstats.py | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/reference/sql/rs_zonalstats.qmd b/docs/reference/sql/rs_zonalstats.qmd index e792ed4800..91fa536442 100644 --- a/docs/reference/sql/rs_zonalstats.qmd +++ b/docs/reference/sql/rs_zonalstats.qmd @@ -18,7 +18,7 @@ title: RS_ZonalStats description: > - Computes a single summary statistic of the raster pixels covered by a roi + Computes a single summary statistic of the raster pixels covered by a region of interest geometry. kernels: - returns: double @@ -90,7 +90,7 @@ kernels: ## Description `RS_ZonalStats` returns one summary statistic of the pixels of a single band -that a roi geometry covers. A pixel +that a region of interest (ROI) geometry covers. A pixel is included when its center falls inside the roi (or, with `all_touched`, when the roi touches it at all). By default the band's nodata pixels are excluded. @@ -99,8 +99,8 @@ The statistic is one of `count`, `sum`, `mean`, `median`, `mode`, `stddev`, statistics are floating point. Variance and standard deviation are the sample (n-1) values, and `mode` breaks ties toward the larger value. -When the roi overlaps the raster but selects no pixel, `count` is 0 and every -other statistic is NULL. When the roi does not intersect the raster at all, the +When the ROI overlaps the raster but selects no pixel, `count` is 0 and every +other statistic is NULL. When the ROI does not intersect the raster at all, the result is NULL under the default `lenient` behavior, or an error when `lenient` is set to false. diff --git a/docs/reference/sql/rs_zonalstatsall.qmd b/docs/reference/sql/rs_zonalstatsall.qmd index 4bc83669e2..2db54b6e9f 100644 --- a/docs/reference/sql/rs_zonalstatsall.qmd +++ b/docs/reference/sql/rs_zonalstatsall.qmd @@ -18,7 +18,7 @@ title: RS_ZonalStatsAll description: > - Computes every summary statistic of the raster pixels covered by a roi + Computes every summary statistic of the raster pixels covered by a region of interest geometry and returns them as a struct. kernels: - returns: struct diff --git a/python/sedonadb/tests/functions/test_rs_zonalstats.py b/python/sedonadb/tests/functions/test_rs_zonalstats.py index 3b78b7af1e..554d19842e 100644 --- a/python/sedonadb/tests/functions/test_rs_zonalstats.py +++ b/python/sedonadb/tests/functions/test_rs_zonalstats.py @@ -39,7 +39,7 @@ pytest.importorskip("rasterio") -from sedonadb.raster_testing import ( # noqa: E402 +from sedonadb.raster_testing import ( random_raster_data, write_geotiff, ) From 7211ba9f8398a0265c8989560f44b28c1f43204b Mon Sep 17 00:00:00 2001 From: jameswillis Date: Tue, 28 Jul 2026 14:43:32 -0700 Subject: [PATCH 11/16] test(python/sedonadb): hoist shapely import to module top --- python/sedonadb/tests/functions/test_rs_zonalstats.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/python/sedonadb/tests/functions/test_rs_zonalstats.py b/python/sedonadb/tests/functions/test_rs_zonalstats.py index 554d19842e..70e037b992 100644 --- a/python/sedonadb/tests/functions/test_rs_zonalstats.py +++ b/python/sedonadb/tests/functions/test_rs_zonalstats.py @@ -36,6 +36,8 @@ import numpy as np import pyarrow as pa import pytest +import shapely +from shapely.geometry import box pytest.importorskip("rasterio") @@ -96,7 +98,6 @@ def numpy_reference(band, wkt, *, all_touched, exclude_no_data): the selection is empty (the caller maps that to count 0 / NULLs). """ import rasterio.features - import shapely from rasterio.transform import Affine geom = shapely.from_wkt(wkt) @@ -242,9 +243,6 @@ def test_no_intersection_is_null_when_lenient_and_errors_when_strict(con, tmp_pa def test_bbox_overlapping_but_geometry_disjoint_is_no_intersection(con, tmp_path): - import shapely - from shapely.geometry import box - path, _ = fixture_raster(tmp_path) # Premise: the roi's bounding box overlaps the raster extent, but the From dea9476b0d2714537a506a053b9e2117d9e01cbd Mon Sep 17 00:00:00 2001 From: jameswillis Date: Tue, 28 Jul 2026 14:44:57 -0700 Subject: [PATCH 12/16] refactor(rust/sedona-raster-gdal): share raster geotransform helper --- rust/sedona-raster-gdal/src/gdal_common.rs | 10 ++++++++++ rust/sedona-raster-gdal/src/rs_clip.rs | 6 ++---- rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 12 ++---------- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/rust/sedona-raster-gdal/src/gdal_common.rs b/rust/sedona-raster-gdal/src/gdal_common.rs index ebb9d6277c..f082fa1f86 100644 --- a/rust/sedona-raster-gdal/src/gdal_common.rs +++ b/rust/sedona-raster-gdal/src/gdal_common.rs @@ -65,6 +65,16 @@ impl ToGdalGeoTransform for T { } } +/// A raster's stored six-coefficient GDAL geo-transform as a fixed array, +/// erroring when the transform is not exactly six elements. +/// +/// GDAL geo-transforms are +/// `[origin_x, pixel_width, rotation_x, origin_y, rotation_y, pixel_height]`. +pub fn raster_geo_transform(raster: &R) -> Result { + <[f64; 6]>::try_from(raster.transform()) + .map_err(|_| exec_datafusion_err!("expected a 6-element geotransform")) +} + /// Reconstruct raster metadata from a GDAL six-element geo-transform and raster dimensions. pub(crate) trait RasterMetadataFromGdalGeoTransform { fn to_raster_metadata(&self, width: usize, height: usize) -> RasterMetadata; diff --git a/rust/sedona-raster-gdal/src/rs_clip.rs b/rust/sedona-raster-gdal/src/rs_clip.rs index 4a6a889c4e..e1cfb80334 100644 --- a/rust/sedona-raster-gdal/src/rs_clip.rs +++ b/rust/sedona-raster-gdal/src/rs_clip.rs @@ -34,7 +34,6 @@ 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; use arrow_schema::DataType; use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; @@ -50,7 +49,7 @@ 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::{raster_geo_transform, with_gdal}; use crate::gdal_dataset_provider::configure_thread_local_options; use crate::mask::{envelope_window, rasterize_geometry_mask, PixelWindow}; use sedona_raster::traits::nodata_f64_to_bytes; @@ -389,8 +388,7 @@ fn clip_raster( .map_err(|e| exec_datafusion_err!("Failed to parse geometry from WKB: {}", e))?; // GDAL geotransform: [upper_left_x, scale_x, skew_x, upper_left_y, skew_y, scale_y]. - let geotransform: GeoTransform = <[f64; 6]>::try_from(raster.transform()) - .map_err(|_| exec_datafusion_err!("RS_Clip: expected a 6-element geotransform"))?; + let geotransform = raster_geo_transform(raster)?; // The clip window is the geometry's envelope intersected with the raster // extent, snapped outward to the pixel grid — the window PostGIS ST_Clip, diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs index 97ea8f4825..5ae0c99e1c 100644 --- a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -63,7 +63,6 @@ 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}; @@ -74,7 +73,7 @@ use sedona_schema::datatypes::SedonaType; use sedona_schema::matchers::ArgMatcher; use sedona_schema::raster::BandDataType; -use crate::gdal_common::with_gdal; +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}; @@ -644,7 +643,7 @@ fn compute_zonal_stats( let byte_size = data_type.byte_size(); let metadata = raster.metadata(); - let transform = raster_transform(raster)?; + 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()) @@ -746,13 +745,6 @@ fn resolve_band(band: Option, num_bands: usize) -> Result { } } -/// The raster's 6-coefficient GDAL geotransform as a fixed array. -fn raster_transform(raster: &RasterRefImpl<'_>) -> Result { - let t = raster.transform(); - <[f64; 6]>::try_from(t) - .map_err(|_| exec_datafusion_err!("RS_ZonalStats: expected a 6-element geotransform")) -} - /// Append every selected pixel value (masked in, and — when `nodata` is set — /// not byte-equal to the nodata sentinel) to `out` as `f64`. /// From a5fa6fa34ce475af530c386f92eee32b5f688c36 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Tue, 28 Jul 2026 14:45:23 -0700 Subject: [PATCH 13/16] perf(rust/sedona-raster-gdal): reuse a scratch buffer for the rasterized zonal mask --- rust/sedona-raster-gdal/src/mask.rs | 11 ++++- rust/sedona-raster-gdal/src/rs_clip.rs | 10 +++- rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 46 ++++++++++++++++--- 3 files changed, 57 insertions(+), 10 deletions(-) diff --git a/rust/sedona-raster-gdal/src/mask.rs b/rust/sedona-raster-gdal/src/mask.rs index 2385d51d91..ede4f04aa2 100644 --- a/rust/sedona-raster-gdal/src/mask.rs +++ b/rust/sedona-raster-gdal/src/mask.rs @@ -108,6 +108,10 @@ pub fn envelope_window( /// Rasterize `geometry` into a window-sized `u8` mask: 1 where the geometry /// covers a pixel, 0 elsewhere. /// +/// The mask is written into the caller-owned `out` buffer (cleared first), whose +/// allocation is reused across calls so a per-row rasterization does not allocate +/// a fresh `Vec` each time. It ends up `window.width * window.height` bytes long. +/// /// The mask is a MEM UInt8 dataset covering only `window`, with the raster /// geotransform shifted to the window's upper-left corner so pixel indices in /// the mask line up with the same-offset pixels of the source raster. GDAL's @@ -120,7 +124,8 @@ pub fn rasterize_geometry_mask( transform: &GeoTransform, window: &PixelWindow, all_touched: bool, -) -> Result> { + out: &mut Vec, +) -> Result<()> { let mask_dataset = MemDatasetBuilder::create(gdal, window.width, window.height, 1, GdalDataType::UInt8) .map_err(|e| exec_datafusion_err!("raster mask: failed to create mask dataset: {e}"))?; @@ -151,5 +156,7 @@ pub fn rasterize_geometry_mask( None, ) .map_err(|e| exec_datafusion_err!("raster mask: failed to read mask: {e}"))?; - Ok(mask_buffer.data().to_vec()) + out.clear(); + out.extend_from_slice(mask_buffer.data()); + Ok(()) } diff --git a/rust/sedona-raster-gdal/src/rs_clip.rs b/rust/sedona-raster-gdal/src/rs_clip.rs index e1cfb80334..f6467210c6 100644 --- a/rust/sedona-raster-gdal/src/rs_clip.rs +++ b/rust/sedona-raster-gdal/src/rs_clip.rs @@ -401,7 +401,15 @@ fn clip_raster( // Rasterize the geometry into a window-sized 0/1 mask: 1 inside, 0 outside // (moves `geometry`, whose only remaining use is the burn). - let mask = rasterize_geometry_mask(gdal, geometry, &geotransform, &window, all_touched)?; + let mut mask = Vec::new(); + rasterize_geometry_mask( + gdal, + geometry, + &geotransform, + &window, + all_touched, + &mut mask, + )?; // The envelope may overlap the raster while the geometry itself selects no // pixel (e.g. it falls between pixel centers); that is still the diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs index 5ae0c99e1c..71b49934bf 100644 --- a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -314,6 +314,7 @@ impl SedonaScalarKernel for RsZonalStats { let mut builder = Float64Builder::with_capacity(num_iterations); let mut scratch: Vec = Vec::new(); + let mut mask_scratch: Vec = Vec::new(); // The executor only sees (raster, roi); the option columns are advanced // in lockstep below. @@ -359,7 +360,14 @@ impl SedonaScalarKernel for RsZonalStats { "raster", engine, )?; - match compute_zonal_stats(gdal, raster, &geom_wkb, ¶ms, &mut scratch)? { + match compute_zonal_stats( + gdal, + raster, + &geom_wkb, + ¶ms, + &mut scratch, + &mut mask_scratch, + )? { Some(stats) => match stats.get(stat_type) { Some(value) => builder.append_value(value), None => builder.append_null(), @@ -485,6 +493,7 @@ impl SedonaScalarKernel for RsZonalStatsAll { let mut builder = StructBuilder::from_fields(zonal_stats_struct_fields(), num_iterations); let mut scratch: Vec = Vec::new(); + let mut mask_scratch: Vec = Vec::new(); let exec_arg_types = [arg_types[0].clone(), arg_types[1].clone()]; let exec_args = [args[0].clone(), args[1].clone()]; @@ -519,7 +528,14 @@ impl SedonaScalarKernel for RsZonalStatsAll { "raster", engine, )?; - match compute_zonal_stats(gdal, raster, &geom_wkb, ¶ms, &mut scratch)? { + match compute_zonal_stats( + gdal, + raster, + &geom_wkb, + ¶ms, + &mut scratch, + &mut mask_scratch, + )? { Some(stats) => append_struct_stats(&mut builder, &stats)?, None if params.lenient => append_struct_null(&mut builder)?, None => return no_intersection_err(), @@ -618,14 +634,16 @@ fn append_struct_stats(builder: &mut StructBuilder, stats: &ZonalStatistics) -> /// that intersects the footprint but whose selected pixels are all outside the /// geometry or all nodata yields `Ok(Some(..))` with `count = 0`. /// -/// `scratch` is a reused buffer for the selected pixel values so the per-row -/// collection does not allocate a fresh `Vec` each call. +/// `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 compute_zonal_stats( gdal: &Gdal, raster: &RasterRefImpl<'_>, geom_wkb: &[u8], params: &ZonalStatsParams, scratch: &mut Vec, + mask_scratch: &mut Vec, ) -> Result> { let num_bands = raster.num_bands(); let band_num = resolve_band(params.band, num_bands)?; @@ -671,8 +689,16 @@ fn compute_zonal_stats( }; // Rasterize the roi into a window-sized 0/1 mask (moves `geometry`, whose - // only remaining use is the burn). - let mask = rasterize_geometry_mask(gdal, geometry, &transform, &window, params.all_touched)?; + // 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 @@ -710,7 +736,13 @@ fn compute_zonal_stats( scratch.clear(); collect_masked_values( - band_bytes, data_type, width, &window, &mask, nodata, scratch, + band_bytes, + data_type, + width, + &window, + mask_scratch, + nodata, + scratch, ); Ok(Some(compute_statistics(scratch))) From 4b08dcee240944b85c2473d04cb8aaa3fec00c12 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Tue, 28 Jul 2026 14:58:46 -0700 Subject: [PATCH 14/16] refactor(rust/sedona-raster-gdal): streamline zonal-stats computation and output Address three RS_ZonalStats review comments, all behavior-preserving: - Mode: compute it from the already-sorted values in a single longest-run pass (tie -> larger) instead of a HashMap keyed on the value bit pattern. - Single statistic: split the shared masking/collection (collect_zonal_values) from the computation so RS_ZonalStats computes only the requested statistic (compute_single_statistic sorts only for median/mode), while RS_ZonalStatsAll still computes every statistic. - Struct output: assemble the RS_ZonalStatsAll struct from one typed builder per field plus an outer validity buffer (ZonalStatsBuilders), building the StructArray once at the end rather than downcasting a StructBuilder's field builders on every row. The sample (two-pass) variance, median, and NaN/empty semantics are unchanged, factored into shared helpers. A test pins that the single-stat path returns exactly what the full computation does. --- rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 384 ++++++++++++------ 1 file changed, 250 insertions(+), 134 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs index 71b49934bf..c179b06fd5 100644 --- a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -48,11 +48,11 @@ //! 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_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; @@ -154,22 +154,14 @@ struct ZonalStatistics { max: Option, } -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 { - 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, - } - } +/// 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, } // ============================================================================= @@ -360,7 +352,7 @@ impl SedonaScalarKernel for RsZonalStats { "raster", engine, )?; - match compute_zonal_stats( + match collect_zonal_values( gdal, raster, &geom_wkb, @@ -368,14 +360,17 @@ impl SedonaScalarKernel for RsZonalStats { &mut scratch, &mut mask_scratch, )? { - Some(stats) => match stats.get(stat_type) { - Some(value) => builder.append_value(value), - None => builder.append_null(), - }, + // 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. - None if params.lenient => builder.append_null(), - None => return no_intersection_err(), + RoiCoverage::NoIntersection if params.lenient => builder.append_null(), + RoiCoverage::NoIntersection => return no_intersection_err(), } Ok(()) }) @@ -491,7 +486,7 @@ impl SedonaScalarKernel for RsZonalStatsAll { 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 builders = ZonalStatsBuilders::with_capacity(num_iterations); let mut scratch: Vec = Vec::new(); let mut mask_scratch: Vec = Vec::new(); @@ -510,12 +505,12 @@ impl SedonaScalarKernel for RsZonalStatsAll { &mut exclude_no_data_iter, &mut lenient_iter, ) else { - append_struct_null(&mut builder)?; + builders.push_null(); return Ok(()); }; let (Some(raster), Some(wkb)) = (raster_opt, wkb_opt) else { - append_struct_null(&mut builder)?; + builders.push_null(); return Ok(()); }; @@ -528,7 +523,7 @@ impl SedonaScalarKernel for RsZonalStatsAll { "raster", engine, )?; - match compute_zonal_stats( + match collect_zonal_values( gdal, raster, &geom_wkb, @@ -536,15 +531,17 @@ impl SedonaScalarKernel for RsZonalStatsAll { &mut scratch, &mut mask_scratch, )? { - Some(stats) => append_struct_stats(&mut builder, &stats)?, - None if params.lenient => append_struct_null(&mut builder)?, - None => return no_intersection_err(), + 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(builder.finish()); + let out: ArrayRef = Arc::new(builders.finish()); RasterExecutor::finish_over(args, out) }) } @@ -571,80 +568,103 @@ fn zonal_stats_struct_fields() -> Fields { ]) } -/// 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::(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::(i) else { - return sedona_internal_err!("RS_ZonalStats: stat field {i} is not a Float64 builder"); - }; - field.append_null(); - } - builder.append(false); - Ok(()) +/// 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, } -/// Append one computed-stats row. The float fields carry through the `Option` -/// so an empty roi records `count = 0` with the rest NULL. -fn append_struct_stats(builder: &mut StructBuilder, stats: &ZonalStatistics) -> Result<()> { - let Some(count) = builder.field_builder::(0) else { - return sedona_internal_err!("RS_ZonalStats: count field is not an Int64 builder"); - }; - count.append_value(stats.count); - for (i, value) in [ - stats.sum, - stats.mean, - stats.median, - stats.mode, - stats.stddev, - stats.variance, - stats.min, - stats.max, - ] - .into_iter() - .enumerate() - { - let Some(field) = builder.field_builder::(i + 1) else { - return sedona_internal_err!( - "RS_ZonalStats: stat field {} is not a Float64 builder", - i + 1 - ); - }; - field.append_option(value); +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 = 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)) } - builder.append(true); - Ok(()) } // ============================================================================= // Core computation // ============================================================================= -/// Compute the statistics of the pixels a roi geometry selects on one band. +/// 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). /// -/// Returns `Ok(None)` when the roi geometry does not intersect the raster's -/// footprint. This is 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 `None` 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 yields `Ok(Some(..))` with `count = 0`. +/// 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 compute_zonal_stats( +fn collect_zonal_values( gdal: &Gdal, raster: &RasterRefImpl<'_>, geom_wkb: &[u8], params: &ZonalStatsParams, scratch: &mut Vec, mask_scratch: &mut Vec, -) -> Result> { +) -> Result { let num_bands = raster.num_bands(); let band_num = resolve_band(params.band, num_bands)?; @@ -673,7 +693,7 @@ fn compute_zonal_stats( // 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(None); + return Ok(RoiCoverage::NoIntersection); } // Parse the roi and clamp its envelope to the raster grid for the pixel @@ -685,7 +705,7 @@ fn compute_zonal_stats( .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(Some(compute_statistics(scratch))); + return Ok(RoiCoverage::Collected); }; // Rasterize the roi into a window-sized 0/1 mask (moves `geometry`, whose @@ -745,7 +765,7 @@ fn compute_zonal_stats( scratch, ); - Ok(Some(compute_statistics(scratch))) + Ok(RoiCoverage::Collected) } /// Resolve the 1-based band to use. `Some(b)` must be a valid 1-based index; @@ -830,17 +850,18 @@ fn collect_masked_values( } } -/// Compute every statistic from the selected pixel values. +/// 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 for the median 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. +/// 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); the caller owns it as reusable -/// scratch. +/// `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() { @@ -875,30 +896,17 @@ fn compute_statistics(values: &mut [f64]) -> ZonalStatistics { }; } - let n = values.len(); let sum: f64 = values.iter().sum(); - let mean = sum / n as f64; + 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 = if n > 1 { - let sum_sq: f64 = values.iter().map(|&v| (v - mean).powi(2)).sum(); - sum_sq / (n as f64 - 1.0) - } else { - 0.0 - }; + let variance = sample_variance(values); let stddev = variance.sqrt(); - let mode = compute_mode(values); - - // Median needs the values sorted; do it in place on the scratch buffer. - values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); - let mid = n / 2; - let median = if n.is_multiple_of(2) { - (values[mid - 1] + values[mid]) / 2.0 - } else { - values[mid] - }; + // 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, @@ -913,21 +921,97 @@ fn compute_statistics(values: &mut [f64]) -> ZonalStatistics { } } -/// The most frequent value, breaking ties toward the larger value (matching -/// Sedona Spark's `StatUtils.mode`, which returns the largest of the tied -/// modes). Values are keyed by their exact bit pattern, so integer-valued -/// pixels compare exactly. -fn compute_mode(values: &[f64]) -> f64 { - let mut counts: HashMap = HashMap::new(); - for &v in values { - *counts.entry(v.to_bits()).or_insert(0) += 1; +/// 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 { + 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); } - let best_count = counts.values().copied().max().unwrap_or(0); - counts - .into_iter() - .filter(|&(_, c)| c == best_count) - .map(|(bits, _)| f64::from_bits(bits)) - .fold(f64::NEG_INFINITY, f64::max) + 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::() / 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::() / n as f64; + let sum_sq: f64 = values.iter().map(|&v| (v - mean).powi(2)).sum(); + sum_sq / (n as f64 - 1.0) +} + +/// Sort `values` ascending in place. Callers exclude NaN beforehand, so the +/// `partial_cmp` fallback is never exercised. +fn sort_values(values: &mut [f64]) { + values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); +} + +/// The linear-interpolated 50th percentile of `sorted` (ascending): the middle +/// element for an odd count, the mean of the two central elements for an even +/// count. `sorted` must be non-empty. +fn median_of_sorted(sorted: &[f64]) -> f64 { + let mid = sorted.len() / 2; + if sorted.len().is_multiple_of(2) { + (sorted[mid - 1] + sorted[mid]) / 2.0 + } else { + sorted[mid] + } +} + +/// The most frequent value in `sorted` (ascending), breaking ties toward the +/// larger value (matching Sedona Spark's `StatUtils.mode`). Equal values are +/// adjacent once sorted, so a single pass over the runs finds the mode without a +/// map. `sorted` must be non-empty. +fn mode_of_sorted(sorted: &[f64]) -> f64 { + let mut best_val = sorted[0]; + let mut best_len = 1usize; + let mut run_len = 1usize; + for i in 1..sorted.len() { + run_len = if sorted[i] == sorted[i - 1] { + run_len + 1 + } else { + 1 + }; + // `>=` keeps the later — and, since ascending, larger — value on a tie. + if run_len >= best_len { + best_len = run_len; + best_val = sorted[i]; + } + } + best_val } // ============================================================================= @@ -1074,6 +1158,35 @@ mod tests { assert_eq!(s.stddev, Some(2.5_f64.sqrt())); } + #[test] + fn single_statistic_matches_full_computation() { + // RS_ZonalStats' single-stat path must return exactly what the full + // RS_ZonalStatsAll path computes, only without the others. A tie (two + // 4s) and an even count exercise mode and median. + let sample = [4.0, 1.0, 4.0, 2.0, 5.0, 3.0]; + let mut all = sample; + let full = compute_statistics(&mut all); + let cases = [ + (StatType::Count, Some(full.count as f64)), + (StatType::Sum, full.sum), + (StatType::Mean, full.mean), + (StatType::Median, full.median), + (StatType::Mode, full.mode), + (StatType::StdDev, full.stddev), + (StatType::Variance, full.variance), + (StatType::Min, full.min), + (StatType::Max, full.max), + ]; + for (stat, expected) in cases { + let mut work = sample; + assert_eq!( + compute_single_statistic(&mut work, stat), + expected, + "single statistic {stat:?} diverged from the full computation" + ); + } + } + #[test] fn statistics_empty_is_zero_count_and_nulls() { let mut values: Vec = vec![]; @@ -1085,10 +1198,13 @@ mod tests { assert_eq!(s.min, None); assert_eq!(s.max, None); assert_eq!(s.variance, None); - // A single-stat lookup returns 0 for count, NULL for the rest. - assert_eq!(s.get(StatType::Count), Some(0.0)); - assert_eq!(s.get(StatType::Sum), None); - assert_eq!(s.get(StatType::Mean), None); + // The single-stat path returns 0 for count, NULL for the rest. + assert_eq!( + compute_single_statistic(&mut values, StatType::Count), + Some(0.0) + ); + assert_eq!(compute_single_statistic(&mut values, StatType::Sum), None); + assert_eq!(compute_single_statistic(&mut values, StatType::Mean), None); } #[test] From a129b729b81b4bde536554c76da0e916a9251ed2 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Wed, 29 Jul 2026 12:13:18 -0700 Subject: [PATCH 15/16] test(python/sedonadb): pin RS_ZonalStats NaN/infinity handling against numpy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add rasterio+numpy parity coverage for the NaN and infinity edge cases Dewey's review flagged. A float64 fixture plants a NaN (then +inf) pixel inside the roi (not the nodata sentinel, so it is not excluded), and the resulting statistics are compared field-by-field against the numpy reference over the same masked selection: - NaN poisons every statistic except count (numpy semantics); - +inf flows through — sum/mean/max/mode become +inf, min/median stay finite, variance/stddev become NaN. A NaN/inf-aware comparator handles the fact that NaN never equals itself. This validates the existing behavior against a trusted reference; no kernel change. --- .../tests/functions/test_rs_zonalstats.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/python/sedonadb/tests/functions/test_rs_zonalstats.py b/python/sedonadb/tests/functions/test_rs_zonalstats.py index 70e037b992..d2f965a847 100644 --- a/python/sedonadb/tests/functions/test_rs_zonalstats.py +++ b/python/sedonadb/tests/functions/test_rs_zonalstats.py @@ -33,6 +33,8 @@ when it is unavailable rather than importing it at module scope. """ +import math + import numpy as np import pyarrow as pa import pytest @@ -91,6 +93,24 @@ def fixture_raster(tmp_path): return path, data[0] +def float_fixture_raster(tmp_path, *, planted_value, name): + """A single-band float64 raster with `planted_value` (NaN or inf) at an + interior pixel that GEOM_RECT selects, so the roi statistics must reckon with + it. The planted value is not the nodata sentinel, so nodata exclusion (on by + default) does not drop it. + + Returns `(path, band)` where `band` is the `(HEIGHT, WIDTH)` numpy array. + """ + data = random_raster_data( + "float64", bands=BANDS, height=HEIGHT, width=WIDTH, seed=7 + ) + # Pixel centre (105, 492.5) sits inside GEOM_RECT. + data[0][2, 2] = planted_value + path = tmp_path / name + write_geotiff(path, data, gdal_transform=GDAL_TRANSFORM, nodata=NODATA) + return path, data[0] + + def numpy_reference(band, wkt, *, all_touched, exclude_no_data): """Reference statistics over the pixels the roi selects, via rasterio+numpy. @@ -131,6 +151,21 @@ def numpy_reference(band, wkt, *, all_touched, exclude_no_data): } +def _assert_stat_equal(stat, got, expected): + """Compare one statistic against the numpy reference with NaN/inf-aware + equality: NaN never equals itself, so match it explicitly; inf compares + exactly; otherwise use exact equality for the integer-exact statistics and a + tolerance for the float-accumulating ones.""" + if math.isnan(expected): + assert got is not None and math.isnan(got), f"{stat}: expected NaN, got {got!r}" + elif math.isinf(expected): + assert got == expected, f"{stat}: expected {expected}, got {got!r}" + elif stat in EXACT_STATS: + assert got == expected, f"{stat}: {got!r} != {expected!r}" + else: + assert got == pytest.approx(expected), f"{stat}: {got!r} !~ {expected!r}" + + def _one_row(con, path, wkt): """A one-row frame with the raster and roi as columns, so the kernel runs its real per-row array path rather than constant-folding a scalar.""" @@ -202,6 +237,42 @@ def test_all_struct_matches_numpy(con, tmp_path): assert got[stat] == pytest.approx(expected[stat]) +def test_nan_pixel_poisons_every_statistic_like_numpy(con, tmp_path): + """A NaN pixel that is not the nodata sentinel poisons every statistic (numpy + semantics): count stays a real tally, everything else is NaN. Pinned against + rasterio+numpy over the same masked selection, so the NaN handling is + validated against a trusted reference rather than asserted on faith.""" + path, band = float_fixture_raster( + tmp_path, planted_value=float("nan"), name="zonal_nan.tif" + ) + with np.errstate(all="ignore"): + expected = numpy_reference( + band, GEOM_RECT, all_touched=False, exclude_no_data=True + ) + assert expected != "empty", "GEOM_RECT should select the planted pixel" + got = zonal_stats_all(con, path, GEOM_RECT, [1]) + for stat in STATS: + _assert_stat_equal(stat, got[stat], expected[stat]) + + +def test_infinity_pixel_matches_numpy(con, tmp_path): + """A +inf pixel flows through (not the nodata sentinel, not NaN): sum, mean, + max and mode go to +inf, min and median stay finite, and variance/stddev + become NaN (inf - inf). Pinned against rasterio+numpy over the same + selection.""" + path, band = float_fixture_raster( + tmp_path, planted_value=float("inf"), name="zonal_inf.tif" + ) + with np.errstate(all="ignore"): + expected = numpy_reference( + band, GEOM_RECT, all_touched=False, exclude_no_data=True + ) + assert expected != "empty", "GEOM_RECT should select the planted pixel" + got = zonal_stats_all(con, path, GEOM_RECT, [1]) + for stat in STATS: + _assert_stat_equal(stat, got[stat], expected[stat]) + + def test_exclude_no_data_default_and_disabled(con, tmp_path): path, band = fixture_raster(tmp_path) # Default excludes nodata; disabling it keeps those pixels, raising count. From 8aa9f092f910b99ed3535c0fab54886a8bcf5dd8 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Fri, 31 Jul 2026 10:06:37 -0700 Subject: [PATCH 16/16] perf(rust/sedona-raster-gdal): RS_ZonalStatsAll reuses its sort and mean for min/max/variance The all-statistics path takes min and max from the single sort it already performs and reuses the computed mean for the variance pass, replacing two folds and a redundant mean computation. The single-statistic path keeps its O(n) folds so the stats that need no ordering do not pay for a sort. --- rust/sedona-raster-gdal/src/rs_zonal_stats.rs | 37 +++++++++++-------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs index c179b06fd5..a82bb80285 100644 --- a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -860,8 +860,8 @@ fn collect_masked_values( /// 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. +/// `values` is sorted in place (for min, max, 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() { @@ -898,13 +898,14 @@ fn compute_statistics(values: &mut [f64]) -> ZonalStatistics { 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 variance = sample_variance(values, mean); let stddev = variance.sqrt(); - // Median and mode both read the values in sorted order; sort once in place. + // min, max, median, and mode all read the values in sorted order: sort once + // in place, then take the extremes from the ends instead of folding again. sort_values(values); + let min = values[0]; + let max = values[values.len() - 1]; let median = median_of_sorted(values); let mode = mode_of_sorted(values); @@ -943,11 +944,11 @@ fn compute_single_statistic(values: &mut [f64], stat: StatType) -> Option { // Count is handled before the value checks above. StatType::Count => unreachable!("count returns early"), StatType::Sum => values.iter().sum(), - StatType::Mean => values.iter().sum::() / values.len() as f64, + StatType::Mean => mean_of(values), 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::Variance => sample_variance(values, mean_of(values)), + StatType::StdDev => sample_variance(values, mean_of(values)).sqrt(), StatType::Median => { sort_values(values); median_of_sorted(values) @@ -959,16 +960,22 @@ fn compute_single_statistic(values: &mut [f64], stat: StatType) -> Option { }) } -/// 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 { +/// The arithmetic mean of `values`, which must be non-empty. +fn mean_of(values: &[f64]) -> f64 { + values.iter().sum::() / values.len() as f64 +} + +/// The sample (n-1) variance of `values` about their precomputed `mean`, +/// matching Sedona Spark; 0 for a single value. Taking the mean as an argument +/// lets the all-statistics path reuse the mean it already computed rather than +/// summing again; the squared-deviation form (rather than the naive +/// sum-of-squares) avoids catastrophic cancellation. `values` must be non-empty +/// and NaN-free. +fn sample_variance(values: &[f64], mean: f64) -> f64 { let n = values.len(); if n <= 1 { return 0.0; } - let mean = values.iter().sum::() / n as f64; let sum_sq: f64 = values.iter().map(|&v| (v - mean).powi(2)).sum(); sum_sq / (n as f64 - 1.0) }