fix: preserve scalar UDF bounds across FFI

datafusion/ffi/src/tests/udf_udaf_udwf.rs

@@ -21,6 +21,7 @@ use arrow_schema::DataType;
 use datafusion_catalog::TableFunctionImpl;
 use datafusion_common::ScalarValue;
 use datafusion_common::config::ConfigOptions;
+use datafusion_expr::interval_arithmetic::Interval;
 use datafusion_expr::sort_properties::ExprProperties;
 use datafusion_expr::{
     AggregateUDF, ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDF,
@@ -169,6 +170,20 @@ impl ScalarUDFImpl for PlacementUDF {
     ) -> datafusion_common::Result<bool> {
         Ok(inputs.iter().all(|input| input.preserves_lex_ordering))
     }
+
+    fn evaluate_bounds(
+        &self,
+        inputs: &[&Interval],
+    ) -> datafusion_common::Result<Interval> {
+        inputs
+            .first()
+            .map(|interval| (*interval).clone())
+            .ok_or_else(|| {
+                datafusion_common::DataFusionError::Internal(
+                    "expected one input".to_string(),
+                )
+            })
+    }
 }
 
 pub(crate) extern "C" fn create_placement_func() -> FFI_ScalarUDF {

datafusion/ffi/src/udf/mod.rs

@@ -26,6 +26,7 @@ use arrow::ffi::{FFI_ArrowSchema, from_ffi, to_ffi};
 use arrow_schema::FieldRef;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::{DataFusionError, Result, internal_err};
+use datafusion_expr::interval_arithmetic::Interval;
 use datafusion_expr::sort_properties::ExprProperties;
 use datafusion_expr::type_coercion::functions::fields_with_udf;
 use datafusion_expr::{
@@ -43,6 +44,7 @@ use crate::arrow_wrappers::{WrappedArray, WrappedSchema};
 use crate::config::FFI_ConfigOptions;
 use crate::expr::columnar_value::FFI_ColumnarValue;
 use crate::expr::expr_properties::FFI_ExprProperties;
+use crate::expr::interval::FFI_Interval;
 use crate::placement::FFI_ExpressionPlacement;
 use crate::util::{
     FFI_Option, FFI_Result, rvec_wrapped_to_vec_datatype, vec_datatype_to_rvec_wrapped,
@@ -124,6 +126,12 @@ pub struct FFI_ScalarUDF {
         inputs: SVec<FFI_ExprProperties>,
     ) -> FFI_Result<bool>,
 
+    /// FFI equivalent to [`ScalarUDFImpl::evaluate_bounds`].
+    pub evaluate_bounds: unsafe extern "C" fn(
+        udf: &Self,
+        inputs: SVec<FFI_Interval>,
+    ) -> FFI_Result<FFI_Interval>,
+
     /// FFI equivalent to [`ScalarUDFImpl::with_updated_config`].
     pub with_updated_config:
         unsafe extern "C" fn(
@@ -221,6 +229,25 @@ unsafe extern "C" fn with_updated_config_fn_wrapper(
     FFI_Result::Ok(updated.into())
 }
 
+unsafe extern "C" fn evaluate_bounds_fn_wrapper(
+    udf: &FFI_ScalarUDF,
+    inputs: SVec<FFI_Interval>,
+) -> FFI_Result<FFI_Interval> {
+    let inputs = sresult_return!(
+        inputs
+            .into_iter()
+            .map(Interval::try_from)
+            .collect::<Result<Vec<_>>>()
+    );
+    let inputs = inputs.iter().collect::<Vec<_>>();
+
+    sresult!(
+        udf.inner()
+            .evaluate_bounds(&inputs)
+            .and_then(FFI_Interval::try_from)
+    )
+}
+
 unsafe extern "C" fn invoke_with_args_fn_wrapper(
     udf: &FFI_ScalarUDF,
     args: SVec<WrappedArray>,
@@ -320,6 +347,7 @@ impl From<Arc<ScalarUDF>> for FFI_ScalarUDF {
             private_data: Box::into_raw(private_data).cast::<c_void>(),
             library_marker_id: crate::get_library_marker_id,
             preserves_lex_ordering: preserves_lex_ordering_fn_wrapper,
+            evaluate_bounds: evaluate_bounds_fn_wrapper,
             with_updated_config: with_updated_config_fn_wrapper,
         }
     }
@@ -548,6 +576,18 @@ impl ScalarUDFImpl for ForeignScalarUDF {
 
         Some(ScalarUDF::new_from_shared_impl(updated.into()))
     }
+
+    fn evaluate_bounds(&self, inputs: &[&Interval]) -> Result<Interval> {
+        let inputs = inputs
+            .iter()
+            .map(|interval| FFI_Interval::try_from(*interval))
+            .collect::<Result<SVec<_>>>()?;
+
+        unsafe {
+            df_result!((self.udf.evaluate_bounds)(&self.udf, inputs))
+                .and_then(Interval::try_from)
+        }
+    }
 }
 
 #[cfg(test)]
@@ -602,6 +642,15 @@ mod tests {
                 signature: self.signature.clone(),
             }))
         }
+
+        fn evaluate_bounds(&self, inputs: &[&Interval]) -> Result<Interval> {
+            inputs
+                .first()
+                .map(|interval| (*interval).clone())
+                .ok_or_else(|| {
+                    DataFusionError::Internal("expected one input".to_string())
+                })
+        }
     }
 
     #[test]
@@ -694,6 +743,13 @@ mod tests {
         );
         assert!(foreign_udf.preserves_lex_ordering(&[]).is_err());
 
+        let interval = Interval::try_new(
+            datafusion_common::ScalarValue::Int64(Some(2)),
+            datafusion_common::ScalarValue::Int64(Some(8)),
+        )?;
+        assert_eq!(foreign_udf.evaluate_bounds(&[&interval])?, interval);
+        assert!(foreign_udf.evaluate_bounds(&[]).is_err());
+
         let updated = foreign_udf
             .with_updated_config(&ConfigOptions::default())
             .expect("provider should return an updated UDF");

datafusion/ffi/tests/ffi_udf.rs

@@ -25,7 +25,9 @@ mod tests {
     use datafusion::error::Result;
     use datafusion::logical_expr::{ExpressionPlacement, ScalarUDF, ScalarUDFImpl};
     use datafusion::prelude::{SessionContext, col};
+    use datafusion::scalar::ScalarValue;
     use datafusion_execution::config::SessionConfig;
+    use datafusion_expr::interval_arithmetic::Interval;
     use datafusion_expr::lit;
     use datafusion_expr::sort_properties::ExprProperties;
     use datafusion_ffi::tests::create_record_batch;
@@ -119,6 +121,11 @@ mod tests {
         assert!(foreign_func.preserves_lex_ordering(std::slice::from_ref(&preserves))?);
         assert!(!foreign_func.preserves_lex_ordering(&[preserves, does_not_preserve])?);
 
+        let interval =
+            Interval::try_new(ScalarValue::Int64(Some(2)), ScalarValue::Int64(Some(8)))?;
+        assert_eq!(foreign_func.evaluate_bounds(&[&interval])?, interval);
+        assert!(foreign_func.evaluate_bounds(&[]).is_err());
+
         Ok(())
     }