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 new file mode 100644 index 0000000000..eb77faef3e --- /dev/null +++ b/docs/reference/sql/rs_asgeotiff.qmd @@ -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. + +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. +--- + +::: 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 +bytes as a binary value — the inverse of reading a GeoTIFF into a raster. The +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 +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(), 16) IS NOT NULL; +``` + +```sql +SELECT RS_AsGeoTiff(RS_Example(), 'DEFLATE', 0.85) IS NOT NULL; +``` + +```sql +SELECT RS_AsGeoTiff(RS_Example(), 'LZW', 0.85, 16, 16) IS NOT NULL; +``` diff --git a/python/sedonadb/tests/functions/test_raster_functions.py b/python/sedonadb/tests/functions/test_raster_functions.py index 5905a97589..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 @@ -283,6 +280,61 @@ def test_rs_setbandnodatavalue_two_arg_requires_single_band(): ) +# 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() + + +# 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): + pytest.importorskip("rasterio") + from rasterio.io import MemoryFile + + 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 + + def _rs_as_raster_sql( pixel_type, all_touched, burn_value, nodata_value, use_geometry_extent ): @@ -348,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 @@ -426,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( """ diff --git a/rust/sedona-raster-gdal/Cargo.toml b/rust/sedona-raster-gdal/Cargo.toml index e30ad1da4e..de0cf3ce89 100644 --- a/rust/sedona-raster-gdal/Cargo.toml +++ b/rust/sedona-raster-gdal/Cargo.toml @@ -60,6 +60,11 @@ sedona-testing = { workspace = true, features = ["criterion"] } 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_clip" 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..1478f70059 --- /dev/null +++ b/rust/sedona-raster-gdal/benches/rs_as_geotiff.rs @@ -0,0 +1,120 @@ +// 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 std::sync::Arc; + +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_schema::raster::BandDataType; +use sedona_testing::raster_spec::RasterSpec; +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())) + }); + } + + // 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(); +} + +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); + // 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)); + 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 8892224c1c..74065b7055 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_as_raster; mod rs_clip; mod rs_frompath; @@ -47,6 +48,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_as_raster::rs_as_raster_udf; pub use rs_clip::rs_clip_udf; pub use rs_frompath::rs_frompath_udf; diff --git a/rust/sedona-raster-gdal/src/register.rs b/rust/sedona-raster-gdal/src/register.rs index b8789825d7..a30a13a05a 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_as_raster::rs_as_raster_udf()); function_set.insert_scalar_udf(crate::rs_clip::rs_clip_udf()); function_set.insert_scalar_udf(crate::rs_frompath::rs_frompath_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..cffcc033d5 --- /dev/null +++ b/rust/sedona-raster-gdal/src/rs_as_geotiff.rs @@ -0,0 +1,964 @@ +// 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::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; +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; +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 { + 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, + quality: Option, + tile_width: Option, + tile_height: Option, + ) -> Result { + 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 = 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))?; + + // 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_buffer_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); + } +} + +/// 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}")) +} + +/// 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 { + 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::BinaryView)); + 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 = BinaryViewBuilder::with_capacity(num_iterations); + + with_gdal(|gdal| { + configure_thread_local_options(gdal, config_options)?; + let provider = thread_local_provider(gdal) + .map_err(|e| exec_datafusion_err!("Failed to init GDAL provider: {e}"))?; + 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, + &provider, + raster, + compression, + quality, + tile_width, + tile_height, + )?; + append_geotiff_view(&mut builder, bytes)?; + + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + }) + } +} + +#[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; + 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::BinaryView(Some(bytes)) = result else { + panic!("expected a BinaryView 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 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 + // 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 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(); + 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()); + // 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 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 + // 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(); + // 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)); + 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(); + + 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(0.75), + 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(); + } +}