paleolimbot commented on code in PR #1028: URL: https://github.com/apache/sedona-db/pull/1028#discussion_r3576094365
########## docs/reference/sql/rs_asgeotiff.qmd: ########## @@ -0,0 +1,99 @@ +--- +# 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 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 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 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 in pixels; must be a multiple of 16. +--- + +## Description + +`RS_AsGeoTiff` encodes a raster as a GeoTIFF (via GDAL) and returns the file Review Comment: This is another one with a lot of options that we should mark as experimental so we can improve the options handling and break compatibility ########## 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. [email protected]( Review Comment: You have rasterio here...you can also check that its contents is identical when decoded and parameterize on data type to make sure it works for all of them with the default parameters (e.g., with default compression heuristics) ########## 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) Review Comment: This copy (and then the second copy when this is appended to the Binary output array) seems like something we could improve. Probably wrapping the vsimem file as a Bytes would put some pressure on the vsimem system that I'm not sure it's designed for, but we could at least get rid of the second copy by returning a binary view. ########## 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: Are all of these error pathways tested? ########## 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: Just putting it out there that with options as a single argument, there are only ever two kernels and `serde_json` takes care of deserializing the options for you -- 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]
