Use return_field instead of return_type for calling aggregates via FFI (apache#17407)

* Use return_field instead of return_type for calling aggregates via FFI

* Add note in upgrade guide

* typo

Author: timsaucer

datafusion/ffi/src/udaf/mod.rs

@@ -69,11 +69,11 @@ pub struct FFI_AggregateUDF {
     /// FFI equivalent to the `volatility` of a [`AggregateUDF`]
     pub volatility: FFI_Volatility,
 
-    /// Determines the return type of the underlying [`AggregateUDF`] based on the
-    /// argument types.
-    pub return_type: unsafe extern "C" fn(
+    /// Determines the return field of the underlying [`AggregateUDF`] based on the
+    /// argument fields.
+    pub return_field: unsafe extern "C" fn(
         udaf: &Self,
-        arg_types: RVec<WrappedSchema>,
+        arg_fields: RVec<WrappedSchema>,
     ) -> RResult<WrappedSchema, RString>,
 
     /// FFI equivalent to the `is_nullable` of a [`AggregateUDF`]
@@ -160,20 +160,22 @@ impl FFI_AggregateUDF {
     }
 }
 
-unsafe extern "C" fn return_type_fn_wrapper(
+unsafe extern "C" fn return_field_fn_wrapper(
     udaf: &FFI_AggregateUDF,
-    arg_types: RVec<WrappedSchema>,
+    arg_fields: RVec<WrappedSchema>,
 ) -> RResult<WrappedSchema, RString> {
     let udaf = udaf.inner();
 
-    let arg_types = rresult_return!(rvec_wrapped_to_vec_datatype(&arg_types));
+    let arg_fields = rresult_return!(rvec_wrapped_to_vec_fieldref(&arg_fields));
 
-    let return_type = udaf
-        .return_type(&arg_types)
-        .and_then(|v| FFI_ArrowSchema::try_from(v).map_err(DataFusionError::from))
+    let return_field = udaf
+        .return_field(&arg_fields)
+        .and_then(|v| {
+            FFI_ArrowSchema::try_from(v.as_ref()).map_err(DataFusionError::from)
+        })
         .map(WrappedSchema);
 
-    rresult!(return_type)
+    rresult!(return_field)
 }
 
 unsafe extern "C" fn accumulator_fn_wrapper(
@@ -346,7 +348,7 @@ impl From<Arc<AggregateUDF>> for FFI_AggregateUDF {
             is_nullable,
             volatility,
             aliases,
-            return_type: return_type_fn_wrapper,
+            return_field: return_field_fn_wrapper,
             accumulator: accumulator_fn_wrapper,
             create_sliding_accumulator: create_sliding_accumulator_fn_wrapper,
             create_groups_accumulator: create_groups_accumulator_fn_wrapper,
@@ -425,14 +427,22 @@ impl AggregateUDFImpl for ForeignAggregateUDF {
         &self.signature
     }
 
-    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        let arg_types = vec_datatype_to_rvec_wrapped(arg_types)?;
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        unimplemented!()
+    }
+
+    fn return_field(&self, arg_fields: &[FieldRef]) -> Result<FieldRef> {
+        let arg_fields = vec_fieldref_to_rvec_wrapped(arg_fields)?;
 
-        let result = unsafe { (self.udaf.return_type)(&self.udaf, arg_types) };
+        let result = unsafe { (self.udaf.return_field)(&self.udaf, arg_fields) };
 
         let result = df_result!(result);
 
-        result.and_then(|r| (&r.0).try_into().map_err(DataFusionError::from))
+        result.and_then(|r| {
+            Field::try_from(&r.0)
+                .map(Arc::new)
+                .map_err(DataFusionError::from)
+        })
     }
 
     fn is_nullable(&self) -> bool {
@@ -608,9 +618,43 @@ mod tests {
         physical_expr::PhysicalSortExpr, physical_plan::expressions::col,
         scalar::ScalarValue,
     };
+    use std::any::Any;
+    use std::collections::HashMap;
 
     use super::*;
 
+    #[derive(Default, Debug, Hash, Eq, PartialEq)]
+    struct SumWithCopiedMetadata {
+        inner: Sum,
+    }
+
+    impl AggregateUDFImpl for SumWithCopiedMetadata {
+        fn as_any(&self) -> &dyn Any {
+            self
+        }
+
+        fn name(&self) -> &str {
+            self.inner.name()
+        }
+
+        fn signature(&self) -> &Signature {
+            self.inner.signature()
+        }
+
+        fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+            unimplemented!()
+        }
+
+        fn return_field(&self, arg_fields: &[FieldRef]) -> Result<FieldRef> {
+            // Copy the input field, so any metadata gets returned
+            Ok(Arc::clone(&arg_fields[0]))
+        }
+
+        fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+            self.inner.accumulator(acc_args)
+        }
+    }
+
     fn create_test_foreign_udaf(
         original_udaf: impl AggregateUDFImpl + 'static,
     ) -> Result<AggregateUDF> {
@@ -644,8 +688,11 @@ mod tests {
         let foreign_udaf =
             create_test_foreign_udaf(Sum::new())?.with_aliases(["my_function"]);
 
-        let return_type = foreign_udaf.return_type(&[DataType::Float64])?;
-        assert_eq!(return_type, DataType::Float64);
+        let return_field =
+            foreign_udaf
+                .return_field(&[Field::new("a", DataType::Float64, true).into()])?;
+        let return_type = return_field.data_type();
+        assert_eq!(return_type, &DataType::Float64);
         Ok(())
     }
 
@@ -673,6 +720,31 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_round_trip_udaf_metadata() -> Result<()> {
+        let original_udaf = SumWithCopiedMetadata::default();
+        let original_udaf = Arc::new(AggregateUDF::from(original_udaf));
+
+        // Convert to FFI format
+        let local_udaf: FFI_AggregateUDF = Arc::clone(&original_udaf).into();
+
+        // Convert back to native format
+        let foreign_udaf: ForeignAggregateUDF = (&local_udaf).try_into()?;
+        let foreign_udaf: AggregateUDF = foreign_udaf.into();
+
+        let metadata: HashMap<String, String> =
+            [("a_key".to_string(), "a_value".to_string())]
+                .into_iter()
+                .collect();
+        let input_field = Arc::new(
+            Field::new("a", DataType::Float64, false).with_metadata(metadata.clone()),
+        );
+        let return_field = foreign_udaf.return_field(&[input_field])?;
+
+        assert_eq!(&metadata, return_field.metadata());
+        Ok(())
+    }
+
     #[test]
     fn test_beneficial_ordering() -> Result<()> {
         let foreign_udaf = create_test_foreign_udaf(

docs/source/library-user-guide/upgrading.md

@@ -285,6 +285,24 @@ If you have custom implementations of `FileOpener` or work directly with `FileOp
 
 [#17397]: https://github.com/apache/datafusion/pull/17397
 
+### FFI user defined aggregate function signature change
+
+The Foreign Function Interface (FFI) signature for user defined aggregate functions
+has been updated to call `return_field` instead of `return_type` on the underlying
+aggregate function. This is to support metadata handling with these aggregate functions.
+This change should be transparent to most users. If you have written unit tests to call
+`return_type` directly, you may need to change them to calling `return_field` instead.
+
+This update is a breaking change to the FFI API. The current best practice when using the
+FFI crate is to ensure that all libraries that are interacting are using the same
+underlying Rust version. Issue [#17374] has been opened to discuss stabilization of
+this interface so that these libraries can be used across different DataFusion versions.
+
+See [#17407] for details.
+
+[#17407]: https://github.com/apache/datafusion/pull/17407
+[#17374]: https://github.com/apache/datafusion/issues/17374
+
 ### Added `PhysicalExpr::is_volatile_node`
 
 We added a method to `PhysicalExpr` to mark a `PhysicalExpr` as volatile: