Kontinuation commented on code in PR #678: URL: https://github.com/apache/sedona-db/pull/678#discussion_r2878832546
########## c/sedona-gdal/src/mem.rs: ########## @@ -0,0 +1,422 @@ +// 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. + +//! High-level builder for creating in-memory (MEM) GDAL datasets. +//! +//! [`MemDatasetBuilder`] provides a fluent, type-safe API for constructing GDAL MEM +//! datasets with zero-copy band attachment, optional geo-transform, projection, and +//! per-band nodata values. +//! +//! # Example +//! +//! ```rust,ignore +//! use sedona_gdal::mem::{MemDatasetBuilder, Nodata}; +//! use sedona_gdal::GdalDataType; +//! +//! let data: Vec<u8> = vec![0u8; 256 * 256]; +//! let dataset = unsafe { +//! MemDatasetBuilder::new(256, 256) +//! .add_band(GdalDataType::UInt8, data.as_ptr()) +//! .geo_transform([0.0, 1.0, 0.0, 0.0, 0.0, -1.0]) +//! .projection("EPSG:4326") +//! .build() +//! .unwrap() +//! }; +//! assert_eq!(dataset.raster_count(), 1); +//! ``` + +use crate::dataset::Dataset; +use crate::errors::{GdalError, Result}; +use crate::gdal_api::call_gdal_api; +use crate::gdal_dyn_bindgen::CE_Failure; +use crate::raster::types::GdalDataType; + +/// Nodata value for a raster band. +/// +/// GDAL has three separate APIs for setting nodata depending on the band data type: +/// - [`f64`] for most types (UInt8 through Float64, excluding Int64/UInt64) +/// - [`i64`] for Int64 bands +/// - [`u64`] for UInt64 bands +/// +/// This enum encapsulates all three variants so callers don't need to match on +/// the band type when setting nodata. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum Nodata { + F64(f64), + I64(i64), + U64(u64), +} + +/// A band specification for [`MemDatasetBuilder`]. +struct MemBand { + data_type: GdalDataType, + data_ptr: *const u8, + pixel_offset: Option<i64>, + line_offset: Option<i64>, + nodata: Option<Nodata>, +} + +/// A builder for constructing in-memory (MEM) GDAL datasets. +/// +/// This creates datasets using `MEMDataset::Create` (bypassing GDAL's open-dataset-list +/// mutex for better concurrency) and attaches bands via `GDALAddBand` with `DATAPOINTER` +/// options for zero-copy operation. +/// +/// # Safety +/// +/// All `add_band*` methods are `unsafe` because the caller must ensure that the +/// provided data pointers remain valid for the lifetime of the built [`Dataset`]. +pub struct MemDatasetBuilder { + width: usize, + height: usize, + n_owned_bands: usize, + owned_bands_data_type: Option<GdalDataType>, + bands: Vec<MemBand>, + geo_transform: Option<[f64; 6]>, + projection: Option<String>, +} + +impl MemDatasetBuilder { + /// Create a new builder for a MEM dataset with the given dimensions. + pub fn new(width: usize, height: usize) -> Self { + Self { + width, + height, + n_owned_bands: 0, + owned_bands_data_type: None, + bands: Vec::new(), + geo_transform: None, + projection: None, + } + } + + /// Create a new builder for a MEM dataset with the given dimensions and number of owned bands. + pub fn new_with_owned_bands( + width: usize, + height: usize, + n_owned_bands: usize, + owned_bands_data_type: GdalDataType, + ) -> Self { + Self { + width, + height, + n_owned_bands, + owned_bands_data_type: Some(owned_bands_data_type), + bands: Vec::new(), + geo_transform: None, + projection: None, + } + } + + /// Create a MEM dataset with owned bands + pub fn create( + width: usize, + height: usize, + n_owned_bands: usize, + owned_bands_data_type: GdalDataType, + ) -> Result<Dataset> { + unsafe { + Self::new_with_owned_bands(width, height, n_owned_bands, owned_bands_data_type).build() + } + } + + /// Add a zero-copy band from a raw data pointer. + /// + /// Uses default pixel and line offsets (contiguous, row-major layout). + /// + /// # Safety + /// + /// The caller must ensure `data_ptr` points to a valid buffer of at least + /// `height * width * data_type.byte_size()` bytes, and that the buffer + /// outlives the built [`Dataset`]. + pub unsafe fn add_band(self, data_type: GdalDataType, data_ptr: *const u8) -> Self { + self.add_band_with_options(data_type, data_ptr, None, None, None) + } + + /// Add a zero-copy band with custom offsets and optional nodata. + /// + /// # Arguments + /// * `data_type` - The GDAL data type of the band. + /// * `data_ptr` - Pointer to the band pixel data. + /// * `pixel_offset` - Byte offset between consecutive pixels. `None` defaults to + /// the byte size of `data_type`. + /// * `line_offset` - Byte offset between consecutive lines. `None` defaults to + /// `pixel_offset * width`. + /// * `nodata` - Optional nodata value for the band. + /// + /// # Safety + /// + /// The caller must ensure `data_ptr` points to a valid buffer of sufficient size + /// for the given dimensions and offsets, and that the buffer outlives the built + /// [`Dataset`]. + pub unsafe fn add_band_with_options( + mut self, + data_type: GdalDataType, + data_ptr: *const u8, + pixel_offset: Option<i64>, + line_offset: Option<i64>, + nodata: Option<Nodata>, + ) -> Self { + self.bands.push(MemBand { + data_type, + data_ptr, + pixel_offset, + line_offset, + nodata, + }); + self + } + + /// Set the geo-transform for the dataset. + /// + /// The array is `[origin_x, pixel_width, rotation_x, origin_y, rotation_y, pixel_height]`. + pub fn geo_transform(mut self, gt: [f64; 6]) -> Self { + self.geo_transform = Some(gt); + self + } + + /// Set the projection (CRS) for the dataset as a WKT or PROJ string. + pub fn projection(mut self, wkt: impl Into<String>) -> Self { + self.projection = Some(wkt.into()); + self + } + + /// Build the GDAL MEM dataset. + /// + /// This creates an empty MEM dataset using `MEMDataset::Create` (bypassing GDAL's + /// open-dataset-list), then attaches bands, sets the geo-transform, projection, + /// and per-band nodata values. + /// + /// # Safety + /// + /// This method is unsafe because the built dataset references memory provided via + /// the `add_band*` methods. The caller must ensure all data pointers remain valid + /// for the lifetime of the returned [`Dataset`]. + pub unsafe fn build(self) -> Result<Dataset> { + let api = crate::register::get_global_gdal_api() + .map_err(|e| GdalError::BadArgument(e.to_string()))?; + + // Create an initial MEM dataset via MEMDataset::Create directly, + // bypassing GDAL's open-dataset-list mutex. + let empty_filename = c""; + let owned_bands_data_type = self + .owned_bands_data_type + .unwrap_or(GdalDataType::UInt8) + .to_c(); + let handle = unsafe { + call_gdal_api!( + api, + MEMDatasetCreate, + empty_filename.as_ptr(), + self.width as i32, + self.height as i32, + self.n_owned_bands as i32, + owned_bands_data_type, + std::ptr::null_mut() + ) + }; + + if handle.is_null() { + return Err(api.last_cpl_err(CE_Failure as u32)); + } + let dataset = Dataset::new_owned(api, handle); + + // Attach bands (zero-copy via DATAPOINTER). + for band_spec in &self.bands { + dataset.add_band_with_data( + band_spec.data_type, + band_spec.data_ptr, + band_spec.pixel_offset, + band_spec.line_offset, + )?; + } + + // Set geo-transform. + if let Some(gt) = &self.geo_transform { + dataset.set_geo_transform(gt)?; + } + + // Set projection/CRS. + if let Some(proj) = &self.projection { + dataset.set_projection(proj)?; + } + + // Set per-band nodata values. + for (i, band_spec) in self.bands.iter().enumerate() { + if let Some(nodata) = &band_spec.nodata { + let raster_band = dataset.rasterband(i + 1)?; + match nodata { + Nodata::F64(v) => raster_band.set_no_data_value(Some(*v))?, + Nodata::I64(v) => raster_band.set_no_data_value_i64(Some(*v))?, + Nodata::U64(v) => raster_band.set_no_data_value_u64(Some(*v))?, + } + } + } Review Comment: Fixed. -- 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]
