From 1de207735d1f580b577a091a8caefffdd2518c0a Mon Sep 17 00:00:00 2001 From: kontinuation Date: Sun, 14 Jun 2026 16:18:05 +0800 Subject: [PATCH 1/9] feat(rust/sedona-raster-gdal): add RS_Polygonize --- docs/reference/sql/rs_polygonize.qmd | 43 ++ rust/sedona-raster-gdal/Cargo.toml | 5 + .../benches/rs_polygonize.rs | 91 +++++ rust/sedona-raster-gdal/src/lib.rs | 2 + rust/sedona-raster-gdal/src/register.rs | 1 + rust/sedona-raster-gdal/src/rs_polygonize.rs | 378 ++++++++++++++++++ 6 files changed, 520 insertions(+) create mode 100644 docs/reference/sql/rs_polygonize.qmd create mode 100644 rust/sedona-raster-gdal/benches/rs_polygonize.rs create mode 100644 rust/sedona-raster-gdal/src/rs_polygonize.rs 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..71431611a --- /dev/null +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -0,0 +1,378 @@ +// 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::convert::TryInto; +use std::sync::Arc; + +use arrow_array::builder::{BinaryBuilder, Float64Builder, ListBuilder, StructBuilder}; +use arrow_array::{ArrayRef, Int32Array}; +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::gdal::Gdal; +use sedona_gdal::gdal_dyn_bindgen::{OGRFieldType, OGRwkbGeometryType}; +use sedona_gdal::mem::MemDatasetBuilder; +use sedona_raster::array::RasterRefImpl; +use sedona_raster::traits::RasterRef; +use sedona_raster_functions::RasterExecutor; +use sedona_schema::datatypes::SedonaType; +use sedona_schema::matchers::ArgMatcher; + +use crate::gdal_common::with_gdal; +use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider}; + +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_builder = ListBuilder::new(StructBuilder::from_fields( + polygon_value_struct_fields(), + 16, + )); + + with_gdal(|gdal| { + configure_thread_local_options(gdal, config_options)?; + 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, + _ => { + list_builder.append_null(); + return Ok(()); + } + }; + + let band_num: usize = band_opt.unwrap().max(1).try_into().unwrap_or(1); + + let polygon_values = polygonize_raster(gdal, raster, band_num)?; + let struct_builder = list_builder.values(); + for (wkb, value) in polygon_values { + let geom_builder = struct_builder + .field_builder::(0) + .expect("Expected BinaryBuilder for geom field"); + geom_builder.append_value(&wkb); + + let value_builder = struct_builder + .field_builder::(1) + .expect("Expected Float64Builder for value field"); + value_builder.append_value(value); + + struct_builder.append(true); + } + list_builder.append(true); + Ok(()) + })?; + + let result: ArrayRef = Arc::new(list_builder.finish()); + 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![ + Field::new("geom", DataType::Binary, false), + Field::new("value", DataType::Float64, false), + ]) +} + +fn polygonize_raster( + gdal: &Gdal, + raster: &RasterRefImpl<'_>, + band_num: usize, +) -> Result, f64)>> { + let bands = raster.bands(); + if band_num == 0 || band_num > bands.len() { + return exec_err!("Band {} is out of range (1-{})", band_num, bands.len()); + } + + let provider = thread_local_provider(gdal) + .map_err(|e| exec_datafusion_err!("Failed to init GDAL provider: {e}"))?; + let raster_ds = provider + .raster_ref_to_gdal(raster) + .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {e}"))?; + 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))?; + + let mem_ds = unsafe { + MemDatasetBuilder::new(0, 0) + .build(gdal) + .map_err(|e| exec_datafusion_err!("Failed to create memory dataset: {e}"))? + }; + + let spatial_ref = gdal_dataset.spatial_ref().ok(); + let mut layer = mem_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}"))?; + + gdal.polygonize( + &raster_band, + None, + &layer, + 0, + &sedona_gdal::raster::polygonize::PolygonizeOptions::default(), + ) + .map_err(|e| exec_datafusion_err!("GDAL polygonize failed: {e}"))?; + + let mut polygon_values = Vec::new(); + 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 + } + }; + + polygon_values.push((wkb, feature.field_as_double(idx))); + } + + Ok(polygon_values) +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow_array::Array; + use arrow_array::StructArray; + use datafusion_common::cast::{as_list_array, as_struct_array}; + use datafusion_common::ScalarValue; + use datafusion_expr::{ScalarUDF, ScalarUDFImpl}; + use sedona_raster::array::RasterStructArray; + use sedona_raster_functions::rs_example::rs_example_udf; + use sedona_schema::datatypes::RASTER; + use sedona_testing::testers::ScalarUdfTester; + + fn load_test_raster() -> StructArray { + let udf: ScalarUDF = rs_example_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![]); + let result = tester.invoke(vec![]).unwrap(); + match result { + ColumnarValue::Scalar(ScalarValue::Struct(array)) => array.as_ref().clone(), + other => panic!("Expected raster struct scalar, got {other:?}"), + } + } + + #[test] + fn test_rs_polygonize_udf_name() { + assert_eq!(rs_polygonize_udf().name(), "rs_polygonize"); + } + + #[test] + fn test_polygon_value_list_type() { + let dt = polygon_value_list_type(); + match dt { + DataType::List(field) => { + assert_eq!(field.name(), "item"); + match field.data_type() { + DataType::Struct(fields) => { + assert_eq!(fields.len(), 2); + assert_eq!(fields[0].name(), "geom"); + assert_eq!(fields[1].name(), "value"); + } + other => panic!("Expected Struct data type, got {other:?}"), + } + } + other => panic!("Expected List data type, got {other:?}"), + } + } + + #[test] + fn test_polygonize_raster() { + let result = with_gdal(|gdal| { + let raster_array = load_test_raster(); + let raster_struct = RasterStructArray::new(&raster_array); + let raster = raster_struct.get(0).unwrap(); + polygonize_raster(gdal, &raster, 1) + }) + .unwrap(); + + assert!( + !result.is_empty(), + "Polygonize should return at least one polygon" + ); + for (wkb, value) in &result { + 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 raster_array = load_test_raster(); + let kernel = RsPolygonize; + let arg_types = vec![RASTER, SedonaType::Arrow(DataType::Int32)]; + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(raster_array))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ]; + + let result = kernel + .invoke_batch_from_args( + &arg_types, + &args, + &SedonaType::Arrow(DataType::Null), + 0, + None, + ) + .unwrap(); + + match result { + ColumnarValue::Scalar(ScalarValue::List(list)) => { + assert!(!list.is_empty()); + } + other => panic!("Expected scalar list result, got {other:?}"), + } + } + + #[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 raster_array = ScalarValue::Struct(Arc::new(load_test_raster())) + .to_array_of_size(2) + .unwrap(); + let band_array = Arc::new(Int32Array::from(vec![Some(1), None])); + let result = tester + .invoke_arrays(vec![raster_array, band_array]) + .unwrap(); + let list_array = as_list_array(&result).unwrap(); + + assert_eq!(list_array.len(), 2); + assert!(!list_array.is_null(0)); + assert!(list_array.is_null(1)); + + let row0 = list_array.value(0); + let row0 = as_struct_array(&row0).unwrap(); + assert!(!row0.is_empty()); + } + + #[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 raster_array = Arc::new(load_test_raster()); + let err = tester + .invoke_arrays(vec![ + raster_array, + Arc::new(Int32Array::from(vec![Some(99)])), + ]) + .expect_err("out-of-range band should error"); + + assert!(err.to_string().contains("Band 99 is out of range")); + } +} From ba6e4edd0a0bb553a034953de5cd053055a5b84e Mon Sep 17 00:00:00 2001 From: kontinuation Date: Tue, 16 Jun 2026 21:11:18 +0800 Subject: [PATCH 2/9] test(rust/sedona-raster-gdal): simplify RS_Polygonize fixtures --- rust/sedona-raster-gdal/src/rs_polygonize.rs | 61 ++++++++------------ 1 file changed, 24 insertions(+), 37 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_polygonize.rs b/rust/sedona-raster-gdal/src/rs_polygonize.rs index 71431611a..a1e512714 100644 --- a/rust/sedona-raster-gdal/src/rs_polygonize.rs +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -235,23 +235,19 @@ mod tests { use super::*; use arrow_array::Array; - use arrow_array::StructArray; use datafusion_common::cast::{as_list_array, as_struct_array}; use datafusion_common::ScalarValue; use datafusion_expr::{ScalarUDF, ScalarUDFImpl}; use sedona_raster::array::RasterStructArray; - use sedona_raster_functions::rs_example::rs_example_udf; use sedona_schema::datatypes::RASTER; + use sedona_testing::raster_spec::RasterSpec; use sedona_testing::testers::ScalarUdfTester; - fn load_test_raster() -> StructArray { - let udf: ScalarUDF = rs_example_udf().into(); - let tester = ScalarUdfTester::new(udf, vec![]); - let result = tester.invoke(vec![]).unwrap(); - match result { - ColumnarValue::Scalar(ScalarValue::Struct(array)) => array.as_ref().clone(), - other => panic!("Expected raster struct scalar, got {other:?}"), - } + 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) } #[test] @@ -281,7 +277,7 @@ mod tests { #[test] fn test_polygonize_raster() { let result = with_gdal(|gdal| { - let raster_array = load_test_raster(); + let raster_array = test_raster_spec().build(); let raster_struct = RasterStructArray::new(&raster_array); let raster = raster_struct.get(0).unwrap(); polygonize_raster(gdal, &raster, 1) @@ -311,26 +307,15 @@ mod tests { #[test] fn test_polygonize_invoke_scalar() { - let raster_array = load_test_raster(); - let kernel = RsPolygonize; - let arg_types = vec![RASTER, SedonaType::Arrow(DataType::Int32)]; - let args = vec![ - ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(raster_array))), - ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), - ]; - - let result = kernel - .invoke_batch_from_args( - &arg_types, - &args, - &SedonaType::Arrow(DataType::Null), - 0, - None, - ) + 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(); match result { - ColumnarValue::Scalar(ScalarValue::List(list)) => { + ScalarValue::List(list) => { assert!(!list.is_empty()); } other => panic!("Expected scalar list result, got {other:?}"), @@ -342,12 +327,15 @@ mod tests { let udf: ScalarUDF = rs_polygonize_udf().into(); let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); - let raster_array = ScalarValue::Struct(Arc::new(load_test_raster())) - .to_array_of_size(2) - .unwrap(); let band_array = Arc::new(Int32Array::from(vec![Some(1), None])); let result = tester - .invoke_arrays(vec![raster_array, band_array]) + .invoke_arrays(vec![ + Arc::new(sedona_testing::raster_spec::raster_array(vec![ + Some(test_raster_spec()), + None, + ])), + band_array, + ]) .unwrap(); let list_array = as_list_array(&result).unwrap(); @@ -365,12 +353,11 @@ mod tests { let udf: ScalarUDF = rs_polygonize_udf().into(); let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); - let raster_array = Arc::new(load_test_raster()); let err = tester - .invoke_arrays(vec![ - raster_array, - Arc::new(Int32Array::from(vec![Some(99)])), - ]) + .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")); From 0b31e9b2b98925676b9edd4423b207a2ccbff9df Mon Sep 17 00:00:00 2001 From: kontinuation Date: Wed, 17 Jun 2026 21:45:49 +0800 Subject: [PATCH 3/9] test(rust/sedona-raster-gdal): update raster array helpers --- rust/sedona-raster-gdal/src/rs_polygonize.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/sedona-raster-gdal/src/rs_polygonize.rs b/rust/sedona-raster-gdal/src/rs_polygonize.rs index a1e512714..2d9dd37b4 100644 --- a/rust/sedona-raster-gdal/src/rs_polygonize.rs +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -278,7 +278,7 @@ mod tests { fn test_polygonize_raster() { let result = with_gdal(|gdal| { let raster_array = test_raster_spec().build(); - let raster_struct = RasterStructArray::new(&raster_array); + let raster_struct = RasterStructArray::try_new(&raster_array).unwrap(); let raster = raster_struct.get(0).unwrap(); polygonize_raster(gdal, &raster, 1) }) From 7c72117b6f6cbb9af13ce768aff5e463a505334c Mon Sep 17 00:00:00 2001 From: kontinuation Date: Thu, 18 Jun 2026 09:50:46 +0800 Subject: [PATCH 4/9] fix(rust/sedona-raster-gdal): use vector Memory driver for polygonize output --- rust/sedona-raster-gdal/src/rs_polygonize.rs | 153 +++++++++++-------- 1 file changed, 89 insertions(+), 64 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_polygonize.rs b/rust/sedona-raster-gdal/src/rs_polygonize.rs index 2d9dd37b4..5b9c96c08 100644 --- a/rust/sedona-raster-gdal/src/rs_polygonize.rs +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -31,7 +31,6 @@ use datafusion_expr::{ColumnarValue, Volatility}; use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; use sedona_gdal::gdal::Gdal; use sedona_gdal::gdal_dyn_bindgen::{OGRFieldType, OGRwkbGeometryType}; -use sedona_gdal::mem::MemDatasetBuilder; use sedona_raster::array::RasterRefImpl; use sedona_raster::traits::RasterRef; use sedona_raster_functions::RasterExecutor; @@ -171,63 +170,68 @@ fn polygonize_raster( .rasterband(band_num) .map_err(|e| exec_datafusion_err!("Failed to get band {}: {}", band_num, e))?; - let mem_ds = unsafe { - MemDatasetBuilder::new(0, 0) - .build(gdal) - .map_err(|e| exec_datafusion_err!("Failed to create memory dataset: {e}"))? - }; - - let spatial_ref = gdal_dataset.spatial_ref().ok(); - let mut layer = mem_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}"))?; - - gdal.polygonize( - &raster_band, - None, - &layer, - 0, - &sedona_gdal::raster::polygonize::PolygonizeOptions::default(), - ) - .map_err(|e| exec_datafusion_err!("GDAL polygonize failed: {e}"))?; - - let mut polygon_values = Vec::new(); - 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 polygon_values = (|| { + let driver = gdal + .get_driver_by_name("Memory") + .map_err(|e| exec_datafusion_err!("Failed to get Memory driver: {e}"))?; + let vector_ds = 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}"))?; + + gdal.polygonize( + &raster_band, + None, + &layer, + 0, + &sedona_gdal::raster::polygonize::PolygonizeOptions::default(), + ) + .map_err(|e| exec_datafusion_err!("GDAL polygonize failed: {e}"))?; - polygon_values.push((wkb, feature.field_as_double(idx))); - } + let mut polygon_values = Vec::new(); + 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 + } + }; + + polygon_values.push((wkb, feature.field_as_double(idx))); + } + + Ok::<_, datafusion_common::DataFusionError>(polygon_values) + })(); - Ok(polygon_values) + polygon_values } #[cfg(test)] @@ -238,10 +242,12 @@ mod tests { use datafusion_common::cast::{as_list_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::RasterSpec; use sedona_testing::testers::ScalarUdfTester; + use tempfile::tempdir; fn test_raster_spec() -> RasterSpec { RasterSpec::d2(3, 3) @@ -250,6 +256,28 @@ mod tests { .nodata(255u8) } + 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"); @@ -277,7 +305,7 @@ mod tests { #[test] fn test_polygonize_raster() { let result = with_gdal(|gdal| { - let raster_array = test_raster_spec().build(); + let raster_array = build_polygonize_test_raster(); let raster_struct = RasterStructArray::try_new(&raster_array).unwrap(); let raster = raster_struct.get(0).unwrap(); polygonize_raster(gdal, &raster, 1) @@ -328,14 +356,11 @@ mod tests { let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); let band_array = Arc::new(Int32Array::from(vec![Some(1), None])); + let raster_array = ScalarValue::Struct(Arc::new(build_polygonize_test_raster())) + .to_array_of_size(2) + .unwrap(); let result = tester - .invoke_arrays(vec![ - Arc::new(sedona_testing::raster_spec::raster_array(vec![ - Some(test_raster_spec()), - None, - ])), - band_array, - ]) + .invoke_arrays(vec![raster_array, band_array]) .unwrap(); let list_array = as_list_array(&result).unwrap(); From ed954fd4091b46dffd414640d52669f3726e40fc Mon Sep 17 00:00:00 2001 From: kontinuation Date: Sun, 21 Jun 2026 19:49:09 +0800 Subject: [PATCH 5/9] Slightly tuned the code to make it look better --- rust/sedona-raster-gdal/src/rs_polygonize.rs | 170 +++++++++---------- 1 file changed, 77 insertions(+), 93 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_polygonize.rs b/rust/sedona-raster-gdal/src/rs_polygonize.rs index 5b9c96c08..d4c6742a7 100644 --- a/rust/sedona-raster-gdal/src/rs_polygonize.rs +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -29,16 +29,19 @@ 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_raster::array::RasterRefImpl; +use sedona_gdal::raster::polygonize::PolygonizeOptions; use sedona_raster::traits::RasterRef; use sedona_raster_functions::RasterExecutor; use sedona_schema::datatypes::SedonaType; use sedona_schema::matchers::ArgMatcher; use crate::gdal_common::with_gdal; -use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider}; +use crate::gdal_dataset_provider::{ + configure_thread_local_options, thread_local_provider, RasterDataset, +}; pub fn rs_polygonize_udf() -> SedonaScalarUDF { SedonaScalarUDF::new( @@ -83,11 +86,17 @@ impl SedonaScalarKernel for RsPolygonize { let mut list_builder = ListBuilder::new(StructBuilder::from_fields( polygon_value_struct_fields(), - 16, + 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}"))?; + let mem_driver = gdal + .get_driver_by_name("MEM") + .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"); @@ -101,7 +110,16 @@ impl SedonaScalarKernel for RsPolygonize { let band_num: usize = band_opt.unwrap().max(1).try_into().unwrap_or(1); - let polygon_values = polygonize_raster(gdal, raster, band_num)?; + let bands = raster.bands(); + if band_num == 0 || band_num > bands.len() { + return exec_err!("Band {} is out of range (1-{})", band_num, bands.len()); + } + + let raster_ds = provider + .raster_ref_to_gdal(raster) + .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {e}"))?; + + let polygon_values = polygonize_raster(gdal, &mem_driver, raster_ds, band_num)?; let struct_builder = list_builder.values(); for (wkb, value) in polygon_values { let geom_builder = struct_builder @@ -149,89 +167,71 @@ fn polygon_value_struct_fields() -> Fields { ]) } +/// Perform GDAL polygonize on the given raster dataset and band, returning a list of (WKB, value) tuples. fn polygonize_raster( gdal: &Gdal, - raster: &RasterRefImpl<'_>, + mem_driver: &Driver, + raster_ds: RasterDataset<'_>, band_num: usize, ) -> Result, f64)>> { - let bands = raster.bands(); - if band_num == 0 || band_num > bands.len() { - return exec_err!("Band {} is out of range (1-{})", band_num, bands.len()); - } - - let provider = thread_local_provider(gdal) - .map_err(|e| exec_datafusion_err!("Failed to init GDAL provider: {e}"))?; - let raster_ds = provider - .raster_ref_to_gdal(raster) - .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {e}"))?; 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))?; - let polygon_values = (|| { - let driver = gdal - .get_driver_by_name("Memory") - .map_err(|e| exec_datafusion_err!("Failed to get Memory driver: {e}"))?; - let vector_ds = 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}"))?; - - gdal.polygonize( - &raster_band, - None, - &layer, - 0, - &sedona_gdal::raster::polygonize::PolygonizeOptions::default(), - ) - .map_err(|e| exec_datafusion_err!("GDAL polygonize failed: {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 mut polygon_values = Vec::new(); - 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 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}"))?; - polygon_values.push((wkb, feature.field_as_double(idx))); - } + // Polygonize the raster band into the vector layer, using the "value" field to store pixel values. + 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 polygon_values = Vec::new(); + 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 + } + }; - Ok::<_, datafusion_common::DataFusionError>(polygon_values) - })(); + polygon_values.push((wkb, feature.field_as_double(idx))); + } - polygon_values + Ok(polygon_values) } #[cfg(test)] @@ -283,32 +283,16 @@ mod tests { assert_eq!(rs_polygonize_udf().name(), "rs_polygonize"); } - #[test] - fn test_polygon_value_list_type() { - let dt = polygon_value_list_type(); - match dt { - DataType::List(field) => { - assert_eq!(field.name(), "item"); - match field.data_type() { - DataType::Struct(fields) => { - assert_eq!(fields.len(), 2); - assert_eq!(fields[0].name(), "geom"); - assert_eq!(fields[1].name(), "value"); - } - other => panic!("Expected Struct data type, got {other:?}"), - } - } - other => panic!("Expected List data type, got {other:?}"), - } - } - #[test] fn test_polygonize_raster() { let result = with_gdal(|gdal| { + let provider = thread_local_provider(gdal).unwrap(); + let mem_driver = gdal.get_driver_by_name("MEM").unwrap(); let raster_array = build_polygonize_test_raster(); let raster_struct = RasterStructArray::try_new(&raster_array).unwrap(); let raster = raster_struct.get(0).unwrap(); - polygonize_raster(gdal, &raster, 1) + let raster_ds = provider.raster_ref_to_gdal(&raster).unwrap(); + polygonize_raster(gdal, &mem_driver, raster_ds, 1) }) .unwrap(); From bc9daf3218b383611626b4af57ea6d67ebb0063c Mon Sep 17 00:00:00 2001 From: kontinuation Date: Sun, 21 Jun 2026 23:02:51 +0800 Subject: [PATCH 6/9] Switch back from MEM to Memory to be compatible with old versions of GDAL --- rust/sedona-raster-gdal/src/rs_polygonize.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_polygonize.rs b/rust/sedona-raster-gdal/src/rs_polygonize.rs index d4c6742a7..87dd1bb8d 100644 --- a/rust/sedona-raster-gdal/src/rs_polygonize.rs +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -93,8 +93,9 @@ impl SedonaScalarKernel for RsPolygonize { 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("MEM") + .get_driver_by_name("Memory") .map_err(|e| exec_datafusion_err!("Failed to get Memory driver: {e}"))?; executor.execute_raster_void(|_, raster_opt| { @@ -287,7 +288,7 @@ mod tests { fn test_polygonize_raster() { let result = with_gdal(|gdal| { let provider = thread_local_provider(gdal).unwrap(); - let mem_driver = gdal.get_driver_by_name("MEM").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(); From 4f59212ccfad238fff60152d1666a86ec81817aa Mon Sep 17 00:00:00 2001 From: kontinuation Date: Sun, 28 Jun 2026 17:04:04 +0800 Subject: [PATCH 7/9] Refactor rs_polygonize to propagate CRS and amortize allocations via builders --- rust/sedona-raster-gdal/src/rs_polygonize.rs | 178 +++++++++++++------ 1 file changed, 127 insertions(+), 51 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_polygonize.rs b/rust/sedona-raster-gdal/src/rs_polygonize.rs index 87dd1bb8d..f15a3467b 100644 --- a/rust/sedona-raster-gdal/src/rs_polygonize.rs +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -20,8 +20,10 @@ use std::convert::TryInto; use std::sync::Arc; -use arrow_array::builder::{BinaryBuilder, Float64Builder, ListBuilder, StructBuilder}; -use arrow_array::{ArrayRef, Int32Array}; +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; @@ -35,7 +37,7 @@ use sedona_gdal::gdal_dyn_bindgen::{OGRFieldType, OGRwkbGeometryType}; use sedona_gdal::raster::polygonize::PolygonizeOptions; use sedona_raster::traits::RasterRef; use sedona_raster_functions::RasterExecutor; -use sedona_schema::datatypes::SedonaType; +use sedona_schema::datatypes::{SedonaType, WKB_GEOMETRY_ITEM_CRS}; use sedona_schema::matchers::ArgMatcher; use crate::gdal_common::with_gdal; @@ -84,10 +86,11 @@ impl SedonaScalarKernel for RsPolygonize { let band_array = band_argument_array(&args[1], num_iterations)?; let mut band_iter = band_array.iter(); - let mut list_builder = ListBuilder::new(StructBuilder::from_fields( - polygon_value_struct_fields(), - num_iterations, - )); + 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)?; @@ -104,7 +107,8 @@ impl SedonaScalarKernel for RsPolygonize { let raster = match (raster_opt, band_opt) { (Some(raster), Some(_)) => raster, _ => { - list_builder.append_null(); + valid_list_items.append_null(); + list_offsets.push_length(0); return Ok(()); } }; @@ -120,26 +124,30 @@ impl SedonaScalarKernel for RsPolygonize { .raster_ref_to_gdal(raster) .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {e}"))?; - let polygon_values = polygonize_raster(gdal, &mem_driver, raster_ds, band_num)?; - let struct_builder = list_builder.values(); - for (wkb, value) in polygon_values { - let geom_builder = struct_builder - .field_builder::(0) - .expect("Expected BinaryBuilder for geom field"); - geom_builder.append_value(&wkb); - - let value_builder = struct_builder - .field_builder::(1) - .expect("Expected Float64Builder for value field"); + 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(()) + })?; - struct_builder.append(true); - } - list_builder.append(true); + valid_list_items.append_non_null(); + list_offsets.push_length(num_polygons); Ok(()) })?; - let result: ArrayRef = Arc::new(list_builder.finish()); + 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) }) } @@ -163,18 +171,65 @@ fn polygon_value_list_type() -> DataType { fn polygon_value_struct_fields() -> Fields { Fields::from(vec![ - Field::new("geom", DataType::Binary, false), + WKB_GEOMETRY_ITEM_CRS + .to_storage_field("geom", false) + .unwrap(), Field::new("value", DataType::Float64, false), ]) } -/// Perform GDAL polygonize on the given raster dataset and band, returning a list of (WKB, value) tuples. -fn polygonize_raster( +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, -) -> Result, f64)>> { + mut callback: F, +) -> Result<()> +where + F: FnMut(&[u8], f64) -> Result<()>, +{ let gdal_dataset = raster_ds.as_dataset(); let raster_band = gdal_dataset @@ -208,7 +263,6 @@ fn polygonize_raster( .map_err(|e| exec_datafusion_err!("GDAL polygonize failed: {e}"))?; // Extract the WKB geometry and value for each feature in the layer. - let mut polygon_values = Vec::new(); let mut value_field_idx: Option = None; for feature in layer.features() { let geom = feature @@ -229,10 +283,12 @@ fn polygonize_raster( } }; - polygon_values.push((wkb, feature.field_as_double(idx))); + let value = feature.field_as_double(idx); + + callback(&wkb, value)?; } - Ok(polygon_values) + Ok(()) } #[cfg(test)] @@ -240,13 +296,13 @@ mod tests { use super::*; use arrow_array::Array; - use datafusion_common::cast::{as_list_array, as_struct_array}; + use datafusion_common::cast::{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::RasterSpec; + use sedona_testing::raster_spec::{raster_array, RasterSpec}; use sedona_testing::testers::ScalarUdfTester; use tempfile::tempdir; @@ -286,22 +342,26 @@ mod tests { #[test] fn test_polygonize_raster() { - let result = with_gdal(|gdal| { + 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) + polygonize_raster(gdal, &mem_driver, raster_ds, 1, |wkb, value| { + results.push((wkb.to_vec(), value)); + Ok(()) + }) }) .unwrap(); assert!( - !result.is_empty(), + !results.is_empty(), "Polygonize should return at least one polygon" ); - for (wkb, value) in &result { + 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"); } @@ -327,11 +387,16 @@ mod tests { .invoke_scalar_scalar(test_raster_spec(), ScalarValue::Int32(Some(1))) .unwrap(); - match result { - ScalarValue::List(list) => { - assert!(!list.is_empty()); - } - other => panic!("Expected scalar list result, got {other:?}"), + 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"); } } @@ -340,22 +405,33 @@ mod tests { let udf: ScalarUDF = rs_polygonize_udf().into(); let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); - let band_array = Arc::new(Int32Array::from(vec![Some(1), None])); - let raster_array = ScalarValue::Struct(Arc::new(build_polygonize_test_raster())) - .to_array_of_size(2) - .unwrap(); + let rasters = vec![Some(test_raster_spec()), None, Some(test_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(), 2); - assert!(!list_array.is_null(0)); - assert!(list_array.is_null(1)); + assert_eq!(list_array.len(), 3); + + // Row 0, 2: non-null, non-empty list of polygons + for k in [0, 2] { + assert!(!list_array.is_null(k)); + let row = list_array.value(k); + let row = as_struct_array(&row).unwrap(); + assert!(!row.is_empty()); + let geom = as_struct_array(row.column(0)).unwrap(); + let crs = as_string_view_array(geom.column(1)).unwrap(); + assert!(!crs.is_empty()); + for i in 0..crs.len() { + assert_eq!(crs.value(i), "OGC:CRS84"); + } + } - let row0 = list_array.value(0); - let row0 = as_struct_array(&row0).unwrap(); - assert!(!row0.is_empty()); + // Row 1: null list (None) + assert!(list_array.is_null(1)); } #[test] From b7d8224d560253af092678e0b1205c59b3fd8327 Mon Sep 17 00:00:00 2001 From: kontinuation Date: Sun, 28 Jun 2026 17:33:29 +0800 Subject: [PATCH 8/9] Support GDALFPolygonize for floating-point raster bands and fix negative band indexing --- c/sedona-gdal/src/gdal.rs | 22 +++++++++++++- c/sedona-gdal/src/raster/polygonize.rs | 32 ++++++++++++++++++++ rust/sedona-raster-gdal/src/rs_polygonize.rs | 22 +++++++++----- 3 files changed, 68 insertions(+), 8 deletions(-) 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/rust/sedona-raster-gdal/src/rs_polygonize.rs b/rust/sedona-raster-gdal/src/rs_polygonize.rs index f15a3467b..1fc124d49 100644 --- a/rust/sedona-raster-gdal/src/rs_polygonize.rs +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -17,7 +17,6 @@ //! RS_Polygonize UDF - Convert a raster band to vector polygons. -use std::convert::TryInto; use std::sync::Arc; use arrow_array::builder::{ @@ -35,6 +34,7 @@ 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}; @@ -113,12 +113,12 @@ impl SedonaScalarKernel for RsPolygonize { } }; - let band_num: usize = band_opt.unwrap().max(1).try_into().unwrap_or(1); - + let band_index = band_opt.unwrap(); let bands = raster.bands(); - if band_num == 0 || band_num > bands.len() { - return exec_err!("Band {} is out of range (1-{})", band_num, bands.len()); + 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) @@ -259,8 +259,16 @@ where .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. - gdal.polygonize(&raster_band, None, &layer, 0, &PolygonizeOptions::default()) - .map_err(|e| exec_datafusion_err!("GDAL polygonize failed: {e}"))?; + 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; From 9bce7bf407968a62a6191ae2f294a865fe43e015 Mon Sep 17 00:00:00 2001 From: kontinuation Date: Sun, 28 Jun 2026 18:43:50 +0800 Subject: [PATCH 9/9] Add unit test verifying float and integer band polygonization precision --- rust/sedona-raster-gdal/src/rs_polygonize.rs | 61 +++++++++++++++----- 1 file changed, 47 insertions(+), 14 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_polygonize.rs b/rust/sedona-raster-gdal/src/rs_polygonize.rs index 1fc124d49..1c09c272d 100644 --- a/rust/sedona-raster-gdal/src/rs_polygonize.rs +++ b/rust/sedona-raster-gdal/src/rs_polygonize.rs @@ -304,7 +304,9 @@ mod tests { use super::*; use arrow_array::Array; - use datafusion_common::cast::{as_list_array, as_string_view_array, as_struct_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; @@ -321,6 +323,13 @@ mod tests { .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"); @@ -413,7 +422,11 @@ mod tests { 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_raster_spec())]; + 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)])); @@ -424,22 +437,42 @@ mod tests { assert_eq!(list_array.len(), 3); - // Row 0, 2: non-null, non-empty list of polygons - for k in [0, 2] { - assert!(!list_array.is_null(k)); - let row = list_array.value(k); - let row = as_struct_array(&row).unwrap(); - assert!(!row.is_empty()); - let geom = as_struct_array(row.column(0)).unwrap(); - let crs = as_string_view_array(geom.column(1)).unwrap(); - assert!(!crs.is_empty()); - for i in 0..crs.len() { - assert_eq!(crs.value(i), "OGC:CRS84"); - } + // 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]