diff --git a/datafusion/ffi/src/tests/udf_udaf_udwf.rs b/datafusion/ffi/src/tests/udf_udaf_udwf.rs index 830c639c743d6..c7d9c2841c7c9 100644 --- a/datafusion/ffi/src/tests/udf_udaf_udwf.rs +++ b/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 { Ok(inputs.iter().all(|input| input.preserves_lex_ordering)) } + + fn evaluate_bounds( + &self, + inputs: &[&Interval], + ) -> datafusion_common::Result { + 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 { diff --git a/datafusion/ffi/src/udf/mod.rs b/datafusion/ffi/src/udf/mod.rs index d14614f1474a3..38f0798b20c40 100644 --- a/datafusion/ffi/src/udf/mod.rs +++ b/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_Result, + /// FFI equivalent to [`ScalarUDFImpl::evaluate_bounds`]. + pub evaluate_bounds: unsafe extern "C" fn( + udf: &Self, + inputs: SVec, + ) -> FFI_Result, + /// 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_Result { + let inputs = sresult_return!( + inputs + .into_iter() + .map(Interval::try_from) + .collect::>>() + ); + let inputs = inputs.iter().collect::>(); + + 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, @@ -320,6 +347,7 @@ impl From> for FFI_ScalarUDF { private_data: Box::into_raw(private_data).cast::(), 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 { + let inputs = inputs + .iter() + .map(|interval| FFI_Interval::try_from(*interval)) + .collect::>>()?; + + 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 { + 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"); diff --git a/datafusion/ffi/tests/ffi_udf.rs b/datafusion/ffi/tests/ffi_udf.rs index 73fafef13415b..99b258967a9d7 100644 --- a/datafusion/ffi/tests/ffi_udf.rs +++ b/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(()) }