From bbdeb8cafd36c179658339844f29dac847d8aea8 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 6 Jul 2026 17:51:01 -0700 Subject: [PATCH 1/8] feat(rust/sedona-raster-gdal): add RS_AsGeoTiff Export a raster to GeoTIFF bytes (binary), with optional compression and tiling overloads. Ported from the GDAL raster draft PR #704 reusing that implementation; reworked the tests onto the RasterSpec harness (dropping the external .tiff fixtures and the roundtrip's dependency on the not-yet-split RS_FromGDALRaster) and added a RasterSpec-based benchmark and SQL docs. --- docs/reference/sql/rs_asgeotiff.qmd | 86 +++ rust/sedona-raster-gdal/Cargo.toml | 5 + .../benches/rs_as_geotiff.rs | 103 ++++ rust/sedona-raster-gdal/src/lib.rs | 2 + rust/sedona-raster-gdal/src/register.rs | 1 + rust/sedona-raster-gdal/src/rs_as_geotiff.rs | 533 ++++++++++++++++++ 6 files changed, 730 insertions(+) create mode 100644 docs/reference/sql/rs_asgeotiff.qmd create mode 100644 rust/sedona-raster-gdal/benches/rs_as_geotiff.rs create mode 100644 rust/sedona-raster-gdal/src/rs_as_geotiff.rs diff --git a/docs/reference/sql/rs_asgeotiff.qmd b/docs/reference/sql/rs_asgeotiff.qmd new file mode 100644 index 0000000000..99f9573a79 --- /dev/null +++ b/docs/reference/sql/rs_asgeotiff.qmd @@ -0,0 +1,86 @@ +--- +# 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_AsGeoTiff +description: > + Encodes a raster as a GeoTIFF and returns the bytes as binary, with optional + compression and tiling. +kernels: + - returns: binary + args: + - raster + - returns: binary + args: + - raster + - {name: tile_size, type: integer} + - returns: binary + args: + - raster + - {name: compression, type: string} + - {name: quality, type: double} + - returns: binary + args: + - raster + - {name: compression, type: string} + - {name: quality, type: double} + - {name: tile_size, type: integer} + - returns: binary + args: + - raster + - name: compression + type: string + description: > + GeoTIFF compression codec (case-insensitive): None, PackBits, Deflate, + Huffman, LZW, or JPEG. Deflate and LZW additionally enable a horizontal + predictor to improve the compression ratio. + - name: quality + type: double + description: Quality factor for lossy JPEG compression; ignored by the other codecs. + - name: tile_width + type: integer + description: > + Tile width. When given together with tile_height, the GeoTIFF is written + tiled rather than stripped. + - name: tile_height + type: integer + description: Tile height. +--- + +## Description + +`RS_AsGeoTiff` encodes a raster as a GeoTIFF (via GDAL) and returns the file +bytes as a binary value — the inverse of reading a GeoTIFF into a raster. The +result is `NULL` for a `NULL` raster. + +The two-argument form `RS_AsGeoTiff(raster, tile_size)` writes a tiled GeoTIFF +with square `tile_size` blocks. Compression is selected by name; a JPEG quality +factor and explicit tile dimensions can be supplied via the longer overloads. + +## Examples + +```sql +SELECT RS_AsGeoTiff(RS_Example()) IS NOT NULL; +``` + +```sql +SELECT RS_AsGeoTiff(RS_Example(), 'DEFLATE', 6) IS NOT NULL; +``` + +```sql +SELECT RS_AsGeoTiff(RS_Example(), 'LZW', 75, 256, 256) IS NOT NULL; +``` diff --git a/rust/sedona-raster-gdal/Cargo.toml b/rust/sedona-raster-gdal/Cargo.toml index 03a82513d8..fa33479ba1 100644 --- a/rust/sedona-raster-gdal/Cargo.toml +++ b/rust/sedona-raster-gdal/Cargo.toml @@ -58,6 +58,11 @@ sedona-testing = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } +[[bench]] +harness = false +name = "rs_as_geotiff" +path = "benches/rs_as_geotiff.rs" + [[bench]] harness = false name = "rs_frompath" diff --git a/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs b/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs new file mode 100644 index 0000000000..cc3ff8e9b2 --- /dev/null +++ b/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs @@ -0,0 +1,103 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for the RS_AsGeoTiff UDF. +//! +//! Exercises the raster → GeoTIFF export path (`RS_Example` rasters, no external +//! fixtures): the uncompressed default and the compression axis (none/lzw/deflate). + +use std::hint::black_box; + +use arrow_array::ArrayRef; +use arrow_schema::DataType; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use datafusion_common::ScalarValue; +use datafusion_expr::{ColumnarValue, 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![]); + match tester.invoke(vec![]).unwrap() { + ColumnarValue::Scalar(value) => value.to_array_of_size(rows).unwrap(), + ColumnarValue::Array(array) => array, + } +} + +fn bench_rs_as_geotiff_basic(c: &mut Criterion) { + let udf: ScalarUDF = sedona_raster_gdal::rs_as_geotiff_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![RASTER]); + + let mut group = c.benchmark_group("rs_as_geotiff"); + for rows in [1usize, 32] { + let rasters = raster_array(rows); + group.throughput(Throughput::Elements(rows as u64)); + group.bench_with_input(BenchmarkId::new("basic", rows), &rasters, |b, input| { + b.iter(|| black_box(tester.invoke_arrays(vec![input.clone()]).unwrap())) + }); + } + group.finish(); +} + +fn bench_rs_as_geotiff_compression(c: &mut Criterion) { + let udf: ScalarUDF = sedona_raster_gdal::rs_as_geotiff_udf().into(); + let tester = ScalarUdfTester::new( + udf, + vec![ + RASTER, + SedonaType::Arrow(DataType::Utf8), + SedonaType::Arrow(DataType::Float64), + ], + ); + + let rasters = raster_array(32); + let quality = ColumnarValue::Scalar(ScalarValue::Float64(Some(75.0))); + + let mut group = c.benchmark_group("rs_as_geotiff_compression"); + group.throughput(Throughput::Elements(rasters.len() as u64)); + for compression in ["none", "lzw", "deflate"] { + let comp = ColumnarValue::Scalar(ScalarValue::Utf8(Some(compression.to_string()))); + group.bench_with_input( + BenchmarkId::new("compression", compression), + &(&rasters, &comp, &quality), + |b, (rasters, comp, quality)| { + b.iter(|| { + black_box( + tester + .invoke(vec![ + ColumnarValue::Array((*rasters).clone()), + (*comp).clone(), + (*quality).clone(), + ]) + .unwrap(), + ) + }) + }, + ); + } + group.finish(); +} + +criterion_group!( + benches, + bench_rs_as_geotiff_basic, + bench_rs_as_geotiff_compression +); +criterion_main!(benches); diff --git a/rust/sedona-raster-gdal/src/lib.rs b/rust/sedona-raster-gdal/src/lib.rs index f878f34a12..2c8a4103ab 100644 --- a/rust/sedona-raster-gdal/src/lib.rs +++ b/rust/sedona-raster-gdal/src/lib.rs @@ -33,6 +33,7 @@ mod gdal_common; mod gdal_dataset_provider; mod raster_loader; +mod rs_as_geotiff; mod rs_frompath; mod rs_metadata; mod rs_polygonize; @@ -45,6 +46,7 @@ pub use gdal_common::{ nodata_bytes_to_f64, nodata_f64_to_bytes, GdalBandLayout, GdalBandPlan, }; pub use raster_loader::{GdalLoader, GDAL_FORMAT}; +pub use rs_as_geotiff::rs_as_geotiff_udf; pub use rs_frompath::rs_frompath_udf; pub use rs_metadata::rs_metadata_udf; pub use rs_polygonize::rs_polygonize_udf; diff --git a/rust/sedona-raster-gdal/src/register.rs b/rust/sedona-raster-gdal/src/register.rs index 52f7dfd488..e756b16d4c 100644 --- a/rust/sedona-raster-gdal/src/register.rs +++ b/rust/sedona-raster-gdal/src/register.rs @@ -20,6 +20,7 @@ use sedona_expr::function_set::FunctionSet; /// Export the set of GDAL-backed functions defined in this crate. pub fn default_function_set() -> FunctionSet { let mut function_set = FunctionSet::new(); + function_set.insert_scalar_udf(crate::rs_as_geotiff::rs_as_geotiff_udf()); 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()); diff --git a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs new file mode 100644 index 0000000000..1d5a22b5f4 --- /dev/null +++ b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs @@ -0,0 +1,533 @@ +// 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_AsGeoTiff UDF - Export raster as GeoTiff binary +//! +//! Returns a binary DataFrame from a Raster DataFrame with multiple overloads: +//! - RS_AsGeoTiff(raster) +//! - RS_AsGeoTiff(raster, tileSize) +//! - RS_AsGeoTiff(raster, compressionType, imageQuality) +//! - RS_AsGeoTiff(raster, compressionType, imageQuality, tileSize) +//! - RS_AsGeoTiff(raster, compressionType, imageQuality, tileWidth, tileHeight) + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use crate::gdal_common::with_gdal; +use arrow_array::builder::BinaryBuilder; +use arrow_schema::DataType; +use datafusion_common::cast::{as_float64_array, as_string_array, as_uint32_array}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::error::Result; +use datafusion_common::{exec_datafusion_err, ScalarValue}; +use datafusion_expr::{ColumnarValue, Volatility}; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_raster::array::RasterRefImpl; +use sedona_raster_functions::RasterExecutor; +use sedona_schema::datatypes::SedonaType; +use sedona_schema::matchers::ArgMatcher; + +// Use thread-local provider to create GDAL datasets from `RasterRef`. +use crate::gdal_dataset_provider::configure_thread_local_options; + +/// Counter for generating unique VSI memory file names +static VSI_FILE_COUNTER: AtomicUsize = AtomicUsize::new(0); + +/// Compression types supported for GeoTiff output +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CompressionType { + None, + PackBits, + Deflate, + Huffman, + Lzw, + Jpeg, +} + +impl CompressionType { + /// Parse compression type from string (case-insensitive) + pub fn parse(s: &str) -> Option { + match s.to_lowercase().as_str() { + "none" => Some(CompressionType::None), + "packbits" => Some(CompressionType::PackBits), + "deflate" => Some(CompressionType::Deflate), + "huffman" => Some(CompressionType::Huffman), + "lzw" => Some(CompressionType::Lzw), + "jpeg" => Some(CompressionType::Jpeg), + _ => None, + } + } + + /// Get GDAL compression option value + pub fn gdal_value(&self) -> &'static str { + match self { + CompressionType::None => "NONE", + CompressionType::PackBits => "PACKBITS", + CompressionType::Deflate => "DEFLATE", + CompressionType::Huffman => "CCITTRLE", + CompressionType::Lzw => "LZW", + CompressionType::Jpeg => "JPEG", + } + } +} + +/// RS_AsGeoTiff() scalar UDF implementation +/// +/// Returns a binary DataFrame from a Raster DataFrame +pub fn rs_as_geotiff_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_asgeotiff", + vec![ + Arc::new(RsAsGeoTiff::new(Variant::Basic)), // RS_AsGeoTiff(raster) + Arc::new(RsAsGeoTiff::new(Variant::WithTileSize)), // RS_AsGeoTiff(raster, tileSize) + Arc::new(RsAsGeoTiff::new(Variant::WithCompressionQuality)), // RS_AsGeoTiff(raster, compression, quality) + Arc::new(RsAsGeoTiff::new(Variant::WithCompressionQualityTileSize)), // RS_AsGeoTiff(raster, compression, quality, tileSize) + Arc::new(RsAsGeoTiff::new(Variant::WithCompressionQualityTileWH)), // RS_AsGeoTiff(raster, compression, quality, tileWidth, tileHeight) + ], + Volatility::Immutable, + ) +} + +/// Variants for different overloads +#[derive(Debug, Clone, Copy)] +enum Variant { + Basic, // (raster) + WithTileSize, // (raster, tileSize) + WithCompressionQuality, // (raster, compression, quality) + WithCompressionQualityTileSize, // (raster, compression, quality, tileSize) + WithCompressionQualityTileWH, // (raster, compression, quality, tileWidth, tileHeight) +} + +/// Kernel implementation for RS_AsGeoTiff +#[derive(Debug)] +struct RsAsGeoTiff { + variant: Variant, +} + +impl RsAsGeoTiff { + fn new(variant: Variant) -> Self { + Self { variant } + } + + /// Generate a unique VSI memory file path + fn generate_vsi_path() -> String { + let counter = VSI_FILE_COUNTER.fetch_add(1, Ordering::SeqCst); + let thread_id = std::thread::current().id(); + format!("/vsimem/rs_as_geotiff_{:?}_{}.tif", thread_id, counter) + } + + /// Convert raster to GeoTiff bytes + fn raster_to_geotiff( + gdal: &sedona_gdal::gdal::Gdal, + raster: &RasterRefImpl, + compression: Option, + quality: Option, + tile_width: Option, + tile_height: Option, + ) -> Result> { + let provider = crate::gdal_dataset_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 source_dataset = raster_ds.as_dataset(); + + let driver = gdal + .get_driver_by_name("GTiff") + .map_err(|e| exec_datafusion_err!("Failed to get GTiff driver: {}", e))?; + + // Build creation options as string list + let mut options_list: Vec = Vec::new(); + + // Add compression option + if let Some(comp) = compression { + options_list.push(format!("COMPRESS={}", comp.gdal_value())); + + // Add quality for JPEG + if comp == CompressionType::Jpeg { + if let Some(q) = quality { + // JPEG quality is 1-100, we receive 0.0-1.0 + let jpeg_quality = (q * 100.0).round() as i32; + options_list.push(format!("JPEG_QUALITY={}", jpeg_quality.clamp(1, 100))); + } + } + + // Add predictor for Deflate/LZW (improves compression) + if comp == CompressionType::Deflate || comp == CompressionType::Lzw { + options_list.push("PREDICTOR=2".to_string()); + } + } + + // Add tiling options + if let (Some(tw), Some(th)) = (tile_width, tile_height) { + options_list.push("TILED=YES".to_string()); + options_list.push(format!("BLOCKXSIZE={}", tw)); + options_list.push(format!("BLOCKYSIZE={}", th)); + } + + // Convert to creation options slice + let options_refs: Vec<&str> = options_list.iter().map(|s| s.as_str()).collect(); + + // Generate VSI path for output + let vsi_path = Self::generate_vsi_path(); + + // Create copy to VSI memory file + let _output_dataset = source_dataset + .create_copy(&driver, &vsi_path, &options_refs) + .map_err(|e| exec_datafusion_err!("Failed to create GeoTiff: {}", e))?; + + // Close the output dataset to flush data + drop(_output_dataset); + + // Read bytes from VSI memory file and clean up + let bytes = gdal.get_vsi_mem_file_bytes_owned(&vsi_path).map_err(|e| { + let _ = gdal.unlink_mem_file(&vsi_path); + exec_datafusion_err!("Failed to read GeoTiff bytes: {}", e) + })?; + + // Clean up VSI file + let _ = gdal.unlink_mem_file(&vsi_path); + + Ok(bytes) + } +} + +impl SedonaScalarKernel for RsAsGeoTiff { + fn return_type(&self, args: &[SedonaType]) -> Result> { + let matchers = match self.variant { + Variant::Basic => vec![ArgMatcher::is_raster()], + Variant::WithTileSize => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_integer(), // tileSize + ], + Variant::WithCompressionQuality => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_string(), // compressionType + ArgMatcher::is_numeric(), // imageQuality + ], + Variant::WithCompressionQualityTileSize => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_string(), // compressionType + ArgMatcher::is_numeric(), // imageQuality + ArgMatcher::is_integer(), // tileSize + ], + Variant::WithCompressionQualityTileWH => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_string(), // compressionType + ArgMatcher::is_numeric(), // imageQuality + ArgMatcher::is_integer(), // tileWidth + ArgMatcher::is_integer(), // tileHeight + ], + }; + + let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(DataType::Binary)); + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + self.invoke_batch_from_args(arg_types, args, &SedonaType::Arrow(DataType::Null), 0, None) + } + + fn invoke_batch_from_args( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + _return_type: &SedonaType, + _num_rows: usize, + config_options: Option<&ConfigOptions>, + ) -> Result { + let executor = RasterExecutor::new(arg_types, args); + let num_iterations = executor.num_iterations(); + + // Convert variant-specific args to arrays upfront via into_array. + // For variants that don't use a parameter, create null-filled default arrays. + let (compression_array, quality_array, tile_width_array, tile_height_array) = + match self.variant { + Variant::Basic => { + // No extra args → all null arrays + let compression = ScalarValue::Utf8(None).to_array_of_size(num_iterations)?; + let quality = ScalarValue::Float64(None).to_array_of_size(num_iterations)?; + let tile_width = ScalarValue::UInt32(None).to_array_of_size(num_iterations)?; + let tile_height = ScalarValue::UInt32(None).to_array_of_size(num_iterations)?; + (compression, quality, tile_width, tile_height) + } + Variant::WithTileSize => { + // args[1] → tile_width AND tile_height + let compression = ScalarValue::Utf8(None).to_array_of_size(num_iterations)?; + let quality = ScalarValue::Float64(None).to_array_of_size(num_iterations)?; + let tile_size = args[1] + .clone() + .cast_to(&DataType::UInt32, None)? + .into_array(num_iterations)?; + (compression, quality, tile_size.clone(), tile_size) + } + Variant::WithCompressionQuality => { + // args[1] → compression, args[2] → quality + let compression = args[1] + .clone() + .cast_to(&DataType::Utf8, None)? + .into_array(num_iterations)?; + let quality = args[2] + .clone() + .cast_to(&DataType::Float64, None)? + .into_array(num_iterations)?; + let tile_width = ScalarValue::UInt32(None).to_array_of_size(num_iterations)?; + let tile_height = ScalarValue::UInt32(None).to_array_of_size(num_iterations)?; + (compression, quality, tile_width, tile_height) + } + Variant::WithCompressionQualityTileSize => { + // args[1] → compression, args[2] → quality, args[3] → tile_width AND tile_height + let compression = args[1] + .clone() + .cast_to(&DataType::Utf8, None)? + .into_array(num_iterations)?; + let quality = args[2] + .clone() + .cast_to(&DataType::Float64, None)? + .into_array(num_iterations)?; + let tile_size = args[3] + .clone() + .cast_to(&DataType::UInt32, None)? + .into_array(num_iterations)?; + (compression, quality, tile_size.clone(), tile_size) + } + Variant::WithCompressionQualityTileWH => { + // args[1] → compression, args[2] → quality, args[3] → tile_width, args[4] → tile_height + let compression = args[1] + .clone() + .cast_to(&DataType::Utf8, None)? + .into_array(num_iterations)?; + let quality = args[2] + .clone() + .cast_to(&DataType::Float64, None)? + .into_array(num_iterations)?; + let tile_width = args[3] + .clone() + .cast_to(&DataType::UInt32, None)? + .into_array(num_iterations)?; + let tile_height = args[4] + .clone() + .cast_to(&DataType::UInt32, None)? + .into_array(num_iterations)?; + (compression, quality, tile_width, tile_height) + } + }; + + // Downcast all parameter arrays once before the loop + let compression_array = as_string_array(&compression_array)?; + let quality_array = as_float64_array(&quality_array)?; + let tile_width_array = as_uint32_array(&tile_width_array)?; + let tile_height_array = as_uint32_array(&tile_height_array)?; + + // Create iterators for each parameter array + let mut compression_iter = compression_array.iter(); + let mut quality_iter = quality_array.iter(); + let mut tile_width_iter = tile_width_array.iter(); + let mut tile_height_iter = tile_height_array.iter(); + + // Build output binary array + let mut builder = BinaryBuilder::with_capacity(num_iterations, num_iterations * 1024); + + with_gdal(|gdal| { + configure_thread_local_options(gdal, config_options)?; + executor.execute_raster_void(|_i, raster_opt| { + let compression_opt = compression_iter.next().unwrap(); + let quality_opt = quality_iter.next().unwrap(); + let tile_width_opt = tile_width_iter.next().unwrap(); + let tile_height_opt = tile_height_iter.next().unwrap(); + + let raster = match raster_opt { + Some(raster) => raster, + None => { + builder.append_null(); + return Ok(()); + } + }; + + let compression = match compression_opt { + Some(comp_str) => Some(CompressionType::parse(comp_str).ok_or_else(|| { + exec_datafusion_err!( + "Unknown compression type: {}. Valid values: None, PackBits, Deflate, Huffman, LZW, JPEG", + comp_str + ) + })?), + None => None, + }; + + let quality = quality_opt; + let tile_width = tile_width_opt; + let tile_height = tile_height_opt; + + let bytes = Self::raster_to_geotiff( + gdal, + raster, + compression, + quality, + tile_width, + tile_height, + )?; + builder.append_value(&bytes); + + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion_expr::ScalarUDF; + use sedona_gdal::gdal_dyn_bindgen::{GDAL_OF_RASTER, GDAL_OF_READONLY}; + use sedona_gdal::raster::types::DatasetOptions; + use sedona_raster::array::RasterStructArray; + use sedona_raster::traits::RasterRef; + use sedona_schema::datatypes::RASTER; + use sedona_testing::raster_spec::RasterSpec; + use sedona_testing::testers::ScalarUdfTester; + + /// A 3x3 single-band UInt8 raster with a CRS (RasterSpec defaults to one), + /// so GDAL export sets a projection. + 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, 2, 3, 4, 5, 6, 7, 8, 9]) + } + + /// Build a one-row raster `StructArray` from a spec. + fn as_raster_array(spec: RasterSpec) -> arrow_array::StructArray { + spec.build() + } + + #[test] + fn test_compression_type_parse() { + assert_eq!(CompressionType::parse("none"), Some(CompressionType::None)); + assert_eq!(CompressionType::parse("NONE"), Some(CompressionType::None)); + assert_eq!( + CompressionType::parse("deflate"), + Some(CompressionType::Deflate) + ); + assert_eq!( + CompressionType::parse("DEFLATE"), + Some(CompressionType::Deflate) + ); + assert_eq!(CompressionType::parse("lzw"), Some(CompressionType::Lzw)); + assert_eq!(CompressionType::parse("jpeg"), Some(CompressionType::Jpeg)); + assert_eq!(CompressionType::parse("invalid"), None); + } + + #[test] + fn test_generate_vsi_path() { + let path1 = RsAsGeoTiff::generate_vsi_path(); + let path2 = RsAsGeoTiff::generate_vsi_path(); + + assert!(path1.starts_with("/vsimem/rs_as_geotiff_")); + assert!(path1.ends_with(".tif")); + assert!(path2.starts_with("/vsimem/rs_as_geotiff_")); + assert_ne!(path1, path2); + } + + #[test] + fn udf_as_geotiff() { + let udf: datafusion_expr::ScalarUDF = rs_as_geotiff_udf().into(); + assert_eq!(udf.name(), "rs_asgeotiff"); + } + + #[test] + fn as_geotiff_produces_valid_tiff() { + // End-to-end through the UDF: a raster scalar in, GeoTIFF binary out. + let udf: ScalarUDF = rs_as_geotiff_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![RASTER]); + let result = tester.invoke_scalar(test_raster_spec()).unwrap(); + let ScalarValue::Binary(Some(bytes)) = result else { + panic!("expected a Binary result, got {result:?}"); + }; + assert!(bytes.len() > 4, "GeoTIFF should have content"); + assert!( + &bytes[0..2] == b"II" || &bytes[0..2] == b"MM", + "should be a valid TIFF header" + ); + } + + #[test] + fn as_geotiff_roundtrips_dimensions() { + // Export to GeoTIFF, reopen the bytes with GDAL, and confirm the raster + // dimensions survive the round trip. Uses only merged reader helpers + // (open + dataset_to_indb_raster) — no dependency on RS_FromGDALRaster. + with_gdal(|gdal| { + let arr = as_raster_array(test_raster_spec()); + let rasters = RasterStructArray::try_new(&arr).unwrap(); + let raster = rasters.get(0).unwrap(); + + let bytes = RsAsGeoTiff::raster_to_geotiff(gdal, &raster, None, None, None, None)?; + assert!(&bytes[0..2] == b"II" || &bytes[0..2] == b"MM"); + + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("roundtrip.tif"); + std::fs::write(&path, &bytes).unwrap(); + let dataset = gdal + .open_ex_with_options( + path.to_str().unwrap(), + DatasetOptions { + open_flags: GDAL_OF_RASTER | GDAL_OF_READONLY, + ..Default::default() + }, + ) + .map_err(crate::gdal_common::convert_gdal_err)?; + let roundtrip = crate::utils::dataset_to_indb_raster(&dataset)?; + let rt = RasterStructArray::try_new(&roundtrip).unwrap(); + let rt_raster = rt.get(0).unwrap(); + + assert_eq!(rt_raster.metadata().width(), raster.metadata().width()); + assert_eq!(rt_raster.metadata().height(), raster.metadata().height()); + assert_eq!(rt_raster.bands().len(), raster.bands().len()); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn as_geotiff_with_compression() { + // LZW and DEFLATE both produce valid, non-empty GeoTIFFs. + with_gdal(|gdal| { + let arr = as_raster_array(test_raster_spec()); + let rasters = RasterStructArray::try_new(&arr).unwrap(); + let raster = rasters.get(0).unwrap(); + + for comp in [CompressionType::Lzw, CompressionType::Deflate] { + let bytes = RsAsGeoTiff::raster_to_geotiff( + gdal, + &raster, + Some(comp), + Some(75.0), + None, + None, + )?; + assert!(!bytes.is_empty(), "{comp:?} GeoTIFF should have content"); + assert!(&bytes[0..2] == b"II" || &bytes[0..2] == b"MM"); + } + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } +} From dc7c2dbf077ff12644c1e2591fc061a125ac86da Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 13 Jul 2026 14:33:00 -0700 Subject: [PATCH 2/8] fix(rust/sedona-raster-gdal): validate RS_AsGeoTiff quality and pick the predictor by band type Review follow-ups: - quality is a 0.0-1.0 fraction (Sedona/GeoTools scale); values outside the range now error for every codec instead of silently clamping to maximum JPEG quality, and the docs state the scale with runnable 0.85 examples - PREDICTOR=3 (floating-point prediction) for float bands, 2 for integer - a drop guard unlinks the vsimem file on every exit path, including a failed create_copy, so failures no longer leak process-lifetime memory - hoist thread_local_provider once per batch, matching rs_metadata/ rs_polygonize - document NULL-option semantics, the Huffman/CCITTRLE 1-bit caveat, and the multiple-of-16 tile constraint; add a two-arg tile_size example - tests: tile overloads via reopened block size, per-row option broadcast, NULL raster, quality mapping/range, float predictor roundtrip, and band values compared in the existing roundtrip --- docs/reference/sql/rs_asgeotiff.qmd | 31 +- .../benches/rs_as_geotiff.rs | 3 +- rust/sedona-raster-gdal/src/rs_as_geotiff.rs | 391 ++++++++++++++++-- 3 files changed, 388 insertions(+), 37 deletions(-) diff --git a/docs/reference/sql/rs_asgeotiff.qmd b/docs/reference/sql/rs_asgeotiff.qmd index 99f9573a79..558b89f2d9 100644 --- a/docs/reference/sql/rs_asgeotiff.qmd +++ b/docs/reference/sql/rs_asgeotiff.qmd @@ -46,26 +46,35 @@ kernels: type: string description: > GeoTIFF compression codec (case-insensitive): None, PackBits, Deflate, - Huffman, LZW, or JPEG. Deflate and LZW additionally enable a horizontal - predictor to improve the compression ratio. + Huffman, LZW, or JPEG. Deflate and LZW additionally enable a predictor + to improve the compression ratio (horizontal differencing for integer + bands, floating-point prediction for float bands). Huffman maps to + CCITT RLE, which GDAL only accepts for 1-bit single-band data — it is + kept for Apache Sedona parity but errors on ordinary rasters. - name: quality type: double - description: Quality factor for lossy JPEG compression; ignored by the other codecs. + description: > + Quality factor for lossy JPEG compression as a fraction between 0.0 + and 1.0 (e.g. 0.75 for JPEG quality 75, matching Apache Sedona); + values outside that range raise an error. Ignored by the other codecs. - name: tile_width type: integer description: > - Tile width. When given together with tile_height, the GeoTIFF is written - tiled rather than stripped. + Tile width in pixels. When given together with tile_height, the + GeoTIFF is written tiled rather than stripped. TIFF tile dimensions + must be multiples of 16. - name: tile_height type: integer - description: Tile height. + description: Tile height in pixels; must be a multiple of 16. --- ## Description `RS_AsGeoTiff` encodes a raster as a GeoTIFF (via GDAL) and returns the file bytes as a binary value — the inverse of reading a GeoTIFF into a raster. The -result is `NULL` for a `NULL` raster. +result is `NULL` for a `NULL` raster. A `NULL` option argument means "not +specified": e.g. a `NULL` compression writes an uncompressed GeoTIFF rather +than propagating `NULL` to the result. The two-argument form `RS_AsGeoTiff(raster, tile_size)` writes a tiled GeoTIFF with square `tile_size` blocks. Compression is selected by name; a JPEG quality @@ -78,9 +87,13 @@ SELECT RS_AsGeoTiff(RS_Example()) IS NOT NULL; ``` ```sql -SELECT RS_AsGeoTiff(RS_Example(), 'DEFLATE', 6) IS NOT NULL; +SELECT RS_AsGeoTiff(RS_Example(), 16) IS NOT NULL; +``` + +```sql +SELECT RS_AsGeoTiff(RS_Example(), 'DEFLATE', 0.85) IS NOT NULL; ``` ```sql -SELECT RS_AsGeoTiff(RS_Example(), 'LZW', 75, 256, 256) IS NOT NULL; +SELECT RS_AsGeoTiff(RS_Example(), 'LZW', 0.85, 16, 16) IS NOT NULL; ``` diff --git a/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs b/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs index cc3ff8e9b2..93d85a0536 100644 --- a/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs +++ b/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs @@ -68,7 +68,8 @@ fn bench_rs_as_geotiff_compression(c: &mut Criterion) { ); let rasters = raster_array(32); - let quality = ColumnarValue::Scalar(ScalarValue::Float64(Some(75.0))); + // Quality is a 0.0-1.0 fraction (ignored by these codecs, but must be valid). + let quality = ColumnarValue::Scalar(ScalarValue::Float64(Some(0.75))); let mut group = c.benchmark_group("rs_as_geotiff_compression"); group.throughput(Throughput::Elements(rasters.len() as u64)); diff --git a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs index 1d5a22b5f4..f30ca86a4f 100644 --- a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs +++ b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs @@ -33,16 +33,20 @@ use arrow_schema::DataType; use datafusion_common::cast::{as_float64_array, as_string_array, as_uint32_array}; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; -use datafusion_common::{exec_datafusion_err, ScalarValue}; +use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue}; use datafusion_expr::{ColumnarValue, Volatility}; use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; 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 sedona_schema::raster::BandDataType; // Use thread-local provider to create GDAL datasets from `RasterRef`. -use crate::gdal_dataset_provider::configure_thread_local_options; +use crate::gdal_dataset_provider::{ + configure_thread_local_options, thread_local_provider, GDALDatasetProvider, +}; /// Counter for generating unique VSI memory file names static VSI_FILE_COUNTER: AtomicUsize = AtomicUsize::new(0); @@ -123,9 +127,10 @@ impl RsAsGeoTiff { Self { variant } } - /// Generate a unique VSI memory file path + /// Generate a unique VSI memory file path. `Relaxed` suffices: the counter + /// only has to hand out distinct values, no ordering with other memory. fn generate_vsi_path() -> String { - let counter = VSI_FILE_COUNTER.fetch_add(1, Ordering::SeqCst); + let counter = VSI_FILE_COUNTER.fetch_add(1, Ordering::Relaxed); let thread_id = std::thread::current().id(); format!("/vsimem/rs_as_geotiff_{:?}_{}.tif", thread_id, counter) } @@ -133,14 +138,13 @@ impl RsAsGeoTiff { /// Convert raster to GeoTiff bytes fn raster_to_geotiff( gdal: &sedona_gdal::gdal::Gdal, + provider: &GDALDatasetProvider, raster: &RasterRefImpl, compression: Option, quality: Option, tile_width: Option, tile_height: Option, ) -> Result> { - let provider = crate::gdal_dataset_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))?; @@ -150,6 +154,11 @@ impl RsAsGeoTiff { .get_driver_by_name("GTiff") .map_err(|e| exec_datafusion_err!("Failed to get GTiff driver: {}", e))?; + // Validate and map the quality up front so an out-of-range value errors + // for every codec, not only JPEG (the codecs that ignore quality should + // not silently accept nonsense either). + let jpeg_quality = quality.map(jpeg_quality_option).transpose()?; + // Build creation options as string list let mut options_list: Vec = Vec::new(); @@ -159,16 +168,18 @@ impl RsAsGeoTiff { // Add quality for JPEG if comp == CompressionType::Jpeg { - if let Some(q) = quality { - // JPEG quality is 1-100, we receive 0.0-1.0 - let jpeg_quality = (q * 100.0).round() as i32; - options_list.push(format!("JPEG_QUALITY={}", jpeg_quality.clamp(1, 100))); + if let Some(q) = jpeg_quality { + options_list.push(format!("JPEG_QUALITY={}", q)); } } - // Add predictor for Deflate/LZW (improves compression) + // Add a predictor for Deflate/LZW (improves compression): horizontal + // differencing (2) for integer samples, floating-point prediction (3) + // for float samples — predictor 2 on float data is legal but usually + // hurts the ratio. GTiff requires uniform band types, so the first + // band's type decides for the whole file. if comp == CompressionType::Deflate || comp == CompressionType::Lzw { - options_list.push("PREDICTOR=2".to_string()); + options_list.push(format!("PREDICTOR={}", predictor_for(raster)?)); } } @@ -182,28 +193,85 @@ impl RsAsGeoTiff { // Convert to creation options slice let options_refs: Vec<&str> = options_list.iter().map(|s| s.as_str()).collect(); - // Generate VSI path for output + // Output VSI path, unlinked on every exit path by the guard: without it + // a failed `create_copy` (invalid creation options, incompatible band + // layout, ...) can leave a partially written file in process-lifetime + // vsimem memory, accumulating across failures. let vsi_path = Self::generate_vsi_path(); + let guard = VsiMemFileGuard { + gdal, + path: &vsi_path, + }; - // Create copy to VSI memory file - let _output_dataset = source_dataset + // Create the copy in the VSI memory file. The returned dataset is + // dropped immediately (end of statement), which closes it and flushes + // the bytes to the vsimem file. + source_dataset .create_copy(&driver, &vsi_path, &options_refs) .map_err(|e| exec_datafusion_err!("Failed to create GeoTiff: {}", e))?; - // Close the output dataset to flush data - drop(_output_dataset); + // Read bytes from the VSI memory file; the guard cleans up. + let bytes = gdal + .get_vsi_mem_file_bytes_owned(&vsi_path) + .map_err(|e| exec_datafusion_err!("Failed to read GeoTiff bytes: {}", e))?; + + drop(guard); + Ok(bytes) + } +} - // Read bytes from VSI memory file and clean up - let bytes = gdal.get_vsi_mem_file_bytes_owned(&vsi_path).map_err(|e| { - let _ = gdal.unlink_mem_file(&vsi_path); - exec_datafusion_err!("Failed to read GeoTiff bytes: {}", e) - })?; +/// Unlinks a vsimem file when dropped, so every exit path of +/// [`RsAsGeoTiff::raster_to_geotiff`] — including failed `create_copy` — +/// releases the process-lifetime vsimem allocation. +struct VsiMemFileGuard<'a> { + gdal: &'a sedona_gdal::gdal::Gdal, + path: &'a str, +} - // Clean up VSI file - let _ = gdal.unlink_mem_file(&vsi_path); +impl Drop for VsiMemFileGuard<'_> { + fn drop(&mut self) { + // Unlinking a file that create_copy never managed to create is a no-op + // error, which is fine to ignore. + let _ = self.gdal.unlink_mem_file(self.path); + } +} - Ok(bytes) +/// Map a quality fraction in `[0.0, 1.0]` to GDAL's 1–100 `JPEG_QUALITY`. +/// +/// The fractional scale matches Apache Sedona (GeoTools' `setCompressionQuality`); +/// a value outside the range errors rather than clamping — silently clamping +/// would turn the most likely mistake (passing a 0–100 quality like `75`) into +/// maximum quality with no warning. +fn jpeg_quality_option(quality: f64) -> Result { + if !(0.0..=1.0).contains(&quality) { + return exec_err!( + "RS_AsGeoTiff: quality must be a fraction between 0.0 and 1.0 (got {quality}); \ + e.g. use 0.75 for JPEG quality 75" + ); } + // Round to 1-100; GDAL rejects 0, so 0.0 maps to the minimum quality 1. + Ok(((quality * 100.0).round() as i32).max(1)) +} + +/// TIFF predictor for Deflate/LZW: 3 (floating-point prediction) for float +/// bands, 2 (horizontal differencing) for integer bands. Decided by the first +/// band's sample type; GTiff creation requires uniform band types anyway. +fn predictor_for(raster: &RasterRefImpl) -> Result { + let bands = raster.bands(); + if bands.is_empty() { + return Ok(2); + } + let band = bands + .band(1) + .map_err(|e| exec_datafusion_err!("RS_AsGeoTiff: {e}"))?; + let data_type = band + .metadata() + .data_type() + .map_err(|e| exec_datafusion_err!("RS_AsGeoTiff: {e}"))?; + Ok(match data_type { + BandDataType::Float32 | BandDataType::Float64 => 3, + _ => 2, + }) } impl SedonaScalarKernel for RsAsGeoTiff { @@ -348,6 +416,8 @@ impl SedonaScalarKernel for RsAsGeoTiff { 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}"))?; executor.execute_raster_void(|_i, raster_opt| { let compression_opt = compression_iter.next().unwrap(); let quality_opt = quality_iter.next().unwrap(); @@ -378,6 +448,7 @@ impl SedonaScalarKernel for RsAsGeoTiff { let bytes = Self::raster_to_geotiff( gdal, + &provider, raster, compression, quality, @@ -397,6 +468,7 @@ impl SedonaScalarKernel for RsAsGeoTiff { #[cfg(test)] mod tests { use super::*; + use arrow_array::Array; use datafusion_expr::ScalarUDF; use sedona_gdal::gdal_dyn_bindgen::{GDAL_OF_RASTER, GDAL_OF_READONLY}; use sedona_gdal::raster::types::DatasetOptions; @@ -479,7 +551,9 @@ mod tests { let rasters = RasterStructArray::try_new(&arr).unwrap(); let raster = rasters.get(0).unwrap(); - let bytes = RsAsGeoTiff::raster_to_geotiff(gdal, &raster, None, None, None, None)?; + let provider = thread_local_provider(gdal).unwrap(); + let bytes = + RsAsGeoTiff::raster_to_geotiff(gdal, &provider, &raster, None, None, None, None)?; assert!(&bytes[0..2] == b"II" || &bytes[0..2] == b"MM"); let tmp = tempfile::tempdir().unwrap(); @@ -501,6 +575,267 @@ mod tests { assert_eq!(rt_raster.metadata().width(), raster.metadata().width()); assert_eq!(rt_raster.metadata().height(), raster.metadata().height()); assert_eq!(rt_raster.bands().len(), raster.bands().len()); + // Pixel values must survive too — this would catch predictor or + // compression corruption that dimension checks cannot. + assert_eq!( + rt_raster + .bands() + .band(1) + .unwrap() + .nd_buffer() + .unwrap() + .as_contiguous() + .unwrap(), + raster + .bands() + .band(1) + .unwrap() + .nd_buffer() + .unwrap() + .as_contiguous() + .unwrap(), + ); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn jpeg_quality_maps_fraction_to_1_100() { + // The subtle bit of the options plumbing: a 0.0-1.0 fraction (Sedona / + // GeoTools scale) maps to GDAL's 1-100 JPEG_QUALITY. + assert_eq!(jpeg_quality_option(0.85).unwrap(), 85); + assert_eq!(jpeg_quality_option(1.0).unwrap(), 100); + // GDAL rejects 0, so the bottom of the range maps to minimum quality 1. + assert_eq!(jpeg_quality_option(0.0).unwrap(), 1); + assert_eq!(jpeg_quality_option(0.004).unwrap(), 1); + } + + #[test] + fn jpeg_quality_out_of_range_errors() { + // A 0-100 style quality (the likely mistake) errors instead of clamping + // to maximum quality silently. + for q in [75.0, -0.1, 1.01, f64::NAN] { + let err = jpeg_quality_option(q).unwrap_err().to_string(); + assert!( + err.contains("between 0.0 and 1.0"), + "unexpected error for {q}: {err}" + ); + } + } + + #[test] + fn out_of_range_quality_errors_for_every_codec() { + // The validation is not JPEG-only: LZW ignores quality, but nonsense + // still errors rather than being silently dropped. + with_gdal(|gdal| { + let arr = as_raster_array(test_raster_spec()); + let rasters = RasterStructArray::try_new(&arr).unwrap(); + let raster = rasters.get(0).unwrap(); + let provider = thread_local_provider(gdal).unwrap(); + let err = RsAsGeoTiff::raster_to_geotiff( + gdal, + &provider, + &raster, + Some(CompressionType::Lzw), + Some(75.0), + None, + None, + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("between 0.0 and 1.0"), + "unexpected error: {err}" + ); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn predictor_matches_band_type() { + // Horizontal differencing (2) for integer bands, floating-point + // prediction (3) for float bands. + let int_arr = as_raster_array(test_raster_spec()); + let int_rasters = RasterStructArray::try_new(&int_arr).unwrap(); + assert_eq!(predictor_for(&int_rasters.get(0).unwrap()).unwrap(), 2); + + let float_arr = as_raster_array( + RasterSpec::d2(3, 3) + .transform([0.0, 1.0, 0.0, 3.0, 0.0, -1.0]) + .band_values(&[1.5f32, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5, 9.5]), + ); + let float_rasters = RasterStructArray::try_new(&float_arr).unwrap(); + assert_eq!(predictor_for(&float_rasters.get(0).unwrap()).unwrap(), 3); + } + + #[test] + fn float_band_exports_with_float_predictor() { + // End-to-end: a Float32 band under DEFLATE goes out with PREDICTOR=3 + // and the values survive a roundtrip (a wrong predictor that libtiff + // rejects, or value corruption, would fail here). + with_gdal(|gdal| { + let spec = RasterSpec::d2(3, 3) + .transform([0.0, 1.0, 0.0, 3.0, 0.0, -1.0]) + .band_values(&[1.5f32, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5, 9.5]); + let arr = as_raster_array(spec); + let rasters = RasterStructArray::try_new(&arr).unwrap(); + let raster = rasters.get(0).unwrap(); + let provider = thread_local_provider(gdal).unwrap(); + + let bytes = RsAsGeoTiff::raster_to_geotiff( + gdal, + &provider, + &raster, + Some(CompressionType::Deflate), + None, + None, + None, + )?; + + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("float_predictor.tif"); + std::fs::write(&path, &bytes).unwrap(); + let dataset = gdal + .open_ex_with_options( + path.to_str().unwrap(), + DatasetOptions { + open_flags: GDAL_OF_RASTER | GDAL_OF_READONLY, + ..Default::default() + }, + ) + .map_err(crate::gdal_common::convert_gdal_err)?; + let roundtrip = crate::utils::dataset_to_indb_raster(&dataset)?; + let rt = RasterStructArray::try_new(&roundtrip).unwrap(); + let rt_raster = rt.get(0).unwrap(); + assert_eq!( + rt_raster + .bands() + .band(1) + .unwrap() + .nd_buffer() + .unwrap() + .as_contiguous() + .unwrap(), + raster + .bands() + .band(1) + .unwrap() + .nd_buffer() + .unwrap() + .as_contiguous() + .unwrap(), + ); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + /// Reopen exported GeoTIFF bytes and return band 1's (block_x, block_y). + fn reopened_block_size(gdal: &sedona_gdal::gdal::Gdal, bytes: &[u8]) -> (usize, usize) { + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("block_size.tif"); + std::fs::write(&path, bytes).unwrap(); + let dataset = gdal + .open_ex_with_options( + path.to_str().unwrap(), + DatasetOptions { + open_flags: GDAL_OF_RASTER | GDAL_OF_READONLY, + ..Default::default() + }, + ) + .unwrap(); + dataset.rasterband(1).unwrap().block_size() + } + + #[test] + fn tile_options_survive_export() { + // The tiling plumbing end to end: TILED=YES + BLOCKXSIZE/BLOCKYSIZE + // reach GDAL, and a reopened dataset exposes them as the block size. + // TIFF tile dimensions must be multiples of 16. + with_gdal(|gdal| { + let arr = as_raster_array(test_raster_spec()); + let rasters = RasterStructArray::try_new(&arr).unwrap(); + let raster = rasters.get(0).unwrap(); + let provider = thread_local_provider(gdal).unwrap(); + + let bytes = RsAsGeoTiff::raster_to_geotiff( + gdal, + &provider, + &raster, + None, + None, + Some(16), + Some(32), + )?; + assert_eq!(reopened_block_size(gdal, &bytes), (16, 32)); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn tile_size_udf_overload_tiles_squares() { + // RS_AsGeoTiff(raster, tile_size) writes square tiles of that size. + let udf: ScalarUDF = rs_as_geotiff_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); + let result = tester + .invoke_arrays(vec![ + Arc::new(as_raster_array(test_raster_spec())) as arrow_array::ArrayRef, + Arc::new(arrow_array::Int32Array::from(vec![Some(16)])), + ]) + .unwrap(); + let binary = result + .as_any() + .downcast_ref::() + .unwrap(); + with_gdal(|gdal| { + assert_eq!(reopened_block_size(gdal, binary.value(0)), (16, 16)); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + + #[test] + fn null_raster_yields_null_output() { + let udf: ScalarUDF = rs_as_geotiff_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![RASTER]); + let rasters = sedona_testing::rasters::generate_test_rasters(1, Some(0)).unwrap(); + let result = tester + .invoke_arrays(vec![Arc::new(rasters) as arrow_array::ArrayRef]) + .unwrap(); + let binary = result + .as_any() + .downcast_ref::() + .unwrap(); + assert!(binary.is_null(0), "NULL raster should export as NULL"); + } + + #[test] + fn per_row_options_broadcast() { + // An array raster with a per-row tile_size column: each row's option + // reaches its own GDAL export (the broadcast path through the upfront + // into_array casts). + let udf: ScalarUDF = rs_as_geotiff_udf().into(); + let tester = ScalarUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int32)]); + + // Two raster rows, tiled 16 and 32 respectively. + let rasters = sedona_testing::rasters::generate_test_rasters(2, None).unwrap(); + let tiles = arrow_array::Int32Array::from(vec![Some(16), Some(32)]); + let result = tester + .invoke_arrays(vec![ + Arc::new(rasters) as arrow_array::ArrayRef, + Arc::new(tiles), + ]) + .unwrap(); + let binary = result + .as_any() + .downcast_ref::() + .unwrap(); + with_gdal(|gdal| { + assert_eq!(reopened_block_size(gdal, binary.value(0)), (16, 16)); + assert_eq!(reopened_block_size(gdal, binary.value(1)), (32, 32)); Ok::<_, datafusion_common::DataFusionError>(()) }) .unwrap(); @@ -514,12 +849,14 @@ mod tests { let rasters = RasterStructArray::try_new(&arr).unwrap(); let raster = rasters.get(0).unwrap(); + let provider = thread_local_provider(gdal).unwrap(); for comp in [CompressionType::Lzw, CompressionType::Deflate] { let bytes = RsAsGeoTiff::raster_to_geotiff( gdal, + &provider, &raster, Some(comp), - Some(75.0), + Some(0.75), None, None, )?; From aece9ce1781e85610793d25de012a98476475e9e Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 13 Jul 2026 15:51:13 -0700 Subject: [PATCH 3/8] test(python/sedonadb): add RS_AsGeoTiff smoke tests Guards registration in the Python build and the GeoTIFF byte output for each overload shape; depth stays in the Rust tests. --- .../tests/functions/test_raster_functions.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/python/sedonadb/tests/functions/test_raster_functions.py b/python/sedonadb/tests/functions/test_raster_functions.py index 4e91b13301..422aec832f 100644 --- a/python/sedonadb/tests/functions/test_raster_functions.py +++ b/python/sedonadb/tests/functions/test_raster_functions.py @@ -268,3 +268,27 @@ def test_rs_setbandnodatavalue_two_arg_requires_single_band(): SedonaDB().assert_query_result( "SELECT RS_SetBandNoDataValue(RS_Example(), 0)", None ) + + +# RS_AsGeoTiff smoke coverage: the GDAL-backed export is tested in depth in +# Rust (rust/sedona-raster-gdal/src/rs_as_geotiff.rs); these only guard that +# the function is registered in the Python build and returns GeoTIFF bytes. +@pytest.mark.parametrize( + "expr", + [ + "RS_AsGeoTiff(RS_Example())", + "RS_AsGeoTiff(RS_Example(), 16)", + "RS_AsGeoTiff(RS_Example(), 'DEFLATE', 0.85)", + "RS_AsGeoTiff(RS_Example(), 'LZW', 0.85, 16, 16)", + ], +) +def test_rs_asgeotiff_returns_tiff_bytes(con, expr): + result = con.sql(f"SELECT {expr} AS t").to_arrow_table()["t"][0].as_py() + assert result[:2] in (b"II", b"MM"), "should start with a TIFF byte-order mark" + + +def test_rs_asgeotiff_out_of_range_quality_errors(con): + # Quality is a 0.0-1.0 fraction; a 0-100 style value errors rather than + # silently clamping to maximum quality. + with pytest.raises(Exception, match="between 0.0 and 1.0"): + con.sql("SELECT RS_AsGeoTiff(RS_Example(), 'JPEG', 75)").to_arrow_table() From 216df3da7949aa41a671a479ce6f09c6dfb4d53d Mon Sep 17 00:00:00 2001 From: jameswillis Date: Tue, 14 Jul 2026 10:02:04 -0700 Subject: [PATCH 4/8] perf(rust/sedona-raster-gdal): return RS_AsGeoTiff bytes without an intermediate Vec copy Review follow-ups: - seize the vsimem buffer via a new Gdal::get_vsi_mem_file_buffer_owned (zero-copy VSIBuffer, freed on drop); the only remaining byte copy is the append into the Binary output builder - mark the docs page experimental so the option handling can evolve - test the failed-create_copy error path (CCITTRLE on an ordinary raster) - rasterio cross-check: decode the exported bytes for every band data type x {uncompressed, DEFLATE, LZW} and compare contents, dtype, and geotransform exactly --- c/sedona-gdal/src/gdal.rs | 9 ++++ c/sedona-gdal/src/vsi.rs | 1 + docs/reference/sql/rs_asgeotiff.qmd | 4 ++ .../tests/functions/test_raster_functions.py | 35 +++++++++++++++ rust/sedona-raster-gdal/src/rs_as_geotiff.rs | 44 +++++++++++++++++-- 5 files changed, 89 insertions(+), 4 deletions(-) diff --git a/c/sedona-gdal/src/gdal.rs b/c/sedona-gdal/src/gdal.rs index fe2a35f868..8db64c1f3e 100644 --- a/c/sedona-gdal/src/gdal.rs +++ b/c/sedona-gdal/src/gdal.rs @@ -171,6 +171,15 @@ impl Gdal { vsi::get_vsi_mem_file_bytes_owned(self.api, file_name) } + /// Take ownership of a VSI in-memory file's buffer without copying: the + /// file is unlinked and its GDAL-allocated bytes are returned as a + /// [`vsi::VSIBuffer`] (freed on drop). Prefer this over + /// [`Self::get_vsi_mem_file_bytes_owned`] when the bytes are consumed as a + /// slice — it skips the `Vec` copy. + pub fn get_vsi_mem_file_buffer_owned(&self, file_name: &str) -> Result { + vsi::get_vsi_mem_file_buffer_owned(self.api, file_name) + } + // -- Raster operations --------------------------------------------------- /// Create a bare in-memory MEM dataset with GDAL-owned bands. diff --git a/c/sedona-gdal/src/vsi.rs b/c/sedona-gdal/src/vsi.rs index f1684eebd6..2c7ad59a70 100644 --- a/c/sedona-gdal/src/vsi.rs +++ b/c/sedona-gdal/src/vsi.rs @@ -28,6 +28,7 @@ use crate::errors::{GdalError, Result}; use crate::gdal_api::{call_gdal_api, GdalApi}; /// An owned GDAL-allocated VSI memory buffer. +#[derive(Debug)] pub struct VSIBuffer { api: &'static GdalApi, ptr: *mut u8, diff --git a/docs/reference/sql/rs_asgeotiff.qmd b/docs/reference/sql/rs_asgeotiff.qmd index 558b89f2d9..eb77faef3e 100644 --- a/docs/reference/sql/rs_asgeotiff.qmd +++ b/docs/reference/sql/rs_asgeotiff.qmd @@ -68,6 +68,10 @@ kernels: description: Tile height in pixels; must be a multiple of 16. --- +::: callout-warning +**Experimental.** This function is experimental; its behavior may change without notice. +::: + ## Description `RS_AsGeoTiff` encodes a raster as a GeoTIFF (via GDAL) and returns the file diff --git a/python/sedonadb/tests/functions/test_raster_functions.py b/python/sedonadb/tests/functions/test_raster_functions.py index 422aec832f..585aeef987 100644 --- a/python/sedonadb/tests/functions/test_raster_functions.py +++ b/python/sedonadb/tests/functions/test_raster_functions.py @@ -292,3 +292,38 @@ def test_rs_asgeotiff_out_of_range_quality_errors(con): # silently clamping to maximum quality. with pytest.raises(Exception, match="between 0.0 and 1.0"): con.sql("SELECT RS_AsGeoTiff(RS_Example(), 'JPEG', 75)").to_arrow_table() + + +# Cross-check RS_AsGeoTiff against rasterio: export a random raster of each +# band data type and confirm rasterio decodes the bytes back to the identical +# array, dtype, and geotransform. The DEFLATE/LZW variants also exercise the +# per-dtype predictor selection (horizontal differencing for integers, +# floating-point prediction for float bands). +@pytest.mark.parametrize("dtype", ["uint8", "uint16", "int32", "float32", "float64"]) +@pytest.mark.parametrize("compression_args", ["", ", 'DEFLATE', 0.85", ", 'LZW', 0.85"]) +def test_rs_asgeotiff_roundtrips_contents(con, dtype, compression_args): + import numpy as np + + pytest.importorskip("rasterio") + from rasterio.io import MemoryFile + + from sedonadb.raster import Raster + + rng = np.random.default_rng(7) + data = (rng.random((5, 4)) * 100).astype(dtype) + gdal_transform = (10.0, 1.0, 0.0, 20.0, 0.0, -1.0) + raster = Raster.from_numpy(data, transform=gdal_transform) + + tiff_bytes = ( + con.sql( + f"SELECT RS_AsGeoTiff($1{compression_args}) AS t", + params=(raster,), + ) + .to_arrow_table()["t"] + .to_pylist()[0] + ) + with MemoryFile(bytes(tiff_bytes)) as mem, mem.open() as src: + decoded = src.read(1) + assert decoded.dtype == data.dtype + np.testing.assert_array_equal(decoded, data) + assert src.transform.to_gdal() == gdal_transform diff --git a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs index f30ca86a4f..824aeaaee7 100644 --- a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs +++ b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs @@ -36,6 +36,7 @@ use datafusion_common::error::Result; use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue}; use datafusion_expr::{ColumnarValue, Volatility}; use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_gdal::vsi::VSIBuffer; use sedona_raster::array::RasterRefImpl; use sedona_raster::traits::RasterRef; use sedona_raster_functions::RasterExecutor; @@ -144,7 +145,7 @@ impl RsAsGeoTiff { quality: Option, tile_width: Option, tile_height: Option, - ) -> Result> { + ) -> Result { let raster_ds = provider .raster_ref_to_gdal(raster) .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {}", e))?; @@ -210,9 +211,13 @@ impl RsAsGeoTiff { .create_copy(&driver, &vsi_path, &options_refs) .map_err(|e| exec_datafusion_err!("Failed to create GeoTiff: {}", e))?; - // Read bytes from the VSI memory file; the guard cleans up. + // Seize the vsimem file's buffer without copying: `VSIBuffer` owns the + // GDAL allocation (freed on drop) and unlinks the file, so the only + // byte copy left is the append into the output builder. The guard's + // unlink becomes a no-op on this path but still cleans up when + // `create_copy` or the seize fails. let bytes = gdal - .get_vsi_mem_file_bytes_owned(&vsi_path) + .get_vsi_mem_file_buffer_owned(&vsi_path) .map_err(|e| exec_datafusion_err!("Failed to read GeoTiff bytes: {}", e))?; drop(guard); @@ -455,7 +460,7 @@ impl SedonaScalarKernel for RsAsGeoTiff { tile_width, tile_height, )?; - builder.append_value(&bytes); + builder.append_value(bytes.as_ref()); Ok(()) })?; @@ -624,6 +629,37 @@ mod tests { } } + #[test] + fn failed_create_copy_surfaces_gdal_error() { + // A creation-option combination GDAL rejects at CreateCopy time — + // CCITTRLE (Huffman) only accepts 1-bit single-band data, so an + // ordinary UInt8 raster fails. The error must surface (and the vsimem + // guard cleans up the partial file rather than leaking it). + with_gdal(|gdal| { + let arr = as_raster_array(test_raster_spec()); + let rasters = RasterStructArray::try_new(&arr).unwrap(); + let raster = rasters.get(0).unwrap(); + let provider = thread_local_provider(gdal).unwrap(); + let err = RsAsGeoTiff::raster_to_geotiff( + gdal, + &provider, + &raster, + Some(CompressionType::Huffman), + None, + None, + None, + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("Failed to create GeoTiff"), + "unexpected error: {err}" + ); + Ok::<_, datafusion_common::DataFusionError>(()) + }) + .unwrap(); + } + #[test] fn out_of_range_quality_errors_for_every_codec() { // The validation is not JPEG-only: LZW ignores quality, but nonsense From b3743bfab87101ac5e55b496aeee58d268a71684 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Wed, 15 Jul 2026 12:35:53 -0700 Subject: [PATCH 5/8] feat(rust/sedona-raster-gdal): RS_AsGeoTiff returns BinaryView; UDF-level unknown-compression test BinaryView output keeps the door open for eliminating the remaining vsimem-to-builder copy by wrapping the VSI buffer as an external Arrow allocation. The new test pins the compression-string parse error as surfaced through the UDF rather than only via CompressionType::parse. --- rust/sedona-raster-gdal/src/rs_as_geotiff.rs | 39 ++++++++++++++++---- 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs index 824aeaaee7..0ef308b50b 100644 --- a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs +++ b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs @@ -28,7 +28,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use crate::gdal_common::with_gdal; -use arrow_array::builder::BinaryBuilder; +use arrow_array::builder::BinaryViewBuilder; use arrow_schema::DataType; use datafusion_common::cast::{as_float64_array, as_string_array, as_uint32_array}; use datafusion_common::config::ConfigOptions; @@ -307,7 +307,7 @@ impl SedonaScalarKernel for RsAsGeoTiff { ], }; - let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(DataType::Binary)); + let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(DataType::BinaryView)); matcher.match_args(args) } @@ -417,7 +417,7 @@ impl SedonaScalarKernel for RsAsGeoTiff { let mut tile_height_iter = tile_height_array.iter(); // Build output binary array - let mut builder = BinaryBuilder::with_capacity(num_iterations, num_iterations * 1024); + let mut builder = BinaryViewBuilder::with_capacity(num_iterations); with_gdal(|gdal| { configure_thread_local_options(gdal, config_options)?; @@ -536,8 +536,8 @@ mod tests { let udf: ScalarUDF = rs_as_geotiff_udf().into(); let tester = ScalarUdfTester::new(udf, vec![RASTER]); let result = tester.invoke_scalar(test_raster_spec()).unwrap(); - let ScalarValue::Binary(Some(bytes)) = result else { - panic!("expected a Binary result, got {result:?}"); + let ScalarValue::BinaryView(Some(bytes)) = result else { + panic!("expected a BinaryView result, got {result:?}"); }; assert!(bytes.len() > 4, "GeoTIFF should have content"); assert!( @@ -546,6 +546,29 @@ mod tests { ); } + #[test] + fn unknown_compression_type_errors_through_udf() { + // The compression-string parse error surfaces through the UDF itself, + // not only from CompressionType::parse in isolation. + let udf: ScalarUDF = rs_as_geotiff_udf().into(); + let tester = ScalarUdfTester::new( + udf, + vec![ + RASTER, + SedonaType::Arrow(DataType::Utf8), + SedonaType::Arrow(DataType::Float64), + ], + ); + let err = tester + .invoke_scalar_scalar_scalar(test_raster_spec(), "GZIP", 0.5) + .unwrap_err() + .to_string(); + assert!( + err.contains("Unknown compression type: GZIP"), + "unexpected error: {err}" + ); + } + #[test] fn as_geotiff_roundtrips_dimensions() { // Export to GeoTIFF, reopen the bytes with GDAL, and confirm the raster @@ -824,7 +847,7 @@ mod tests { .unwrap(); let binary = result .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); with_gdal(|gdal| { assert_eq!(reopened_block_size(gdal, binary.value(0)), (16, 16)); @@ -843,7 +866,7 @@ mod tests { .unwrap(); let binary = result .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); assert!(binary.is_null(0), "NULL raster should export as NULL"); } @@ -867,7 +890,7 @@ mod tests { .unwrap(); let binary = result .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); with_gdal(|gdal| { assert_eq!(reopened_block_size(gdal, binary.value(0)), (16, 16)); From bb195dba086cb08dddbd91f53b8daef709701903 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Wed, 15 Jul 2026 12:35:53 -0700 Subject: [PATCH 6/8] test(python/sedonadb): hoist non-rasterio imports to module top numpy/pandas/Raster now import once at the top of test_raster_functions.py; rasterio imports stay function-local behind importorskip. --- .../tests/functions/test_raster_functions.py | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/python/sedonadb/tests/functions/test_raster_functions.py b/python/sedonadb/tests/functions/test_raster_functions.py index e33880eedf..631e68450f 100644 --- a/python/sedonadb/tests/functions/test_raster_functions.py +++ b/python/sedonadb/tests/functions/test_raster_functions.py @@ -15,8 +15,9 @@ # specific language governing permissions and limitations # under the License. -import pytest import numpy as np +import pandas as pd +import pytest from sedonadb.testing import SedonaDB from sedonadb.raster import Raster @@ -149,15 +150,11 @@ def test_rs_value_matches_rasterio(con): positions per pixel (toward the corners, kept inside the pixel to avoid floor ambiguity at exact boundaries) and a batch of random interior points. """ - import numpy as np - import pandas as pd pytest.importorskip("rasterio") from rasterio.io import MemoryFile from rasterio.transform import Affine - from sedonadb.raster import Raster - rng = np.random.default_rng(42) height, width = 7, 5 data = rng.random((height, width)) * 1000.0 @@ -315,13 +312,9 @@ def test_rs_asgeotiff_out_of_range_quality_errors(con): @pytest.mark.parametrize("dtype", ["uint8", "uint16", "int32", "float32", "float64"]) @pytest.mark.parametrize("compression_args", ["", ", 'DEFLATE', 0.85", ", 'LZW', 0.85"]) def test_rs_asgeotiff_roundtrips_contents(con, dtype, compression_args): - import numpy as np - pytest.importorskip("rasterio") from rasterio.io import MemoryFile - from sedonadb.raster import Raster - rng = np.random.default_rng(7) data = (rng.random((5, 4)) * 100).astype(dtype) gdal_transform = (10.0, 1.0, 0.0, 20.0, 0.0, -1.0) @@ -407,8 +400,6 @@ def test_rs_as_raster_matches_rasterio( def test_rs_as_raster_all_touched_changes_pixels(con, sedona_testing): - import numpy as np - pytest.importorskip("rasterio") from rasterio.features import rasterize from rasterio.transform import Affine @@ -485,8 +476,6 @@ def test_rs_as_raster_rejects_fractional_integer_nodata(con, sedona_testing): def test_rs_as_raster_sets_output_nodata(con, sedona_testing): - import numpy as np - path = sedona_testing / "data/raster/test4.tiff" tab = con.sql( """ From 15cbc8bf4138f4dd2e0df98f9bc428b31e5f8275 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Wed, 15 Jul 2026 13:07:26 -0700 Subject: [PATCH 7/8] perf(rust/sedona-raster-gdal): RS_AsGeoTiff output views the VSI allocation directly The encoded GeoTIFF bytes are no longer copied into the output array: each row's VSIBuffer is wrapped as an external Arrow allocation (Buffer::from_custom_allocation) and appended as a view block, so the GDAL allocation is freed when the last Arrow reference drops. GeoTIFFs above the 4 GiB binary-view element limit error instead of truncating. rs_as_geotiff bench, new 16 MiB uncompressed case: 3.21 ms -> 2.08 ms (4.86 -> 7.52 GiB/s), -35%. --- .../benches/rs_as_geotiff.rs | 16 ++++++++ rust/sedona-raster-gdal/src/rs_as_geotiff.rs | 37 ++++++++++++++++++- 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs b/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs index 93d85a0536..1478f70059 100644 --- a/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs +++ b/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs @@ -21,6 +21,7 @@ //! fixtures): the uncompressed default and the compression axis (none/lzw/deflate). use std::hint::black_box; +use std::sync::Arc; use arrow_array::ArrayRef; use arrow_schema::DataType; @@ -28,6 +29,8 @@ use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Through use datafusion_common::ScalarValue; use datafusion_expr::{ColumnarValue, ScalarUDF}; use sedona_schema::datatypes::{SedonaType, RASTER}; +use sedona_schema::raster::BandDataType; +use sedona_testing::raster_spec::RasterSpec; use sedona_testing::testers::ScalarUdfTester; fn raster_array(rows: usize) -> ArrayRef { @@ -53,6 +56,19 @@ fn bench_rs_as_geotiff_basic(c: &mut Criterion) { b.iter(|| black_box(tester.invoke_arrays(vec![input.clone()]).unwrap())) }); } + + // A single ~16 MB raster: large enough that per-row byte handling + // (vsimem buffer -> output array) registers against the GDAL encode. + const LARGE: i64 = 4096; + let large: ArrayRef = Arc::new( + RasterSpec::d2(LARGE, LARGE) + .band(BandDataType::UInt8) + .build(), + ); + group.throughput(Throughput::Bytes((LARGE * LARGE) as u64)); + group.bench_with_input(BenchmarkId::new("large", "16MiB"), &large, |b, input| { + b.iter(|| black_box(tester.invoke_arrays(vec![input.clone()]).unwrap())) + }); group.finish(); } diff --git a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs index 0ef308b50b..6ecf345bfc 100644 --- a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs +++ b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs @@ -24,11 +24,13 @@ //! - RS_AsGeoTiff(raster, compressionType, imageQuality, tileSize) //! - RS_AsGeoTiff(raster, compressionType, imageQuality, tileWidth, tileHeight) +use std::ptr::NonNull; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use crate::gdal_common::with_gdal; use arrow_array::builder::BinaryViewBuilder; +use arrow_buffer::Buffer; use arrow_schema::DataType; use datafusion_common::cast::{as_float64_array, as_string_array, as_uint32_array}; use datafusion_common::config::ConfigOptions; @@ -247,6 +249,36 @@ impl Drop for VsiMemFileGuard<'_> { /// a value outside the range errors rather than clamping — silently clamping /// would turn the most likely mistake (passing a 0–100 quality like `75`) into /// maximum quality with no warning. +/// Append one encoded GeoTIFF to the output as a view over the GDAL +/// allocation itself — the bytes are not copied. The [`VSIBuffer`] becomes an +/// external Arrow allocation owned by the output array, so its lifetime (and +/// `VSIFree`) follows the array rather than this call. +fn append_geotiff_view(builder: &mut BinaryViewBuilder, bytes: VSIBuffer) -> Result<()> { + let len = bytes.len(); + // A binary-view element addresses at most u32::MAX bytes. + let Ok(view_len) = u32::try_from(len) else { + return exec_err!("RS_AsGeoTiff: {len}-byte GeoTIFF exceeds the 4 GiB binary output limit"); + }; + // Views this small are stored inline in the view struct; wrapping the + // allocation would save nothing (and a zero-length buffer has no pointer). + if len <= 12 { + builder.append_value(bytes.as_ref()); + return Ok(()); + } + let Some(ptr) = NonNull::new(bytes.as_ref().as_ptr() as *mut u8) else { + return exec_err!("RS_AsGeoTiff: GDAL returned a null buffer for a {len}-byte GeoTIFF"); + }; + // SAFETY: `ptr`/`len` describe exactly the allocation owned by `bytes`, + // which stores a raw pointer (the bytes never move) and stays alive inside + // the Arc — freeing via `VSIFree` — until the last Arrow reference to + // `buffer` drops. + let buffer = unsafe { Buffer::from_custom_allocation(ptr, len, Arc::new(bytes)) }; + let block = builder.append_block(buffer); + builder + .try_append_view(block, 0, view_len) + .map_err(|e| exec_datafusion_err!("RS_AsGeoTiff: failed to append binary view: {e}")) +} + fn jpeg_quality_option(quality: f64) -> Result { if !(0.0..=1.0).contains(&quality) { return exec_err!( @@ -460,7 +492,7 @@ impl SedonaScalarKernel for RsAsGeoTiff { tile_width, tile_height, )?; - builder.append_value(bytes.as_ref()); + append_geotiff_view(&mut builder, bytes)?; Ok(()) })?; @@ -892,6 +924,9 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); + // Each row's bytes are a view over its own GDAL allocation (one + // variadic buffer per row), not a copy into a shared builder buffer. + assert_eq!(binary.data_buffers().len(), 2); with_gdal(|gdal| { assert_eq!(reopened_block_size(gdal, binary.value(0)), (16, 16)); assert_eq!(reopened_block_size(gdal, binary.value(1)), (32, 32)); From 3b738e533ef137b43700ad91ffaad7c640a8a86d Mon Sep 17 00:00:00 2001 From: jameswillis Date: Thu, 16 Jul 2026 23:41:28 -0700 Subject: [PATCH 8/8] docs(rust/sedona-raster-gdal): move misplaced jpeg_quality_option doc comment The doc comment describing the quality-fraction to JPEG_QUALITY mapping sat above append_geotiff_view (which had its own doc) while jpeg_quality_option itself was undocumented. Relocate it to the function it describes. --- rust/sedona-raster-gdal/src/rs_as_geotiff.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs index 6ecf345bfc..cffcc033d5 100644 --- a/rust/sedona-raster-gdal/src/rs_as_geotiff.rs +++ b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs @@ -243,12 +243,6 @@ impl Drop for VsiMemFileGuard<'_> { } } -/// Map a quality fraction in `[0.0, 1.0]` to GDAL's 1–100 `JPEG_QUALITY`. -/// -/// The fractional scale matches Apache Sedona (GeoTools' `setCompressionQuality`); -/// a value outside the range errors rather than clamping — silently clamping -/// would turn the most likely mistake (passing a 0–100 quality like `75`) into -/// maximum quality with no warning. /// Append one encoded GeoTIFF to the output as a view over the GDAL /// allocation itself — the bytes are not copied. The [`VSIBuffer`] becomes an /// external Arrow allocation owned by the output array, so its lifetime (and @@ -279,6 +273,12 @@ fn append_geotiff_view(builder: &mut BinaryViewBuilder, bytes: VSIBuffer) -> Res .map_err(|e| exec_datafusion_err!("RS_AsGeoTiff: failed to append binary view: {e}")) } +/// Map a quality fraction in `[0.0, 1.0]` to GDAL's 1–100 `JPEG_QUALITY`. +/// +/// The fractional scale matches Apache Sedona (GeoTools' `setCompressionQuality`); +/// a value outside the range errors rather than clamping — silently clamping +/// would turn the most likely mistake (passing a 0–100 quality like `75`) into +/// maximum quality with no warning. fn jpeg_quality_option(quality: f64) -> Result { if !(0.0..=1.0).contains(&quality) { return exec_err!(