Lazy Raster Loading Support for Zarr and GDAL

rust/sedona-expr/src/scalar_udf.rs

@@ -71,6 +71,13 @@ pub struct SedonaScalarUDF {
     signature: Signature,
     kernels: Vec<ScalarKernelRef>,
     aliases: Vec<String>,
+    /// Class-level metadata: does this UDF read raster pixel bytes from
+    /// any of its inputs? The optimiser/analyzer rule for `RS_EnsureLoaded`
+    /// uses this to decide whether to wrap raster arguments. Default
+    /// `false`; set via [`SedonaScalarUDF::with_needs_bytes`] at
+    /// construction time for UDFs whose kernels call
+    /// `BandRef::nd_buffer()` / `BandRef::contiguous_data()`.
+    needs_bytes: bool,
 }
 
 impl PartialEq for SedonaScalarUDF {
@@ -191,19 +198,31 @@ impl SedonaScalarUDF {
             signature,
             kernels,
             aliases: vec![],
+            needs_bytes: false,
         }
     }
 
     /// Add aliases to an existing SedonaScalarUDF
     pub fn with_aliases(self, aliases: Vec<String>) -> SedonaScalarUDF {
+        Self { aliases, ..self }
+    }
+
+    /// Mark this UDF as one whose kernels read raster pixel bytes from
+    /// their inputs. The optimiser/analyzer rule for `RS_EnsureLoaded`
+    /// reads this to decide whether to wrap raster arguments with the
+    /// async byte-materialisation UDF.
+    pub fn with_needs_bytes(self) -> SedonaScalarUDF {
         Self {
-            name: self.name,
-            signature: self.signature,
-            kernels: self.kernels,
-            aliases,
+            needs_bytes: true,
+            ..self
         }
     }
 
+    /// Returns whether this UDF reads raster pixel bytes from its inputs.
+    pub fn needs_bytes(&self) -> bool {
+        self.needs_bytes
+    }
+
     /// Create a SedonaScalarUDF from a single kernel
     ///
     /// This constructor creates a [Volatility::Immutable] function with no documentation
@@ -334,6 +353,24 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn needs_bytes_defaults_false_and_flips_via_builder() {
+        let udf = SedonaScalarUDF::new("u", vec![], Volatility::Immutable);
+        assert!(!udf.needs_bytes());
+
+        let annotated = udf.with_needs_bytes();
+        assert!(annotated.needs_bytes());
+    }
+
+    #[test]
+    fn needs_bytes_survives_with_aliases() {
+        let udf = SedonaScalarUDF::new("u", vec![], Volatility::Immutable)
+            .with_needs_bytes()
+            .with_aliases(vec!["u_alias".to_string()]);
+        assert!(udf.needs_bytes());
+        assert_eq!(udf.aliases(), &["u_alias".to_string()]);
+    }
+
     #[test]
     fn udf_empty() -> Result<()> {
         // UDF with no implementations

rust/sedona-raster-gdal/Cargo.toml

@@ -33,16 +33,19 @@ result_large_err = "allow"
 [dependencies]
 arrow-array = { workspace = true }
 arrow-buffer = { workspace = true }
+arrow-schema = { workspace = true }
+async-trait = { workspace = true }
 datafusion-common = { workspace = true }
 lru = { workspace = true }
 sedona-common = { workspace = true }
 sedona-gdal = { workspace = true }
 sedona-raster = { workspace = true }
 sedona-schema = { workspace = true }
+tokio = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true }
 sedona-gdal = { workspace = true, features = ["gdal-sys"] }
 sedona-testing = { workspace = true }
 tempfile = { workspace = true }
-tokio = { workspace = true, features = ["rt-multi-thread"] }
+tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }

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

@@ -187,7 +187,7 @@ impl GDALDatasetCache {
         })
     }
 
-    fn get_or_create_outdb_source(
+    pub(crate) fn get_or_create_outdb_source(
         &self,
         gdal: &Gdal,
         path: &str,

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

@@ -32,14 +32,14 @@ mod gdal_common;
 #[allow(dead_code)]
 mod gdal_dataset_provider;
 
-mod utils;
-
-#[cfg(test)]
+mod outdb_loader;
 mod source_uri;
+mod utils;
 
 // Re-export main dataset conversion functions
 pub use gdal_common::{
     band_data_type_to_gdal, bytes_to_f64, gdal_to_band_data_type, gdal_type_byte_size,
     nodata_bytes_to_f64, nodata_f64_to_bytes,
 };
+pub use outdb_loader::{GdalLoader, GDAL_FORMAT};
 pub use utils::{append_as_indb_raster, dataset_to_indb_raster};