diff --git a/c/sedona-gdal/src/gdal.rs b/c/sedona-gdal/src/gdal.rs index 481f9d273..fe2a35f86 100644 --- a/c/sedona-gdal/src/gdal.rs +++ b/c/sedona-gdal/src/gdal.rs @@ -29,7 +29,7 @@ use crate::errors::Result; use crate::gdal_api::GdalApi; use crate::gdal_dyn_bindgen::{GDALOpenFlags, OGRFieldType}; use crate::mem::create_mem_dataset; -use crate::raster::polygonize::{polygonize, PolygonizeOptions}; +use crate::raster::polygonize::{fpolygonize, polygonize, PolygonizeOptions}; use crate::raster::rasterband::RasterBand; use crate::raster::rasterize::{rasterize, RasterizeOptions}; use crate::raster::rasterize_affine::rasterize_affine; @@ -250,4 +250,24 @@ impl Gdal { options, ) } + + /// Polygonize a raster band into a vector layer using float pixels. + /// See also [`fpolygonize`]. + pub fn fpolygonize( + &self, + src_band: &RasterBand<'_>, + mask_band: Option<&RasterBand<'_>>, + out_layer: &Layer<'_>, + pixel_value_field: i32, + options: &PolygonizeOptions, + ) -> Result<()> { + fpolygonize( + self.api, + src_band, + mask_band, + out_layer, + pixel_value_field, + options, + ) + } } diff --git a/c/sedona-gdal/src/raster/polygonize.rs b/c/sedona-gdal/src/raster/polygonize.rs index 60f067ced..35d96d56d 100644 --- a/c/sedona-gdal/src/raster/polygonize.rs +++ b/c/sedona-gdal/src/raster/polygonize.rs @@ -100,6 +100,38 @@ pub fn polygonize( Ok(()) } +/// Polygonize a raster band into a vector layer. +/// This uses `GDALFPolygonize`, which reads source pixels as floats. +pub fn fpolygonize( + api: &'static GdalApi, + src_band: &RasterBand<'_>, + mask_band: Option<&RasterBand<'_>>, + out_layer: &Layer<'_>, + pixel_value_field: i32, + options: &PolygonizeOptions, +) -> Result<()> { + let mask = mask_band.map_or(ptr::null_mut(), |b| b.c_rasterband()); + let csl = options.to_options_list()?; + + let rv = unsafe { + call_gdal_api!( + api, + GDALFPolygonize, + src_band.c_rasterband(), + mask, + out_layer.c_layer(), + pixel_value_field, + csl.as_ptr(), + ptr::null_mut(), // pfnProgress + ptr::null_mut() // pProgressData + ) + }; + if rv != CE_None { + return Err(api.last_cpl_err(rv as u32)); + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/docs/reference/sql/rs_polygonize.qmd b/docs/reference/sql/rs_polygonize.qmd new file mode 100644 index 000000000..4943283c8 --- /dev/null +++ b/docs/reference/sql/rs_polygonize.qmd @@ -0,0 +1,43 @@ +--- +# 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_Polygonize +description: Converts a raster band into polygons grouped by connected pixel value regions. +kernels: + - returns: list + args: + - raster + - name: band + type: integer +--- + +## Description + +Converts the specified raster `band` into vector polygons, returning one item per +connected region of pixels that shares the same value. + +Each returned list item contains: + +- `geom`: the polygon geometry encoded as WKB +- `value`: the pixel value represented by that polygon + +## Examples + +```sql +SELECT RS_Polygonize(RS_Example(), 1); +``` diff --git a/rust/sedona-raster-gdal/Cargo.toml b/rust/sedona-raster-gdal/Cargo.toml index c014d6761..db09a027b 100644 --- a/rust/sedona-raster-gdal/Cargo.toml +++ b/rust/sedona-raster-gdal/Cargo.toml @@ -63,3 +63,8 @@ path = "benches/rs_frompath.rs" harness = false name = "rs_metadata" path = "benches/rs_metadata.rs" + +[[bench]] +harness = false +name = "rs_polygonize" +path = "benches/rs_polygonize.rs" diff --git a/rust/sedona-raster-gdal/benches/rs_polygonize.rs b/rust/sedona-raster-gdal/benches/rs_polygonize.rs new file mode 100644 index 000000000..d3a0ec96c --- /dev/null +++ b/rust/sedona-raster-gdal/benches/rs_polygonize.rs @@ -0,0 +1,91 @@ +// 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 RS_Polygonize UDF. + +use std::{hint::black_box, sync::Arc}; + +use arrow_array::{ArrayRef, Int32Array}; +use arrow_schema::DataType; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use datafusion_expr::ColumnarValue; +use datafusion_expr::ScalarUDF; +use sedona_schema::datatypes::{SedonaType, RASTER}; +use sedona_testing::testers::ScalarUdfTester; + +fn raster_array(rows: usize) -> ArrayRef { + assert!(rows > 0, "benchmark rows must be positive"); + + let example: ScalarUDF = sedona_raster_functions::rs_example::rs_example_udf().into(); + let tester = ScalarUdfTester::new(example, vec![]); + let scalar = tester.invoke(vec![]).unwrap(); + match scalar { + ColumnarValue::Scalar(value) => value.to_array_of_size(rows).unwrap(), + ColumnarValue::Array(array) => array, + } +} + +fn band_array(rows: usize) -> ArrayRef { + Arc::new(Int32Array::from(vec![1; rows])) +} + +fn bench_rs_polygonize(c: &mut Criterion) { + let udf: ScalarUDF = sedona_raster_gdal::rs_polygonize_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); + + let single = raster_array(1); + let batched = raster_array(32); + let single_band = band_array(single.len()); + let batched_band = band_array(batched.len()); + + let mut group = c.benchmark_group("rs_polygonize"); + + group.throughput(Throughput::Elements(single.len() as u64)); + group.bench_with_input( + BenchmarkId::new("fixtures", "single_example"), + &single, + |b, input| { + b.iter(|| { + black_box( + tester + .invoke_arrays(vec![input.clone(), single_band.clone()]) + .unwrap(), + ) + }) + }, + ); + + group.throughput(Throughput::Elements(batched.len() as u64)); + group.bench_with_input( + BenchmarkId::new("fixtures", "batched_example"), + &batched, + |b, input| { + b.iter(|| { + black_box( + tester + .invoke_arrays(vec![input.clone(), batched_band.clone()]) + .unwrap(), + ) + }) + }, + ); + + group.finish(); +} + +criterion_group!(benches, bench_rs_polygonize); +criterion_main!(benches); diff --git a/rust/sedona-raster-gdal/src/lib.rs b/rust/sedona-raster-gdal/src/lib.rs index 6682d41f6..f878f34a1 100644 --- a/rust/sedona-raster-gdal/src/lib.rs +++ b/rust/sedona-raster-gdal/src/lib.rs @@ -35,6 +35,7 @@ mod gdal_dataset_provider; mod raster_loader; mod rs_frompath; mod rs_metadata; +mod rs_polygonize; mod source_uri; mod utils; @@ -46,6 +47,7 @@ pub use gdal_common::{ pub use raster_loader::{GdalLoader, GDAL_FORMAT}; pub use rs_frompath::rs_frompath_udf; pub use rs_metadata::rs_metadata_udf; +pub use rs_polygonize::rs_polygonize_udf; pub use utils::{ append_as_indb_raster, append_as_outdb_raster, append_nd_from_dataset, dataset_to_indb_raster, gdal_dataset_to_nd_raster, diff --git a/rust/sedona-raster-gdal/src/register.rs b/rust/sedona-raster-gdal/src/register.rs index f4e792268..52f7dfd48 100644 --- a/rust/sedona-raster-gdal/src/register.rs +++ b/rust/sedona-raster-gdal/src/register.rs @@ -22,5 +22,6 @@ pub fn default_function_set() -> FunctionSet { let mut function_set = FunctionSet::new(); 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 } diff --git a/rust/sedona-raster-gdal/src/rs_polygonize.rs b/rust/sedona-raster-gdal/src/rs_polygonize.rs new file mode 100644 index 000000000..1c09c272d --- /dev/null +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -0,0 +1,492 @@ +// 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_Polygonize UDF - Convert a raster band to vector polygons. + +use std::sync::Arc; + +use arrow_array::builder::{ + BinaryBuilder, Float64Builder, NullBufferBuilder, OffsetBufferBuilder, StringViewBuilder, +}; +use arrow_array::{ArrayRef, Int32Array, ListArray, StructArray}; +use arrow_schema::{DataType, Field, Fields}; +use datafusion_common::cast::as_int32_array; +use datafusion_common::config::ConfigOptions; +use datafusion_common::error::Result; +use datafusion_common::{exec_datafusion_err, exec_err}; +use datafusion_expr::{ColumnarValue, Volatility}; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_gdal::driver::Driver; +use sedona_gdal::gdal::Gdal; +use sedona_gdal::gdal_dyn_bindgen::{OGRFieldType, OGRwkbGeometryType}; +use sedona_gdal::raster::polygonize::PolygonizeOptions; +use sedona_gdal::raster::types::GdalDataType; +use sedona_raster::traits::RasterRef; +use sedona_raster_functions::RasterExecutor; +use sedona_schema::datatypes::{SedonaType, WKB_GEOMETRY_ITEM_CRS}; +use sedona_schema::matchers::ArgMatcher; + +use crate::gdal_common::with_gdal; +use crate::gdal_dataset_provider::{ + configure_thread_local_options, thread_local_provider, RasterDataset, +}; + +pub fn rs_polygonize_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_polygonize", + vec![Arc::new(RsPolygonize)], + Volatility::Immutable, + ) +} + +#[derive(Debug)] +struct RsPolygonize; + +impl SedonaScalarKernel for RsPolygonize { + fn return_type(&self, args: &[SedonaType]) -> Result> { + ArgMatcher::new( + vec![ArgMatcher::is_raster(), ArgMatcher::is_integer()], + SedonaType::Arrow(polygon_value_list_type()), + ) + .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 executor = RasterExecutor::new(arg_types, args); + let num_iterations = executor.num_iterations(); + let band_array = band_argument_array(&args[1], num_iterations)?; + let mut band_iter = band_array.iter(); + + let mut list_offsets = OffsetBufferBuilder::::new(num_iterations); + let mut geom_item_builder = BinaryBuilder::new(); + let mut geom_crs_builder = StringViewBuilder::new(); + let mut value_builder = Float64Builder::new(); + let mut valid_list_items = NullBufferBuilder::new(num_iterations); + + with_gdal(|gdal| { + configure_thread_local_options(gdal, config_options)?; + let provider = thread_local_provider(gdal) + .map_err(|e| exec_datafusion_err!("Failed to init GDAL provider: {e}"))?; + // Note: We deliberately use "Memory" instead of "MEM" to be compatible with older GDAL versions. + let mem_driver = gdal + .get_driver_by_name("Memory") + .map_err(|e| exec_datafusion_err!("Failed to get Memory driver: {e}"))?; + + executor.execute_raster_void(|_, raster_opt| { + let band_opt = band_iter.next().expect("band iteration should match rows"); + + let raster = match (raster_opt, band_opt) { + (Some(raster), Some(_)) => raster, + _ => { + valid_list_items.append_null(); + list_offsets.push_length(0); + return Ok(()); + } + }; + + let band_index = band_opt.unwrap(); + let bands = raster.bands(); + if band_index <= 0 || band_index as usize > bands.len() { + return exec_err!("Band {} is out of range (1-{})", band_index, bands.len()); + } + let band_num = band_index as usize; + + let raster_ds = provider + .raster_ref_to_gdal(raster) + .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {e}"))?; + + let crs_str = raster.crs(); + let mut num_polygons = 0; + polygonize_raster(gdal, &mem_driver, raster_ds, band_num, |wkb, value| { + geom_item_builder.append_value(wkb); + geom_crs_builder.append_option(crs_str); + value_builder.append_value(value); + num_polygons += 1; + Ok(()) + })?; + + valid_list_items.append_non_null(); + list_offsets.push_length(num_polygons); + Ok(()) + })?; + + let list_array = assemble_polygon_values_list( + geom_item_builder, + geom_crs_builder, + value_builder, + list_offsets, + valid_list_items, + ); + + let result: ArrayRef = Arc::new(list_array); + executor.finish(result) + }) + } +} + +fn band_argument_array(arg: &ColumnarValue, num_iterations: usize) -> Result { + let array = arg + .clone() + .cast_to(&DataType::Int32, None)? + .into_array(num_iterations)?; + Ok(as_int32_array(&array)?.clone()) +} + +fn polygon_value_list_type() -> DataType { + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(polygon_value_struct_fields()), + true, + ))) +} + +fn polygon_value_struct_fields() -> Fields { + Fields::from(vec![ + WKB_GEOMETRY_ITEM_CRS + .to_storage_field("geom", false) + .unwrap(), + Field::new("value", DataType::Float64, false), + ]) +} + +fn assemble_polygon_values_list( + mut geom_item_builder: BinaryBuilder, + mut geom_crs_builder: StringViewBuilder, + mut value_builder: Float64Builder, + list_offsets: OffsetBufferBuilder, + mut valid_list_items: NullBufferBuilder, +) -> ListArray { + let item_array = Arc::new(geom_item_builder.finish()) as ArrayRef; + let crs_array = Arc::new(geom_crs_builder.finish()) as ArrayRef; + + let geom_field = WKB_GEOMETRY_ITEM_CRS + .to_storage_field("geom", false) + .unwrap(); + let geom_fields = match geom_field.data_type() { + DataType::Struct(fields) => fields.clone(), + _ => unreachable!(), + }; + + let geom_struct = StructArray::new(geom_fields, vec![item_array, crs_array], None); + + let value_array = Arc::new(value_builder.finish()) as ArrayRef; + let element_struct = StructArray::new( + polygon_value_struct_fields(), + vec![Arc::new(geom_struct) as ArrayRef, value_array], + None, + ); + + let list_field = Arc::new(Field::new( + "item", + DataType::Struct(polygon_value_struct_fields()), + true, + )); + + ListArray::new( + list_field, + list_offsets.finish(), + Arc::new(element_struct), + valid_list_items.finish(), + ) +} + +/// Perform GDAL polygonize on the given raster dataset and band, executing a callback closure for each polygonized output. +fn polygonize_raster( + gdal: &Gdal, + mem_driver: &Driver, + raster_ds: RasterDataset<'_>, + band_num: usize, + mut callback: F, +) -> Result<()> +where + F: FnMut(&[u8], f64) -> Result<()>, +{ + let gdal_dataset = raster_ds.as_dataset(); + + let raster_band = gdal_dataset + .rasterband(band_num) + .map_err(|e| exec_datafusion_err!("Failed to get band {}: {}", band_num, e))?; + + // Create a memory dataset containing one vector layer to hold the polygonized output. + let vector_ds = mem_driver + .create_vector_only("") + .map_err(|e| exec_datafusion_err!("Failed to create vector dataset: {e}"))?; + + let spatial_ref = gdal_dataset.spatial_ref().ok(); + let mut layer = vector_ds + .create_layer(sedona_gdal::dataset::LayerOptions { + name: "polygons", + srs: spatial_ref.as_ref(), + ty: OGRwkbGeometryType::wkbPolygon, + options: None, + }) + .map_err(|e| exec_datafusion_err!("Failed to create layer: {e}"))?; + + let field_defn = gdal + .create_field_defn("value", OGRFieldType::OFTReal) + .map_err(|e| exec_datafusion_err!("Failed to create field definition: {e}"))?; + layer + .create_field(&field_defn) + .map_err(|e| exec_datafusion_err!("Failed to add field to layer: {e}"))?; + + // Polygonize the raster band into the vector layer, using the "value" field to store pixel values. + let band_type = raster_band.band_type(); + let is_float = matches!(band_type, GdalDataType::Float32 | GdalDataType::Float64); + + if is_float { + gdal.fpolygonize(&raster_band, None, &layer, 0, &PolygonizeOptions::default()) + .map_err(|e| exec_datafusion_err!("GDAL fpolygonize failed: {e}"))?; + } else { + gdal.polygonize(&raster_band, None, &layer, 0, &PolygonizeOptions::default()) + .map_err(|e| exec_datafusion_err!("GDAL polygonize failed: {e}"))?; + } + + // Extract the WKB geometry and value for each feature in the layer. + let mut value_field_idx: Option = None; + for feature in layer.features() { + let geom = feature + .geometry() + .ok_or_else(|| exec_datafusion_err!("Polygonize output feature missing geometry"))?; + let wkb = geom + .wkb() + .map_err(|e| exec_datafusion_err!("Failed to export geometry to WKB: {e}"))?; + + let idx = match value_field_idx { + Some(idx) => idx, + None => { + let idx = feature + .field_index("value") + .map_err(|e| exec_datafusion_err!("Missing 'value' field: {e}"))?; + value_field_idx = Some(idx); + idx + } + }; + + let value = feature.field_as_double(idx); + + callback(&wkb, value)?; + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow_array::Array; + use datafusion_common::cast::{ + as_float64_array, as_list_array, as_string_view_array, as_struct_array, + }; + use datafusion_common::ScalarValue; + use datafusion_expr::{ScalarUDF, ScalarUDFImpl}; + use sedona_gdal::raster::types::Buffer; + use sedona_raster::array::RasterStructArray; + use sedona_schema::datatypes::RASTER; + use sedona_testing::raster_spec::{raster_array, RasterSpec}; + use sedona_testing::testers::ScalarUdfTester; + use tempfile::tempdir; + + fn test_raster_spec() -> RasterSpec { + RasterSpec::d2(3, 3) + .transform([0.0, 1.0, 0.0, 3.0, 0.0, -1.0]) + .band_values(&[1u8, 1, 0, 1, 2, 2, 0, 2, 2]) + .nodata(255u8) + } + + fn test_float_raster_spec() -> RasterSpec { + RasterSpec::d2(3, 3) + .transform([0.0, 1.0, 0.0, 3.0, 0.0, -1.0]) + .band_values(&[1.5f64, 1.5, 0.0, 1.5, 2.7, 2.7, 0.0, 2.7, 2.7]) + .nodata(255.0f64) + } + + fn build_polygonize_test_raster() -> arrow_array::StructArray { + let tmp = tempdir().unwrap(); + let path = tmp.path().join("polygonize.tif"); + let path_str = path.to_string_lossy().to_string(); + + with_gdal(|gdal| { + let driver = gdal.get_driver_by_name("GTiff").unwrap(); + let dataset = driver + .create_with_band_type::(&path_str, 3, 3, 1) + .unwrap(); + dataset + .set_geo_transform(&[0.0, 1.0, 0.0, 3.0, 0.0, -1.0]) + .unwrap(); + let band = dataset.rasterband(1).unwrap(); + band.set_no_data_value(Some(255.0)).unwrap(); + let mut buffer = Buffer::new((3, 3), vec![1u8, 1, 0, 1, 2, 2, 0, 2, 2]); + band.write((0, 0), (3, 3), &mut buffer).unwrap(); + crate::utils::dataset_to_indb_raster(&dataset) + }) + .unwrap() + } + + #[test] + fn test_rs_polygonize_udf_name() { + assert_eq!(rs_polygonize_udf().name(), "rs_polygonize"); + } + + #[test] + fn test_polygonize_raster() { + let mut results = Vec::new(); + with_gdal(|gdal| { + let provider = thread_local_provider(gdal).unwrap(); + let mem_driver = gdal.get_driver_by_name("Memory").unwrap(); + let raster_array = build_polygonize_test_raster(); + let raster_struct = RasterStructArray::try_new(&raster_array).unwrap(); + let raster = raster_struct.get(0).unwrap(); + let raster_ds = provider.raster_ref_to_gdal(&raster).unwrap(); + polygonize_raster(gdal, &mem_driver, raster_ds, 1, |wkb, value| { + results.push((wkb.to_vec(), value)); + Ok(()) + }) + }) + .unwrap(); + + assert!( + !results.is_empty(), + "Polygonize should return at least one polygon" + ); + for (wkb, value) in &results { + assert!(wkb.len() >= 5, "WKB should be at least 5 bytes"); + assert!(value.is_finite(), "Value should be a finite number"); + } + } + + #[test] + fn test_polygonize_kernel_return_type() { + let kernel = RsPolygonize; + let arg_types = vec![RASTER, SedonaType::Arrow(DataType::Int32)]; + let return_type = kernel.return_type(&arg_types).unwrap(); + assert!(matches!( + return_type, + Some(SedonaType::Arrow(DataType::List(_))) + )); + } + + #[test] + fn test_polygonize_invoke_scalar() { + let udf: ScalarUDF = rs_polygonize_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); + + let result = tester + .invoke_scalar_scalar(test_raster_spec(), ScalarValue::Int32(Some(1))) + .unwrap(); + + let array = result.to_array_of_size(1).unwrap(); + let list_array = as_list_array(&array).unwrap(); + let binding = list_array.value(0); + let struct_array = as_struct_array(&binding).unwrap(); + let geom_array = as_struct_array(struct_array.column(0)).unwrap(); + let crs_array = as_string_view_array(geom_array.column(1)).unwrap(); + + assert!(!crs_array.is_empty()); + for i in 0..crs_array.len() { + assert_eq!(crs_array.value(i), "OGC:CRS84"); + } + } + + #[test] + fn test_polygonize_invoke_array_and_nulls() { + let udf: ScalarUDF = rs_polygonize_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); + + let rasters = vec![ + Some(test_raster_spec()), + None, + Some(test_float_raster_spec()), + ]; + let raster_array = Arc::new(raster_array(rasters)); + let band_array = Arc::new(Int32Array::from(vec![Some(1), Some(1), Some(1)])); + + let result = tester + .invoke_arrays(vec![raster_array, band_array]) + .unwrap(); + let list_array = as_list_array(&result).unwrap(); + + assert_eq!(list_array.len(), 3); + + // Row 0: non-null, non-empty list of polygons for integer raster + assert!(!list_array.is_null(0)); + let row0 = list_array.value(0); + let row0 = as_struct_array(&row0).unwrap(); + assert!(!row0.is_empty()); + let geom0 = as_struct_array(row0.column(0)).unwrap(); + let crs0 = as_string_view_array(geom0.column(1)).unwrap(); + for i in 0..crs0.len() { + assert_eq!(crs0.value(i), "OGC:CRS84"); + } + let values0 = as_float64_array(row0.column(1)).unwrap(); + let mut sorted_vals0 = (0..values0.len()) + .map(|i| values0.value(i)) + .collect::>(); + sorted_vals0.sort_by(|a, b| a.partial_cmp(b).unwrap()); + assert_eq!(sorted_vals0, vec![0.0, 0.0, 1.0, 2.0]); + + // Row 1: null list (None) + assert!(list_array.is_null(1)); + + // Row 2: non-null, non-empty list of polygons for float raster (preserving floating point precision) + assert!(!list_array.is_null(2)); + let row2 = list_array.value(2); + let row2 = as_struct_array(&row2).unwrap(); + assert!(!row2.is_empty()); + let geom2 = as_struct_array(row2.column(0)).unwrap(); + let crs2 = as_string_view_array(geom2.column(1)).unwrap(); + for i in 0..crs2.len() { + assert_eq!(crs2.value(i), "OGC:CRS84"); + } + let values2 = as_float64_array(row2.column(1)).unwrap(); + let mut sorted_vals2 = (0..values2.len()) + .map(|i| values2.value(i)) + .collect::>(); + sorted_vals2.sort_by(|a, b| a.partial_cmp(b).unwrap()); + assert_eq!(sorted_vals2, vec![0.0, 0.0, 1.5, 2.7f32 as f64]); + } + + #[test] + fn test_polygonize_invalid_band_errors() { + let udf: ScalarUDF = rs_polygonize_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); + + let err = tester + .invoke_raster_array_scalar( + vec![Some(test_raster_spec())], + ScalarValue::Int32(Some(99)), + ) + .expect_err("out-of-range band should error"); + + assert!(err.to_string().contains("Band 99 is out of range")); + } +}