james-willis commented on code in PR #1028: URL: https://github.com/apache/sedona-db/pull/1028#discussion_r3581634583
########## rust/sedona-raster-gdal/src/rs_as_geotiff.rs: ########## @@ -0,0 +1,870 @@ +// 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, 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, thread_local_provider, GDALDatasetProvider, +}; + +/// 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<Self> { + 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. `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::Relaxed); + 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, + provider: &GDALDatasetProvider, + raster: &RasterRefImpl, + compression: Option<CompressionType>, + quality: Option<f64>, + tile_width: Option<u32>, + tile_height: Option<u32>, + ) -> Result<Vec<u8>> { + 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))?; + + // 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<String> = 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) = jpeg_quality { + options_list.push(format!("JPEG_QUALITY={}", q)); + } + } + + // 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(format!("PREDICTOR={}", predictor_for(raster)?)); + } + } + + // 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(); + + // 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 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))?; + + // 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) + } +} + +/// 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, +} + +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); + } +} + +/// 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<i32> { + 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<i32> { + 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 { + fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> { + 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<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(); + + // 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) + } Review Comment: hmm this brings up the same question from RS_Value where we have to consider if we want to match sedona spark or not. https://sedona.apache.org/latest/api/sql/Raster-Output/RS_AsGeoTiff/ -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
