Add a Dataset::read_as method that reads a whole image at once

This is similar to RasterBand::read_as, but it is reading the whole
image at once. Similar to how GDAL has RasterIO on both the dataset and
band level.

Author: julienr

CHANGES.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+- Added `Dataset::read_as`
+
+  - <https://github.com/georust/gdal/pull/374>
+
 - Added `CslStringList::add_string`
 
   - <https://github.com/georust/gdal/pull/364>

src/dataset.rs

@@ -1,6 +1,6 @@
 use ptr::null_mut;
 use std::convert::TryInto;
-use std::mem::MaybeUninit;
+use std::mem::{size_of, MaybeUninit};
 use std::{
     ffi::NulError,
     ffi::{CStr, CString},
@@ -12,6 +12,7 @@ use std::{
 use crate::cpl::CslStringList;
 use crate::errors::*;
 use crate::raster::RasterCreationOption;
+use crate::raster::{Buffer3D, GdalType, RasterIOExtraArg, ResampleAlg};
 use crate::utils::{_last_cpl_err, _last_null_pointer_err, _path_to_c_string, _string};
 use crate::vector::{sql, Geometry, OwnedLayer};
 use crate::{
@@ -20,10 +21,11 @@ use crate::{
 };
 
 use gdal_sys::{
-    self, CPLErr, GDALAccess, GDALDatasetH, GDALMajorObjectH, OGRErr, OGRLayerH, OGRwkbGeometryType,
+    self, CPLErr, GDALAccess, GDALDatasetH, GDALMajorObjectH, GDALRWFlag, GDALRasterIOExtraArg,
+    OGRErr, OGRLayerH, OGRwkbGeometryType,
 };
 use gdal_sys::{GDALFlushCache, OGRGeometryH};
-use libc::{c_double, c_int, c_uint};
+use libc::{c_double, c_int, c_uint, c_void};
 
 #[cfg(all(major_ge_3, minor_ge_1))]
 use crate::raster::Group;
@@ -702,6 +704,94 @@ impl Dataset {
         Ok(self.child_layer(c_layer))
     }
 
+    /// Read a [`Buffer<T>`] from this dataset, where `T` implements [`GdalType`].
+    ///
+    /// # Arguments
+    /// * `window` - the window position from top left
+    /// * `window_size` - the window size (GDAL will interpolate data if `window_size` != `buffer_size`)
+    /// * `buffer_size` - the desired size of the 'Buffer'
+    /// * `e_resample_alg` - the resample algorithm used for the interpolation. Default: `NearestNeighbor`.
+    ///
+    /// The returned buffer contains the image data in HWC order.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # fn main() -> gdal::errors::Result<()> {
+    /// use gdal::Dataset;
+    /// use gdal::raster::ResampleAlg;
+    /// let dataset = Dataset::open("fixtures/m_3607824_se_17_1_20160620_sub.tif")?;
+    /// let size = 2;
+    /// let buf = dataset.read_as::<u8>((0, 0), dataset.raster_size(), (size, size), Some(ResampleAlg::Bilinear))?;
+    /// assert_eq!(buf.size, (size, size, dataset.raster_count() as usize));
+    /// assert_eq!(buf.data, [103, 116, 101, 169, 92, 108, 94, 163, 92, 112, 93, 179, 89, 109, 91, 181]);
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn read_as<T: Copy + GdalType>(
+        &self,
+        window: (isize, isize),
+        window_size: (usize, usize),
+        buffer_size: (usize, usize),
+        e_resample_alg: Option<ResampleAlg>,
+    ) -> Result<Buffer3D<T>> {
+        let band_count = self.raster_count() as usize;
+        let pixels = buffer_size.0 * buffer_size.1 * band_count;
+        let mut data: Vec<T> = Vec::with_capacity(pixels);
+
+        let resample_alg = e_resample_alg.unwrap_or(ResampleAlg::NearestNeighbour);
+
+        let mut options: GDALRasterIOExtraArg = RasterIOExtraArg {
+            e_resample_alg: resample_alg,
+            ..Default::default()
+        }
+        .into();
+
+        let options_ptr: *mut GDALRasterIOExtraArg = &mut options;
+
+        let mut bands: Vec<i32> = (1_i32..band_count as i32 + 1_i32).collect();
+
+        let size_t = size_of::<T>() as i64;
+        // Safety: the GDALDatasetRasterIOEx writes
+        // exactly pixel elements into the slice, before we
+        // read from this slice. This paradigm is suggested
+        // in the rust std docs
+        // (https://doc.rust-lang.org/std/vec/struct.Vec.html#examples-18)
+        let rv = unsafe {
+            gdal_sys::GDALDatasetRasterIOEx(
+                self.c_dataset,
+                GDALRWFlag::GF_Read,
+                window.0 as c_int,
+                window.1 as c_int,
+                window_size.0 as c_int,
+                window_size.1 as c_int,
+                data.as_mut_ptr() as *mut c_void,
+                buffer_size.0 as c_int,
+                buffer_size.1 as c_int,
+                T::gdal_ordinal(),
+                band_count as i32,
+                bands.as_mut_ptr() as *mut c_int,
+                // We want to read a HWC
+                size_t * band_count as i64,
+                buffer_size.1 as i64 * size_t * band_count as i64,
+                size_t,
+                options_ptr,
+            )
+        };
+        if rv != CPLErr::CE_None {
+            return Err(_last_cpl_err(rv));
+        }
+
+        unsafe {
+            data.set_len(pixels);
+        };
+
+        Ok(Buffer3D {
+            size: (buffer_size.0, buffer_size.1, band_count),
+            data,
+        })
+    }
+
     /// Set the [`Dataset`]'s affine transformation; also called a _geo-transformation_.
     ///
     /// This is like a linear transformation preserves points, straight lines and planes.
@@ -1288,4 +1378,42 @@ mod tests {
         let mut ds = Dataset::open(fixture("roads.geojson")).unwrap();
         assert!(ds.start_transaction().is_err());
     }
+
+    #[test]
+    fn test_dataset_read_as() {
+        let ds = Dataset::open(fixture("m_3607824_se_17_1_20160620_sub.tif")).unwrap();
+        let size: usize = 4;
+        let band_count = ds.raster_count() as usize;
+        // Compare reading the whole dataset at once to reading each band individually
+        let ds_buf = ds
+            .read_as::<u8>(
+                (0, 0),
+                ds.raster_size(),
+                (size, size),
+                Some(ResampleAlg::Bilinear),
+            )
+            .unwrap();
+        assert_eq!(ds_buf.size, (size, size, band_count));
+
+        for band_index in 0..band_count {
+            let band = ds.rasterband(band_index as isize + 1).unwrap();
+            let band_buf = band
+                .read_as::<u8>(
+                    (0, 0),
+                    ds.raster_size(),
+                    (size, size),
+                    Some(ResampleAlg::Bilinear),
+                )
+                .unwrap();
+            assert_eq!(band_buf.size, (size, size));
+            for i in 0..size {
+                for j in 0..size {
+                    assert_eq!(
+                        band_buf.data[i * size + j],
+                        ds_buf.data[i * size * band_count + j * band_count + band_index],
+                    );
+                }
+            }
+        }
+    }
 }

src/raster/mod.rs

@@ -90,7 +90,10 @@ pub use rasterband::{
     PaletteInterpretation, RasterBand, RgbaEntry, StatisticsAll, StatisticsMinMax,
 };
 pub use rasterize::{rasterize, BurnSource, MergeAlgorithm, OptimizeMode, RasterizeOptions};
-pub use types::{AdjustedValue, Buffer, ByteBuffer, GdalDataType, GdalType, ResampleAlg};
+pub use types::{
+    AdjustedValue, Buffer, Buffer3D, ByteBuffer, GdalDataType, GdalType, RasterIOExtraArg,
+    ResampleAlg,
+};
 pub use warp::reproject;
 
 /// Key/value pair for passing driver-specific creation options to

src/raster/types.rs

@@ -460,6 +460,18 @@ impl<T: GdalType> Buffer<T> {
 
 pub type ByteBuffer = Buffer<u8>;
 
+#[derive(Debug)]
+pub struct Buffer3D<T: GdalType> {
+    pub size: (usize, usize, usize),
+    pub data: Vec<T>,
+}
+
+impl<T: GdalType> Buffer3D<T> {
+    pub fn new(size: (usize, usize, usize), data: Vec<T>) -> Buffer3D<T> {
+        Buffer3D { size, data }
+    }
+}
+
 /// Extra options used to read a raster.
 ///
 /// For documentation, see `gdal_sys::GDALRasterIOExtraArg`.