feat: N-dimensional raster type extension (Phase 1)

rust/sedona-raster-functions/src/rs_pixel_functions.rs

@@ -191,7 +191,7 @@ impl SedonaScalarKernel for RsPixelAsCentroid {
                     let grid_x = (col_x - 1) as f64 + 0.5;
                     let grid_y = (row_y - 1) as f64 + 0.5;
 
-                    let affine = AffineMatrix::from_metadata(raster.metadata());
+                    let affine = AffineMatrix::from_metadata(&raster.metadata());
                     let (wx, wy) = affine.transform(grid_x, grid_y);
 
                     write_wkb_point(&mut builder, (wx, wy))

rust/sedona-raster-gdal/src/gdal_common.rs

@@ -222,6 +222,13 @@ pub unsafe fn raster_ref_to_gdal_mem<R: RasterRef + ?Sized>(
             .band(src_band_index)
             .map_err(|e| arrow_datafusion_err!(e))?;
 
+        if !band.is_2d() {
+            return exec_err!(
+                "GDAL backend requires a 2-dim band; got dim_names={:?}",
+                band.dim_names()
+            );
+        }
+
         if band.metadata().storage_type()? != StorageType::InDb {
             return Err(DataFusionError::NotImplemented(
                 "OutDb bands are not supported in raster_to_mem_dataset".to_string(),
@@ -825,4 +832,40 @@ mod tests {
             .unwrap();
         assert!(err.to_string().contains("OutDb bands are not supported"));
     }
+
+    #[test]
+    fn test_raster_ref_to_gdal_mem_rejects_nd_bands() {
+        // Build a 3-D in-db band shaped ["time","y","x"] over a 2-D raster.
+        // The N-D guard should fire before any GDAL call.
+        let mut builder = RasterBuilder::new(1);
+        builder
+            .start_raster_2d(2, 2, 0.0, 2.0, 1.0, -1.0, 0.0, 0.0, None)
+            .unwrap();
+        builder
+            .start_band_nd(
+                None,
+                &["time", "y", "x"],
+                &[3, 2, 2],
+                BandDataType::UInt8,
+                None,
+                None,
+                None,
+            )
+            .unwrap();
+        builder
+            .band_data_writer()
+            .append_value(vec![0u8; 3 * 2 * 2]);
+        builder.finish_band().unwrap();
+        builder.finish_raster().unwrap();
+        let raster_array = builder.finish().unwrap();
+        let raster = single_raster(&raster_array);
+
+        let err = with_gdal(|gdal| unsafe { raster_ref_to_gdal_mem(gdal, &raster, &[1]) })
+            .err()
+            .unwrap();
+        assert!(
+            err.to_string().contains("requires a 2-dim band"),
+            "got: {err}"
+        );
+    }
 }

rust/sedona-raster-gdal/src/gdal_dataset_provider.rs

@@ -259,6 +259,14 @@ impl GDALDatasetCache {
 
         for i in 1..=num_bands {
             let band = bands.band(i).map_err(|e| arrow_datafusion_err!(e))?;
+
+            if !band.is_2d() {
+                return exec_err!(
+                    "GDAL backend requires 2-dim bands; got dim_names={:?}",
+                    band.dim_names()
+                );
+            }
+
             let band_metadata = band.metadata();
             let band_type = band_metadata.data_type()?;
             let gdal_type = band_data_type_to_gdal(&band_type);
@@ -660,6 +668,32 @@ mod tests {
         path_str
     }
 
+    /// Two-band GeoTIFF on disk: band 1 is filled with `band1_fill`, band 2
+    /// with `band2_fill`. Used to exercise `#band=2` selection end-to-end.
+    fn create_two_band_source_tiff(temp_dir: &TempDir, band1_fill: u8, band2_fill: u8) -> String {
+        let path = temp_dir.path().join("two_band.tif");
+        let path_str = path.to_string_lossy().to_string();
+
+        with_gdal(|gdal| {
+            let driver = gdal.get_driver_by_name("GTiff").unwrap();
+            let dataset = driver
+                .create_with_band_type::<u8>(&path_str, 8, 8, 2)
+                .unwrap();
+            dataset
+                .set_geo_transform(&[0.0, 1.0, 0.0, 8.0, 0.0, -1.0])
+                .unwrap();
+            for (i, fill) in [band1_fill, band2_fill].iter().enumerate() {
+                let band = dataset.rasterband(i + 1).unwrap();
+                let mut buffer = Buffer::new((8, 8), vec![*fill; 8 * 8]);
+                band.write((0, 0), (8, 8), &mut buffer).unwrap();
+            }
+            Ok(())
+        })
+        .unwrap();
+
+        path_str
+    }
+
     fn build_outdb_raster(path: &str) -> arrow_array::StructArray {
         let mut builder = RasterBuilder::new(1);
         let metadata = RasterMetadata {
@@ -976,4 +1010,109 @@ mod tests {
 
         assert!(key_a != key_b);
     }
+
+    #[test]
+    fn test_provider_rejects_nd_band_in_vrt_path() {
+        let temp_dir = TempDir::new().unwrap();
+        let path = create_source_tiff(&temp_dir);
+
+        // Build a raster mixing one in-db 3-D band (forces N-D rejection inside
+        // build_vrt_from_sources) with one out-db band.
+        let mut builder = RasterBuilder::new(1);
+        builder
+            .start_raster_2d(8, 8, 0.0, 8.0, 1.0, -1.0, 0.0, 0.0, None)
+            .unwrap();
+        builder
+            .start_band_nd(
+                None,
+                &["time", "y", "x"],
+                &[2, 8, 8],
+                BandDataType::UInt8,
+                None,
+                None,
+                None,
+            )
+            .unwrap();
+        builder
+            .band_data_writer()
+            .append_value(vec![0u8; 2 * 8 * 8]);
+        builder.finish_band().unwrap();
+        builder
+            .start_band_nd(
+                None,
+                &["y", "x"],
+                &[8, 8],
+                BandDataType::UInt8,
+                Some(&[0u8]),
+                Some(&path),
+                Some("geotiff"),
+            )
+            .unwrap();
+        builder.band_data_writer().append_value([]);
+        builder.finish_band().unwrap();
+        builder.finish_raster().unwrap();
+        let raster_struct = builder.finish().unwrap();
+        let raster_array = RasterStructArray::new(&raster_struct);
+        let raster = raster_array.get(0).unwrap();
+        let cache = Rc::new(GDALDatasetCache::try_new(4, 4).unwrap());
+
+        let err = with_gdal(|gdal| {
+            let provider = GDALDatasetProvider::new(gdal, Rc::clone(&cache));
+            provider.raster_ref_to_gdal(&raster)
+        })
+        .err()
+        .unwrap();
+        assert!(err.to_string().contains("2-dim band"), "got: {err}");
+    }
+
+    #[test]
+    fn test_provider_selects_outdb_band_via_band_fragment() {
+        let temp_dir = TempDir::new().unwrap();
+        // Source TIFF: band 1 filled with 7s, band 2 filled with 99s.
+        let path = create_two_band_source_tiff(&temp_dir, 7u8, 99u8);
+
+        // Build a 1-band raster whose single band points at source band 2.
+        let metadata = RasterMetadata {
+            width: 8,
+            height: 8,
+            upperleft_x: 0.0,
+            upperleft_y: 8.0,
+            scale_x: 1.0,
+            scale_y: -1.0,
+            skew_x: 0.0,
+            skew_y: 0.0,
+        };
+        let mut builder = RasterBuilder::new(1);
+        builder.start_raster(&metadata, None).unwrap();
+        builder
+            .start_band(BandMetadata {
+                nodata_value: Some(vec![0u8]),
+                storage_type: StorageType::OutDbRef,
+                datatype: BandDataType::UInt8,
+                outdb_url: Some(path.clone()),
+                outdb_band_id: Some(2),
+            })
+            .unwrap();
+        builder.band_data_writer().append_value([]);
+        builder.finish_band().unwrap();
+        builder.finish_raster().unwrap();
+        let raster_struct = builder.finish().unwrap();
+        let raster_array = RasterStructArray::new(&raster_struct);
+        let raster = raster_array.get(0).unwrap();
+        let cache = Rc::new(GDALDatasetCache::try_new(4, 4).unwrap());
+
+        let dataset = with_gdal(|gdal| {
+            let provider = GDALDatasetProvider::new(gdal, Rc::clone(&cache));
+            provider.raster_ref_to_gdal(&raster)
+        })
+        .unwrap();
+
+        let band = dataset
+            .as_dataset()
+            .rasterband(1)
+            .unwrap()
+            .read_as::<u8>((0, 0), (8, 8), (8, 8), None)
+            .unwrap();
+        assert_eq!(band.data().to_vec(), vec![99u8; 8 * 8]);
+    }
 }

rust/sedona-raster/Cargo.toml

@@ -34,9 +34,11 @@ result_large_err = "allow"
 arrow-schema = { workspace = true }
 arrow-array = { workspace = true }
 arrow-buffer = { workspace = true }
+datafusion-common = { workspace = true }
 sedona-common = { workspace = true }
 sedona-schema = { workspace = true }
 
 [dev-dependencies]
 sedona-testing = { workspace = true }
 approx = { workspace = true }
+arrow-ipc = { workspace = true }

rust/sedona-raster/src/affine_transformation.rs

@@ -108,7 +108,7 @@ pub fn rotation(raster: &dyn RasterRef) -> f64 {
 /// * `y` - Y coordinate in pixel space (row)
 #[inline]
 pub fn to_world_coordinate(raster: &dyn RasterRef, x: i64, y: i64) -> (f64, f64) {
-    AffineMatrix::from_metadata(raster.metadata()).transform(x as f64, y as f64)
+    AffineMatrix::from_metadata(&raster.metadata()).transform(x as f64, y as f64)
 }
 
 /// Performs the inverse affine transformation to convert world coordinates back to raster pixel coordinates.
@@ -124,14 +124,14 @@ pub fn to_raster_coordinate(
     world_y: f64,
 ) -> Result<(i64, i64), ArrowError> {
     let (rx, ry) =
-        AffineMatrix::from_metadata(raster.metadata()).inv_transform(world_x, world_y)?;
+        AffineMatrix::from_metadata(&raster.metadata()).inv_transform(world_x, world_y)?;
     Ok((rx as i64, ry as i64))
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::traits::{MetadataRef, RasterMetadata};
+    use crate::traits::{BandRef, Bands, RasterMetadata};
     use approx::assert_relative_eq;
     use std::f64::consts::FRAC_1_SQRT_2;
     use std::f64::consts::PI;
@@ -141,14 +141,34 @@ mod tests {
     }
 
     impl RasterRef for TestRaster {
-        fn metadata(&self) -> &dyn MetadataRef {
-            &self.metadata
+        fn num_bands(&self) -> usize {
+            0
+        }
+        fn bands(&self) -> Bands<'_> {
+            Bands::new(self)
+        }
+        fn band(&self, index: usize) -> Result<Box<dyn BandRef + '_>, ArrowError> {
+            Err(ArrowError::InvalidArgumentError(format!(
+                "Band index {index} is out of range: this raster has 0 bands"
+            )))
+        }
+        fn band_name(&self, _index: usize) -> Option<&str> {
+            None
         }
         fn crs(&self) -> Option<&str> {
             None
         }
-        fn bands(&self) -> &dyn crate::traits::BandsRef {
-            unimplemented!()
+        fn transform(&self) -> &[f64] {
+            &[]
+        }
+        fn spatial_dims(&self) -> Vec<&str> {
+            vec![]
+        }
+        fn spatial_shape(&self) -> &[i64] {
+            &[]
+        }
+        fn metadata(&self) -> RasterMetadata {
+            self.metadata.clone()
         }
     }