jesspav commented on code in PR #406: URL: https://github.com/apache/sedona-db/pull/406#discussion_r2594100795
########## rust/sedona-gdal/src/raster_io.rs: ########## @@ -0,0 +1,679 @@ +// 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. + +use crate::dataset::{geotransform_components, tile_size, to_banddatatype}; +use arrow::array::StructArray; +use arrow::error::ArrowError; +use gdal::raster::RasterBand; +use gdal::Dataset; +use sedona_raster::builder::RasterBuilder; +use sedona_raster::traits::{BandMetadata, RasterMetadata}; +use sedona_raster::utils::{bytes_per_pixel, f64_to_bandtype_bytes}; +use sedona_schema::raster::{BandDataType, StorageType}; +use std::sync::Arc; + +// ============================================================================ +// Format-specific convenience wrappers +// ============================================================================ + +/// Reads a GeoTIFF file using GDAL and converts it into a StructArray of rasters. +/// +/// This is a convenience wrapper around [`read_raster`] for GeoTIFF files. +/// +/// # Arguments +/// * `filepath` - Path to the GeoTIFF file +/// * `tile_size_opt` - Optional tile size to override dataset metadata +pub fn read_geotiff( + filepath: &str, + tile_size_opt: Option<(usize, usize)>, +) -> Result<Arc<StructArray>, ArrowError> { + // Check that the filepath has a GeoTIFF extension + let filepath_lower = filepath.to_lowercase(); + if !filepath_lower.ends_with(".tif") && !filepath_lower.ends_with(".tiff") { + return Err(ArrowError::InvalidArgumentError(format!( + "Expected GeoTIFF file with .tif or .tiff extension, got: {}", + filepath + ))); + } + read_raster(filepath, tile_size_opt) +} + +/// Writes a tiled raster StructArray to a GeoTIFF file using GDAL. +/// +/// This is a convenience wrapper around [`write_raster`] for GeoTIFF files. +/// +/// # Arguments +/// * `raster_array` - The raster struct array to write +/// * `filepath` - Path to the output GeoTIFF file +pub fn write_geotiff(raster_array: &StructArray, filepath: &str) -> Result<(), ArrowError> { + // Check that the filepath has a GeoTIFF extension + let filepath_lower = filepath.to_lowercase(); + if !filepath_lower.ends_with(".tif") && !filepath_lower.ends_with(".tiff") { + return Err(ArrowError::InvalidArgumentError(format!( + "Expected GeoTIFF file with .tif or .tiff extension, got: {}", + filepath + ))); + } + write_raster(raster_array, filepath, "GTiff") +} + +/// Reads a raster file using GDAL and converts it into a StructArray of rasters. +/// +/// Currently only supports reading rasters into InDb storage type. +/// OutDb storage types are not yet implemented. +/// +/// # Arguments +/// * `filepath` - Path to the raster file +/// * `tile_size_opt` - Optional tile size to override dataset metadata (uses dataset metadata if None) +pub fn read_raster( + filepath: &str, + tile_size_opt: Option<(usize, usize)>, +) -> Result<Arc<StructArray>, ArrowError> { + let dataset = Dataset::open(filepath).map_err(|err| ArrowError::ParseError(err.to_string()))?; + + // Extract geotransform components + let (origin_upperleft_x, origin_upperleft_y, scale_x, scale_y, skew_x, skew_y) = + geotransform_components(&dataset)?; + + let (raster_width, raster_height) = dataset.raster_size(); + + // Use provided tile dimensions or fall back to dataset metadata + let (tile_width, tile_height) = match tile_size_opt { + Some((w, h)) => (w, h), + _ => tile_size(&dataset)?, + }; + + let x_tile_count = raster_width.div_ceil(tile_width); + let y_tile_count = raster_height.div_ceil(tile_height); + + let mut raster_builder = RasterBuilder::new(x_tile_count * y_tile_count); + let band_count = dataset.raster_count(); + + for tile_y in 0..y_tile_count { + for tile_x in 0..x_tile_count { + let x_offset = tile_x * tile_width; + let y_offset = tile_y * tile_height; + + // Calculate geographic coordinates for this tile + // using the geotransform from the original raster + let tile_upperleft_x = + origin_upperleft_x + (x_offset as f64) * scale_x + (y_offset as f64) * skew_x; + let tile_upperleft_y = + origin_upperleft_y + (x_offset as f64) * skew_y + (y_offset as f64) * scale_y; + + // Create raster metadata for this tile with actual geotransform values + let tile_metadata = RasterMetadata { + width: tile_width as u64, + height: tile_height as u64, + upperleft_x: tile_upperleft_x, + upperleft_y: tile_upperleft_y, + scale_x, + scale_y, + skew_x, + skew_y, + }; + + raster_builder.start_raster(&tile_metadata, None)?; + + for band_number in 1..=band_count { + let band: RasterBand = dataset.rasterband(band_number).unwrap(); + + // The band size is always the full raster size, not the tile size + // We read a specific window from the band + let actual_tile_width = tile_width.min(raster_width - x_offset); + let actual_tile_height = tile_height.min(raster_height - y_offset); + + let data_type = to_banddatatype(band.band_type())?; + let nodata_value = band + .no_data_value() + .map(|val| f64_to_bandtype_bytes(val, data_type.clone())); Review Comment: done! -- 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]
