Lazy Raster Loading Support for Zarr and GDAL

c/sedona-extension/src/scalar_kernel.rs

@@ -364,8 +364,8 @@ impl ExportedScalarKernel {
     /// when passed across a boundary.
     pub fn with_function_name(self, function_name: impl AsRef<str>) -> Self {
         Self {
-            inner: self.inner,
             function_name: Some(CString::from_str(function_name.as_ref()).unwrap()),
+            ..self
         }
     }
 

rust/sedona-expr/src/scalar_udf.rs

@@ -14,7 +14,7 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
-use std::{any::Any, fmt::Debug, sync::Arc};
+use std::{any::Any, collections::HashMap, fmt::Debug, sync::Arc};
 
 use arrow_schema::{DataType, FieldRef};
 use datafusion_common::config::ConfigOptions;
@@ -60,6 +60,24 @@ impl<T: SedonaScalarKernel + 'static> IntoScalarKernelRefs for Vec<Arc<T>> {
     }
 }
 
+/// Canonical name of the `RS_EnsureLoaded` async UDF.
+///
+/// Lives here (rather than next to the UDF impl) because two crates that
+/// can't depend on each other both need it: the UDF implementation in
+/// `sedona-raster-functions`, and the logical optimizer rule in
+/// `sedona-query-planner` that wraps raster args with it. Both depend on
+/// `sedona-expr`, so this is their common home.
+pub const RS_ENSURE_LOADED_NAME: &str = "rs_ensureloaded";
+
+/// Well-known [`SedonaScalarUDF`] metadata key marking a UDF whose
+/// kernels read raster pixel bytes from their inputs. Presence with
+/// value `"true"` is what the `RS_EnsureLoaded` optimizer rule keys off
+/// to decide whether to wrap raster arguments. A generic string-keyed
+/// metadata map (rather than a dedicated bool) keeps the UDF metadata
+/// surface — and the eventual cross-cdylib FFI for it — extensible
+/// without a schema change per flag.
+pub const NEEDS_PIXELS_METADATA_KEY: &str = "needs_pixels";
+
 /// Top-level scalar user-defined function
 ///
 /// This struct implements datafusion's ScalarUDF and implements kernel dispatch
@@ -71,6 +89,13 @@ pub struct SedonaScalarUDF {
     signature: Signature,
     kernels: Vec<ScalarKernelRef>,
     aliases: Vec<String>,
+    /// Class-level, string-keyed metadata describing this UDF to the
+    /// planner. Currently carries the [`NEEDS_PIXELS_METADATA_KEY`] flag
+    /// (set via [`SedonaScalarUDF::with_needs_bytes`]) that the
+    /// `RS_EnsureLoaded` optimizer rule reads; the map shape leaves room
+    /// for further planner-visible flags — and a future cross-cdylib FFI
+    /// carrying them — without a new field per flag.
+    metadata: HashMap<String, String>,
 }
 
 impl PartialEq for SedonaScalarUDF {
@@ -191,17 +216,50 @@ impl SedonaScalarUDF {
             signature,
             kernels,
             aliases: vec![],
+            metadata: HashMap::new(),
         }
     }
 
     /// Add aliases to an existing SedonaScalarUDF
     pub fn with_aliases(self, aliases: Vec<String>) -> SedonaScalarUDF {
-        Self {
-            name: self.name,
-            signature: self.signature,
-            kernels: self.kernels,
-            aliases,
-        }
+        Self { aliases, ..self }
+    }
+
+    /// Set a class-level metadata entry on this UDF, returning the
+    /// modified UDF. Metadata is planner-visible (e.g. the
+    /// `RS_EnsureLoaded` optimizer rule reads [`NEEDS_PIXELS_METADATA_KEY`])
+    /// and crosses the `sedona-extension` FFI boundary so plugin-defined
+    /// UDFs can declare it too.
+    pub fn with_metadata(
+        mut self,
+        key: impl Into<String>,
+        value: impl Into<String>,
+    ) -> SedonaScalarUDF {
+        self.metadata.insert(key.into(), value.into());
+        self
+    }
+
+    /// Class-level metadata map describing this UDF to the planner.
+    pub fn metadata(&self) -> &HashMap<String, String> {
+        &self.metadata
+    }
+
+    /// Mark this UDF as one whose kernels read raster pixel bytes from
+    /// their inputs — convenience for setting [`NEEDS_PIXELS_METADATA_KEY`]
+    /// to `"true"`. The `RS_EnsureLoaded` optimizer rule reads this to
+    /// decide whether to wrap raster arguments with the async
+    /// byte-materialisation UDF.
+    pub fn with_needs_bytes(self) -> SedonaScalarUDF {
+        self.with_metadata(NEEDS_PIXELS_METADATA_KEY, "true")
+    }
+
+    /// Returns whether this UDF reads raster pixel bytes from its inputs
+    /// (i.e. carries [`NEEDS_PIXELS_METADATA_KEY`] = `"true"`).
+    pub fn needs_bytes(&self) -> bool {
+        self.metadata
+            .get(NEEDS_PIXELS_METADATA_KEY)
+            .map(String::as_str)
+            == Some("true")
     }
 
     /// Create a SedonaScalarUDF from a single kernel
@@ -334,6 +392,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