james-willis commented on code in PR #1028: URL: https://github.com/apache/sedona-db/pull/1028#discussion_r3590309611
########## 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" + ); Review Comment: Quality-out-of-range (per codec) and CreateCopy failure were; added a UDF-level unknown-compression test ('GZIP' → "Unknown compression type") GDAL internal errors (e.g.failed to get GTiff driver") were not. i dont thin theres an obvious way to trigger that -- 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]
