fix(ffi): forward distinct_handling through FFI_AggregateUDF

datafusion/ffi/src/udaf/mod.rs

@@ -24,7 +24,8 @@ use datafusion_common::{DataFusionError, Result, ffi_datafusion_err};
 use datafusion_expr::function::AggregateFunctionSimplification;
 use datafusion_expr::type_coercion::functions::fields_with_udf;
 use datafusion_expr::{
-    Accumulator, AggregateUDF, AggregateUDFImpl, GroupsAccumulator, Signature,
+    Accumulator, AggregateUDF, AggregateUDFImpl, DistinctHandling, GroupsAccumulator,
+    Signature,
 };
 use datafusion_functions_aggregate_common::accumulator::{
     AccumulatorArgs, StateFieldsArgs,
@@ -149,6 +150,10 @@ pub struct FFI_AggregateUDF {
     /// FFI equivalent to [`AggregateUDF::supports_null_handling_clause`]
     pub supports_null_handling_clause:
         unsafe extern "C" fn(udaf: &FFI_AggregateUDF) -> bool,
+
+    /// FFI equivalent to [`AggregateUDF::distinct_handling`]
+    pub distinct_handling:
+        unsafe extern "C" fn(udaf: &FFI_AggregateUDF) -> FFI_DistinctHandling,
 }
 
 unsafe impl Send for FFI_AggregateUDF {}
@@ -337,6 +342,12 @@ unsafe extern "C" fn supports_null_handling_clause_fn_wrapper(
     unsafe { udaf.inner().supports_null_handling_clause() }
 }
 
+unsafe extern "C" fn distinct_handling_fn_wrapper(
+    udaf: &FFI_AggregateUDF,
+) -> FFI_DistinctHandling {
+    unsafe { udaf.inner().distinct_handling().into() }
+}
+
 unsafe extern "C" fn coerce_types_fn_wrapper(
     udaf: &FFI_AggregateUDF,
     arg_types: SVec<WrappedSchema>,
@@ -412,6 +423,7 @@ impl From<Arc<AggregateUDF>> for FFI_AggregateUDF {
             private_data: Box::into_raw(private_data).cast::<c_void>(),
             library_marker_id: crate::get_library_marker_id,
             supports_null_handling_clause: supports_null_handling_clause_fn_wrapper,
+            distinct_handling: distinct_handling_fn_wrapper,
         }
     }
 }
@@ -610,6 +622,10 @@ impl AggregateUDFImpl for ForeignAggregateUDF {
         unsafe { (self.udaf.supports_null_handling_clause)(&self.udaf) }
     }
 
+    fn distinct_handling(&self) -> DistinctHandling {
+        unsafe { (self.udaf.distinct_handling)(&self.udaf).into() }
+    }
+
     fn simplify(&self) -> Option<AggregateFunctionSimplification> {
         None
     }
@@ -655,6 +671,37 @@ impl From<AggregateOrderSensitivity> for FFI_AggregateOrderSensitivity {
     }
 }
 
+#[repr(C)]
+#[derive(Debug)]
+pub enum FFI_DistinctHandling {
+    Insensitive,
+    Sensitive,
+    Unsupported,
+}
+
+impl From<FFI_DistinctHandling> for DistinctHandling {
+    fn from(value: FFI_DistinctHandling) -> Self {
+        match value {
+            FFI_DistinctHandling::Insensitive => Self::Insensitive,
+            FFI_DistinctHandling::Sensitive => Self::Sensitive,
+            FFI_DistinctHandling::Unsupported => Self::Unsupported,
+        }
+    }
+}
+
+impl From<DistinctHandling> for FFI_DistinctHandling {
+    fn from(value: DistinctHandling) -> Self {
+        match value {
+            DistinctHandling::Insensitive => Self::Insensitive,
+            DistinctHandling::Sensitive => Self::Sensitive,
+            DistinctHandling::Unsupported => Self::Unsupported,
+            // `DistinctHandling` is non_exhaustive, so map any variant added
+            // later to the default.
+            _ => Self::Sensitive,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
 
@@ -798,6 +845,24 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_distinct_handling() -> Result<()> {
+        let min = create_test_foreign_udaf(
+            datafusion::functions_aggregate::min_max::Min::new(),
+        )?;
+        assert_eq!(min.distinct_handling(), DistinctHandling::Insensitive);
+
+        let sum = create_test_foreign_udaf(Sum::new())?;
+        assert_eq!(sum.distinct_handling(), DistinctHandling::Sensitive);
+
+        let stddev = create_test_foreign_udaf(
+            datafusion::functions_aggregate::stddev::Stddev::new(),
+        )?;
+        assert_eq!(stddev.distinct_handling(), DistinctHandling::Unsupported);
+
+        Ok(())
+    }
+
     #[test]
     fn test_beneficial_ordering() -> Result<()> {
         let foreign_udaf = create_test_foreign_udaf(
@@ -870,6 +935,18 @@ mod tests {
         test_round_trip_order_sensitivity(AggregateOrderSensitivity::Beneficial);
     }
 
+    #[test]
+    fn test_round_trip_all_distinct_handlings() {
+        for handling in [
+            DistinctHandling::Insensitive,
+            DistinctHandling::Sensitive,
+            DistinctHandling::Unsupported,
+        ] {
+            let ffi_handling: FFI_DistinctHandling = handling.into();
+            assert_eq!(DistinctHandling::from(ffi_handling), handling);
+        }
+    }
+
     #[test]
     fn test_ffi_udaf_local_bypass() -> Result<()> {
         let original_udaf = Sum::new();

datafusion/ffi/tests/ffi_udaf.rs

@@ -23,7 +23,7 @@ mod tests {
 
     use arrow::array::{Float64Array, record_batch};
     use datafusion::error::Result;
-    use datafusion::logical_expr::{AggregateUDF, AggregateUDFImpl};
+    use datafusion::logical_expr::{AggregateUDF, AggregateUDFImpl, DistinctHandling};
     use datafusion::prelude::{SessionContext, col};
     use datafusion_catalog::MemTable;
     use datafusion_expr::{ScalarUDF, ScalarUDFImpl};
@@ -82,6 +82,27 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_distinct_handling() -> Result<()> {
+        let module = get_module()?;
+
+        let ffi_stddev_func = (module.create_stddev_udaf)();
+        let foreign_stddev_func: Arc<dyn AggregateUDFImpl> = (&ffi_stddev_func).into();
+        assert_eq!(
+            foreign_stddev_func.distinct_handling(),
+            DistinctHandling::Unsupported
+        );
+
+        let ffi_sum_func = (module.create_sum_udaf)();
+        let foreign_sum_func: Arc<dyn AggregateUDFImpl> = (&ffi_sum_func).into();
+        assert_eq!(
+            foreign_sum_func.distinct_handling(),
+            DistinctHandling::Sensitive
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_ffi_grouping_udaf() -> Result<()> {
         let module = get_module()?;