diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index d372dcf9177e6..59bcc861d0567 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -31,8 +31,9 @@ use datafusion_expr::{Expr, TableType}; use datafusion_physical_plan::ExecutionPlan; use sync_provider::create_sync_table_provider; use udf_udaf_udwf::{ - create_ffi_abs_func, create_ffi_random_func, create_ffi_rank_func, - create_ffi_stddev_func, create_ffi_sum_func, create_ffi_table_func, + create_ffi_abs_func, create_ffi_first_value_func, create_ffi_random_func, + create_ffi_rank_func, create_ffi_stddev_func, create_ffi_sum_func, + create_ffi_table_func, }; use crate::catalog_provider::FFI_CatalogProvider; @@ -117,6 +118,9 @@ pub struct ForeignLibraryModule { pub create_context_aware_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, pub version: extern "C" fn() -> u64, + + /// Create an aggregate UDAF using first_value + pub create_first_value_udaf: extern "C" fn() -> FFI_AggregateUDF, } pub fn create_test_schema() -> Arc { @@ -266,5 +270,6 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_context_aware_optimizer_rule: physical_optimizer::create_context_aware_optimizer_rule, version: super::version, + create_first_value_udaf: create_ffi_first_value_func, } } diff --git a/datafusion/ffi/src/tests/udf_udaf_udwf.rs b/datafusion/ffi/src/tests/udf_udaf_udwf.rs index b393f5db3a506..b5ed406420376 100644 --- a/datafusion/ffi/src/tests/udf_udaf_udwf.rs +++ b/datafusion/ffi/src/tests/udf_udaf_udwf.rs @@ -26,6 +26,7 @@ use datafusion_expr::{ }; use datafusion_functions::math::abs::AbsFunc; use datafusion_functions::math::random::RandomFunc; +use datafusion_functions_aggregate::first_last::FirstValue; use datafusion_functions_aggregate::stddev::Stddev; use datafusion_functions_aggregate::sum::Sum; use datafusion_functions_table::generate_series::RangeFunc; @@ -176,6 +177,12 @@ pub(crate) extern "C" fn create_ffi_sum_func() -> FFI_AggregateUDF { udaf.into() } +pub(crate) extern "C" fn create_ffi_first_value_func() -> FFI_AggregateUDF { + let udaf: Arc = Arc::new(FirstValue::new().into()); + + udaf.into() +} + pub(crate) extern "C" fn create_ffi_stddev_func() -> FFI_AggregateUDF { let udaf: Arc = Arc::new(Stddev::new().into()); diff --git a/datafusion/ffi/src/udaf/mod.rs b/datafusion/ffi/src/udaf/mod.rs index c4f8fb1254e84..b3a087e5d0022 100644 --- a/datafusion/ffi/src/udaf/mod.rs +++ b/datafusion/ffi/src/udaf/mod.rs @@ -145,6 +145,10 @@ pub struct FFI_AggregateUDF { /// the foreign interface. See [`crate::get_library_marker_id`] and /// the crate's `README.md` for more information. pub library_marker_id: extern "C" fn() -> usize, + + /// FFI equivalent to [`AggregateUDF::supports_null_handling_clause`] + pub supports_null_handling_clause: + unsafe extern "C" fn(udaf: &FFI_AggregateUDF) -> bool, } unsafe impl Send for FFI_AggregateUDF {} @@ -327,6 +331,12 @@ unsafe extern "C" fn order_sensitivity_fn_wrapper( unsafe { udaf.inner().order_sensitivity().into() } } +unsafe extern "C" fn supports_null_handling_clause_fn_wrapper( + udaf: &FFI_AggregateUDF, +) -> bool { + unsafe { udaf.inner().supports_null_handling_clause() } +} + unsafe extern "C" fn coerce_types_fn_wrapper( udaf: &FFI_AggregateUDF, arg_types: SVec, @@ -401,6 +411,7 @@ impl From> for FFI_AggregateUDF { release: release_fn_wrapper, private_data: Box::into_raw(private_data) as *mut c_void, library_marker_id: crate::get_library_marker_id, + supports_null_handling_clause: supports_null_handling_clause_fn_wrapper, } } } @@ -595,6 +606,10 @@ impl AggregateUDFImpl for ForeignAggregateUDF { unsafe { (self.udaf.order_sensitivity)(&self.udaf).into() } } + fn supports_null_handling_clause(&self) -> bool { + unsafe { (self.udaf.supports_null_handling_clause)(&self.udaf) } + } + fn simplify(&self) -> Option { None } @@ -774,6 +789,19 @@ mod tests { Ok(()) } + #[test] + fn test_supports_null_handling_clause() -> Result<()> { + let first_value = create_test_foreign_udaf( + datafusion::functions_aggregate::first_last::FirstValue::new(), + )?; + assert!(first_value.supports_null_handling_clause()); + + let sum = create_test_foreign_udaf(Sum::new())?; + assert!(!sum.supports_null_handling_clause()); + + Ok(()) + } + #[test] fn test_beneficial_ordering() -> Result<()> { let foreign_udaf = create_test_foreign_udaf( diff --git a/datafusion/ffi/tests/ffi_udaf.rs b/datafusion/ffi/tests/ffi_udaf.rs index 3234f6533df9c..090151416e4e9 100644 --- a/datafusion/ffi/tests/ffi_udaf.rs +++ b/datafusion/ffi/tests/ffi_udaf.rs @@ -66,6 +66,22 @@ mod tests { Ok(()) } + #[test] + fn test_supports_null_handling_clause() -> Result<()> { + let module = get_module()?; + + let ffi_first_value_func = (module.create_first_value_udaf)(); + let foreign_first_value_func: Arc = + (&ffi_first_value_func).into(); + assert!(foreign_first_value_func.supports_null_handling_clause()); + + let ffi_sum_func = (module.create_sum_udaf)(); + let foreign_sum_func: Arc = (&ffi_sum_func).into(); + assert!(!foreign_sum_func.supports_null_handling_clause()); + + Ok(()) + } + #[tokio::test] async fn test_ffi_grouping_udaf() -> Result<()> { let module = get_module()?;