Add support for getting inner ffi types on foreign types that use boxed traits

Author: timsaucer

datafusion/expr-common/src/accumulator.rs

@@ -48,7 +48,7 @@ use std::fmt::Debug;
 /// [`evaluate`]: Self::evaluate
 /// [`merge_batch`]: Self::merge_batch
 /// [window function]: https://en.wikipedia.org/wiki/Window_function_(SQL)
-pub trait Accumulator: Send + Sync + Debug {
+pub trait Accumulator: Send + Sync + Debug + std::any::Any {
     /// Updates the accumulator's state from its input.
     ///
     /// `values` contains the arguments to this aggregate function.

datafusion/expr-common/src/groups_accumulator.rs

@@ -108,7 +108,7 @@ impl EmitTo {
 ///
 /// [`Accumulator`]: crate::accumulator::Accumulator
 /// [Aggregating Millions of Groups Fast blog]: https://arrow.apache.org/blog/2023/08/05/datafusion_fast_grouping/
-pub trait GroupsAccumulator: Send {
+pub trait GroupsAccumulator: Send + std::any::Any {
     /// Updates the accumulator's state from its arguments, encoded as
     /// a vector of [`ArrayRef`]s.
     ///

datafusion/expr/src/partition_evaluator.rs

@@ -90,7 +90,7 @@ use crate::window_state::WindowAggState;
 /// For more background, please also see the [User defined Window Functions in DataFusion blog]
 ///
 /// [User defined Window Functions in DataFusion blog]: https://datafusion.apache.org/blog/2025/04/19/user-defined-window-functions
-pub trait PartitionEvaluator: Debug + Send {
+pub trait PartitionEvaluator: Debug + Send + std::any::Any {
     /// When the window frame has a fixed beginning (e.g UNBOUNDED
     /// PRECEDING), some functions such as FIRST_VALUE, LAST_VALUE and
     /// NTH_VALUE do not need the (unbounded) input once they have

datafusion/ffi/src/udaf/accumulator.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::any::Any;
 use std::ffi::c_void;
 use std::ops::Deref;
 use std::ptr::null_mut;
@@ -204,9 +205,12 @@ unsafe extern "C" fn release_fn_wrapper(accumulator: &mut FFI_Accumulator) {
 
 impl From<Box<dyn Accumulator>> for FFI_Accumulator {
     fn from(accumulator: Box<dyn Accumulator>) -> Self {
-        // if let Some(accumulator) = accumulator.into_any().downcast::<ForeignAccumulator>() {
-        //     return accumulator.accumulator;
-        // }
+        if (accumulator.as_ref() as &dyn Any).is::<ForeignAccumulator>() {
+            let accumulator = (accumulator as Box<dyn Any>)
+                .downcast::<ForeignAccumulator>()
+                .expect("already checked type");
+            return accumulator.accumulator;
+        }
 
         let supports_retract_batch = accumulator.supports_retract_batch();
         let private_data = AccumulatorPrivateData { accumulator };

datafusion/ffi/src/udaf/groups_accumulator.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::any::Any;
 use std::ffi::c_void;
 use std::ops::Deref;
 use std::ptr::null_mut;
@@ -245,9 +246,12 @@ unsafe extern "C" fn release_fn_wrapper(accumulator: &mut FFI_GroupsAccumulator)
 
 impl From<Box<dyn GroupsAccumulator>> for FFI_GroupsAccumulator {
     fn from(accumulator: Box<dyn GroupsAccumulator>) -> Self {
-        // if let Some(accumulator) = accumulator.into_any().downcast_ref::<ForeignGroupsAccumulator>() {
-        //     return accumulator.accumulator;
-        // }
+        if (accumulator.as_ref() as &dyn Any).is::<ForeignGroupsAccumulator>() {
+            let accumulator = (accumulator as Box<dyn Any>)
+                .downcast::<ForeignGroupsAccumulator>()
+                .expect("already checked type");
+            return accumulator.accumulator;
+        }
 
         let supports_convert_to_state = accumulator.supports_convert_to_state();
         let private_data = GroupsAccumulatorPrivateData { accumulator };

datafusion/ffi/src/udwf/partition_evaluator.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::any::Any;
 use std::ffi::c_void;
 use std::ops::Range;
 
@@ -205,9 +206,12 @@ unsafe extern "C" fn release_fn_wrapper(evaluator: &mut FFI_PartitionEvaluator)
 
 impl From<Box<dyn PartitionEvaluator>> for FFI_PartitionEvaluator {
     fn from(evaluator: Box<dyn PartitionEvaluator>) -> Self {
-        // if let Ok(evaluator) = evaluator.into_any().downcast::<ForeignPartitionEvaluator>() {
-        //     return evaluator.evaluator;
-        // }
+        if (evaluator.as_ref() as &dyn Any).is::<ForeignPartitionEvaluator>() {
+            let evaluator = (evaluator as Box<dyn Any>)
+                .downcast::<ForeignPartitionEvaluator>()
+                .expect("already checked type");
+            return evaluator.evaluator;
+        }
 
         let is_causal = evaluator.is_causal();
         let supports_bounded_execution = evaluator.supports_bounded_execution();

datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs

@@ -37,7 +37,7 @@ use super::accumulate::NullState;
 #[derive(Debug)]
 pub struct BooleanGroupsAccumulator<F>
 where
-    F: Fn(bool, bool) -> bool + Send + Sync,
+    F: Fn(bool, bool) -> bool + Send + Sync + 'static,
 {
     /// values per group
     values: BooleanBufferBuilder,
@@ -55,7 +55,7 @@ where
 
 impl<F> BooleanGroupsAccumulator<F>
 where
-    F: Fn(bool, bool) -> bool + Send + Sync,
+    F: Fn(bool, bool) -> bool + Send + Sync + 'static,
 {
     pub fn new(bool_fn: F, identity: bool) -> Self {
         Self {
@@ -69,7 +69,7 @@ where
 
 impl<F> GroupsAccumulator for BooleanGroupsAccumulator<F>
 where
-    F: Fn(bool, bool) -> bool + Send + Sync,
+    F: Fn(bool, bool) -> bool + Send + Sync + 'static,
 {
     fn update_batch(
         &mut self,

datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs

@@ -41,7 +41,7 @@ use super::accumulate::NullState;
 pub struct PrimitiveGroupsAccumulator<T, F>
 where
     T: ArrowPrimitiveType + Send,
-    F: Fn(&mut T::Native, T::Native) + Send + Sync,
+    F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static,
 {
     /// values per group, stored as the native type
     values: Vec<T::Native>,
@@ -62,7 +62,7 @@ where
 impl<T, F> PrimitiveGroupsAccumulator<T, F>
 where
     T: ArrowPrimitiveType + Send,
-    F: Fn(&mut T::Native, T::Native) + Send + Sync,
+    F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static,
 {
     pub fn new(data_type: &DataType, prim_fn: F) -> Self {
         Self {
@@ -84,7 +84,7 @@ where
 impl<T, F> GroupsAccumulator for PrimitiveGroupsAccumulator<T, F>
 where
     T: ArrowPrimitiveType + Send,
-    F: Fn(&mut T::Native, T::Native) + Send + Sync,
+    F: Fn(&mut T::Native, T::Native) + Send + Sync + 'static,
 {
     fn update_batch(
         &mut self,

datafusion/functions-aggregate/src/average.rs

@@ -754,7 +754,7 @@ impl Accumulator for DurationAvgAccumulator {
 struct AvgGroupsAccumulator<T, F>
 where
     T: ArrowNumericType + Send,
-    F: Fn(T::Native, u64) -> Result<T::Native> + Send,
+    F: Fn(T::Native, u64) -> Result<T::Native> + Send + 'static,
 {
     /// The type of the internal sum
     sum_data_type: DataType,
@@ -778,7 +778,7 @@ where
 impl<T, F> AvgGroupsAccumulator<T, F>
 where
     T: ArrowNumericType + Send,
-    F: Fn(T::Native, u64) -> Result<T::Native> + Send,
+    F: Fn(T::Native, u64) -> Result<T::Native> + Send + 'static,
 {
     pub fn new(sum_data_type: &DataType, return_data_type: &DataType, avg_fn: F) -> Self {
         debug!(
@@ -800,7 +800,7 @@ where
 impl<T, F> GroupsAccumulator for AvgGroupsAccumulator<T, F>
 where
     T: ArrowNumericType + Send,
-    F: Fn(T::Native, u64) -> Result<T::Native> + Send,
+    F: Fn(T::Native, u64) -> Result<T::Native> + Send + 'static,
 {
     fn update_batch(
         &mut self,

datafusion/spark/src/function/aggregate/avg.rs

@@ -213,7 +213,7 @@ impl Accumulator for AvgAccumulator {
 struct AvgGroupsAccumulator<T, F>
 where
     T: ArrowNumericType + Send,
-    F: Fn(T::Native, i64) -> Result<T::Native> + Send,
+    F: Fn(T::Native, i64) -> Result<T::Native> + Send + 'static,
 {
     /// The type of the returned average
     return_data_type: DataType,
@@ -231,7 +231,7 @@ where
 impl<T, F> AvgGroupsAccumulator<T, F>
 where
     T: ArrowNumericType + Send,
-    F: Fn(T::Native, i64) -> Result<T::Native> + Send,
+    F: Fn(T::Native, i64) -> Result<T::Native> + Send + 'static,
 {
     pub fn new(return_data_type: &DataType, avg_fn: F) -> Self {
         Self {
@@ -246,7 +246,7 @@ where
 impl<T, F> GroupsAccumulator for AvgGroupsAccumulator<T, F>
 where
     T: ArrowNumericType + Send,
-    F: Fn(T::Native, i64) -> Result<T::Native> + Send,
+    F: Fn(T::Native, i64) -> Result<T::Native> + Send + 'static,
 {
     fn update_batch(
         &mut self,