From 218c574fdd1030272d0dfffee52ce052326de67a Mon Sep 17 00:00:00 2001 From: Amogh Ramesh Date: Sun, 26 Jul 2026 19:38:11 +0530 Subject: [PATCH 1/2] fix(ffi): preserve aggregate null-handling support --- datafusion/ffi/src/tests/udf_udaf_udwf.rs | 60 +++++++++++++++++++++-- datafusion/ffi/src/udaf/mod.rs | 28 +++++++++++ datafusion/ffi/tests/ffi_udaf.rs | 15 ++++++ 3 files changed, 99 insertions(+), 4 deletions(-) diff --git a/datafusion/ffi/src/tests/udf_udaf_udwf.rs b/datafusion/ffi/src/tests/udf_udaf_udwf.rs index b393f5db3a506..e2a5609b635bb 100644 --- a/datafusion/ffi/src/tests/udf_udaf_udwf.rs +++ b/datafusion/ffi/src/tests/udf_udaf_udwf.rs @@ -17,17 +17,21 @@ use std::sync::Arc; -use arrow_schema::DataType; +use arrow_schema::{DataType, FieldRef}; use datafusion_catalog::TableFunctionImpl; use datafusion_common::ScalarValue; use datafusion_expr::{ - AggregateUDF, ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDF, - ScalarUDFImpl, Signature, Volatility, WindowUDF, + Accumulator, AggregateUDF, AggregateUDFImpl, ColumnarValue, ExpressionPlacement, + GroupsAccumulator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + Volatility, WindowUDF, }; use datafusion_functions::math::abs::AbsFunc; use datafusion_functions::math::random::RandomFunc; use datafusion_functions_aggregate::stddev::Stddev; use datafusion_functions_aggregate::sum::Sum; +use datafusion_functions_aggregate_common::accumulator::{ + AccumulatorArgs, StateFieldsArgs, +}; use datafusion_functions_table::generate_series::RangeFunc; use datafusion_functions_window::rank::Rank; @@ -170,8 +174,56 @@ pub(crate) extern "C" fn create_ffi_table_func( FFI_TableFunction::new_with_ffi_codec(udtf, None, codec) } +#[derive(Debug, Default, Hash, Eq, PartialEq)] +struct SumWithNullHandling { + inner: Sum, +} + +impl AggregateUDFImpl for SumWithNullHandling { + fn name(&self) -> &str { + self.inner.name() + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result { + self.inner.return_type(arg_types) + } + + fn accumulator( + &self, + args: AccumulatorArgs, + ) -> datafusion_common::Result> { + self.inner.accumulator(args) + } + + fn state_fields( + &self, + args: StateFieldsArgs, + ) -> datafusion_common::Result> { + self.inner.state_fields(args) + } + + fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + self.inner.groups_accumulator_supported(args) + } + + fn create_groups_accumulator( + &self, + args: AccumulatorArgs, + ) -> datafusion_common::Result> { + self.inner.create_groups_accumulator(args) + } + + fn supports_null_handling_clause(&self) -> bool { + true + } +} + pub(crate) extern "C" fn create_ffi_sum_func() -> FFI_AggregateUDF { - let udaf: Arc = Arc::new(Sum::new().into()); + let udaf: Arc = Arc::new(SumWithNullHandling::default().into()); udaf.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..1f8f47c4a6200 100644 --- a/datafusion/ffi/tests/ffi_udaf.rs +++ b/datafusion/ffi/tests/ffi_udaf.rs @@ -66,6 +66,21 @@ mod tests { Ok(()) } + #[test] + fn test_supports_null_handling_clause() -> Result<()> { + let module = get_module()?; + + 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()); + + let ffi_stddev_func = (module.create_stddev_udaf)(); + let foreign_stddev_func: Arc = (&ffi_stddev_func).into(); + assert!(!foreign_stddev_func.supports_null_handling_clause()); + + Ok(()) + } + #[tokio::test] async fn test_ffi_grouping_udaf() -> Result<()> { let module = get_module()?; From bb1e57e74e03ee756343cdae1cff432bb7ff1aa0 Mon Sep 17 00:00:00 2001 From: Amogh Ramesh Date: Sun, 26 Jul 2026 20:55:57 +0530 Subject: [PATCH 2/2] test(ffi): use real null-handling aggregate --- datafusion/ffi/src/tests/mod.rs | 9 +++- datafusion/ffi/src/tests/udf_udaf_udwf.rs | 63 ++++------------------- datafusion/ffi/tests/ffi_udaf.rs | 11 ++-- 3 files changed, 22 insertions(+), 61 deletions(-) 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 e2a5609b635bb..b5ed406420376 100644 --- a/datafusion/ffi/src/tests/udf_udaf_udwf.rs +++ b/datafusion/ffi/src/tests/udf_udaf_udwf.rs @@ -17,21 +17,18 @@ use std::sync::Arc; -use arrow_schema::{DataType, FieldRef}; +use arrow_schema::DataType; use datafusion_catalog::TableFunctionImpl; use datafusion_common::ScalarValue; use datafusion_expr::{ - Accumulator, AggregateUDF, AggregateUDFImpl, ColumnarValue, ExpressionPlacement, - GroupsAccumulator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, - Volatility, WindowUDF, + AggregateUDF, ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDF, + ScalarUDFImpl, Signature, Volatility, WindowUDF, }; 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_aggregate_common::accumulator::{ - AccumulatorArgs, StateFieldsArgs, -}; use datafusion_functions_table::generate_series::RangeFunc; use datafusion_functions_window::rank::Rank; @@ -174,56 +171,14 @@ pub(crate) extern "C" fn create_ffi_table_func( FFI_TableFunction::new_with_ffi_codec(udtf, None, codec) } -#[derive(Debug, Default, Hash, Eq, PartialEq)] -struct SumWithNullHandling { - inner: Sum, -} - -impl AggregateUDFImpl for SumWithNullHandling { - fn name(&self) -> &str { - self.inner.name() - } - - fn signature(&self) -> &Signature { - self.inner.signature() - } - - fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result { - self.inner.return_type(arg_types) - } - - fn accumulator( - &self, - args: AccumulatorArgs, - ) -> datafusion_common::Result> { - self.inner.accumulator(args) - } - - fn state_fields( - &self, - args: StateFieldsArgs, - ) -> datafusion_common::Result> { - self.inner.state_fields(args) - } - - fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - self.inner.groups_accumulator_supported(args) - } - - fn create_groups_accumulator( - &self, - args: AccumulatorArgs, - ) -> datafusion_common::Result> { - self.inner.create_groups_accumulator(args) - } +pub(crate) extern "C" fn create_ffi_sum_func() -> FFI_AggregateUDF { + let udaf: Arc = Arc::new(Sum::new().into()); - fn supports_null_handling_clause(&self) -> bool { - true - } + udaf.into() } -pub(crate) extern "C" fn create_ffi_sum_func() -> FFI_AggregateUDF { - let udaf: Arc = Arc::new(SumWithNullHandling::default().into()); +pub(crate) extern "C" fn create_ffi_first_value_func() -> FFI_AggregateUDF { + let udaf: Arc = Arc::new(FirstValue::new().into()); udaf.into() } diff --git a/datafusion/ffi/tests/ffi_udaf.rs b/datafusion/ffi/tests/ffi_udaf.rs index 1f8f47c4a6200..090151416e4e9 100644 --- a/datafusion/ffi/tests/ffi_udaf.rs +++ b/datafusion/ffi/tests/ffi_udaf.rs @@ -70,13 +70,14 @@ mod tests { 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()); - - let ffi_stddev_func = (module.create_stddev_udaf)(); - let foreign_stddev_func: Arc = (&ffi_stddev_func).into(); - assert!(!foreign_stddev_func.supports_null_handling_clause()); + assert!(!foreign_sum_func.supports_null_handling_clause()); Ok(()) }