You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by ko...@apache.org on 2026/07/03 10:04:20 UTC
(sedona-db) branch main updated: feat(rust/sedona-raster-gdal): add RS_Polygonize (#955)
This is an automated email from the ASF dual-hosted git repository.
Kontinuation pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new 1eca227f7 feat(rust/sedona-raster-gdal): add RS_Polygonize (#955)
1eca227f7 is described below
commit 1eca227f77d7d13484fde680307b05b7e9abb135
Author: Kristin Cowalcijk <kr...@gmail.com>
AuthorDate: Fri Jul 3 18:04:15 2026 +0800
feat(rust/sedona-raster-gdal): add RS_Polygonize (#955)
Add the GDAL-backed `RS_Polygonize` raster function for extracting polygons from a given raster
---
c/sedona-gdal/src/gdal.rs | 22 +-
c/sedona-gdal/src/raster/polygonize.rs | 32 ++
docs/reference/sql/rs_polygonize.qmd | 43 ++
rust/sedona-raster-gdal/Cargo.toml | 5 +
rust/sedona-raster-gdal/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 | 492 +++++++++++++++++++++++
8 files changed, 687 insertions(+), 1 deletion(-)
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 a8899f4e4..03a82513d 100644
--- a/rust/sedona-raster-gdal/Cargo.toml
+++ b/rust/sedona-raster-gdal/Cargo.toml
@@ -67,3 +67,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<Option<SedonaType>> {
+ 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<ColumnarValue> {
+ self.invoke_batch_from_args(arg_types, args, &SedonaType::Arrow(DataType::Null), 0, None)
+ }
+
+ fn invoke_batch_from_args(
+ &self,
+ arg_types: &[SedonaType],
+ args: &[ColumnarValue],
+ _return_type: &SedonaType,
+ _num_rows: usize,
+ config_options: Option<&ConfigOptions>,
+ ) -> Result<ColumnarValue> {
+ let 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::<i32>::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<Int32Array> {
+ 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<i32>,
+ 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<F>(
+ 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<i32> = 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::<u8>(&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::<Vec<_>>();
+ 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::<Vec<_>>();
+ 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"));
+ }
+}