feat: Implement SchemaBasedAggregateUdf and enhance accumulator argument handling

Author: kosiew

datafusion/core/tests/user_defined/user_defined_aggregates.rs

@@ -1014,6 +1014,51 @@ impl Accumulator for MetadataBasedAccumulator {
     }
 }
 
+#[derive(Debug)]
+struct SchemaBasedAggregateUdf {
+    signature: Signature,
+}
+
+impl SchemaBasedAggregateUdf {
+    fn new() -> Self {
+        Self {
+            signature: Signature::any(1, Volatility::Immutable),
+        }
+    }
+}
+
+impl AggregateUDFImpl for SchemaBasedAggregateUdf {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "schema_based_udf"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::UInt64)
+    }
+
+    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+        let field = acc_args.schema.field(0).clone();
+        let double_output = field
+            .metadata()
+            .get("modify_values")
+            .map(|v| v == "double_output")
+            .unwrap_or(false);
+
+        Ok(Box::new(MetadataBasedAccumulator {
+            double_output,
+            curr_sum: 0,
+        }))
+    }
+}
+
 #[tokio::test]
 async fn test_metadata_based_aggregate() -> Result<()> {
     let data_array = Arc::new(UInt64Array::from(vec![0, 5, 10, 15, 20])) as ArrayRef;
@@ -1170,3 +1215,34 @@ async fn test_metadata_based_aggregate_as_window() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_scalar_udaf_schema_metadata() -> Result<()> {
+    use datafusion_expr::{expr::FieldMetadata, lit_with_metadata};
+    use std::collections::BTreeMap;
+
+    let ctx = SessionContext::new();
+    let udf = AggregateUDF::from(SchemaBasedAggregateUdf::new());
+
+    let metadata = FieldMetadata::new(BTreeMap::from([(
+        "modify_values".to_string(),
+        "double_output".to_string(),
+    )]));
+
+    let expr = udf
+        .call(vec![lit_with_metadata(
+            ScalarValue::UInt64(Some(1)),
+            Some(metadata),
+        )])
+        .alias("res");
+
+    let plan = LogicalPlanBuilder::empty(true)
+        .aggregate(Vec::<Expr>::new(), vec![expr])?
+        .build()?;
+
+    let df = DataFrame::new(ctx.state(), plan);
+    let batches = df.collect().await?;
+    let array = batches[0].column(0).as_primitive::<UInt64Type>();
+    assert_eq!(array.value(0), 2);
+    Ok(())
+}

datafusion/expr/src/udaf.rs

@@ -754,7 +754,12 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
     /// group during query execution.
     ///
     /// acc_args: [`AccumulatorArgs`] contains information about how the
-    /// aggregate function was called.
+    /// aggregate function was called. Use `acc_args.exprs` together with
+    /// `acc_args.schema` to inspect the [`FieldRef`] of each input. When an
+    /// aggregate is invoked with literal values only, `acc_args.schema` is
+    /// synthesized from those literals so that any field metadata (for
+    /// example Arrow extension types) is available to the accumulator
+    /// implementation.
     fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>>;
 
     /// Return the fields used to store the intermediate state of this accumulator.

datafusion/functions-aggregate-common/src/accumulator.rs

@@ -30,7 +30,13 @@ pub struct AccumulatorArgs<'a> {
     /// The return field of the aggregate function.
     pub return_field: FieldRef,
 
-    /// The schema of the input arguments
+    /// The schema of the input arguments.
+    ///
+    /// This schema contains the fields corresponding to [`Self::exprs`]. When
+    /// an aggregate is invoked with only literal values, this schema is
+    /// synthesized from those literals so that field-level metadata (such as
+    /// Arrow extension types) remains available to accumulator
+    /// implementations.
     pub schema: &'a Schema,
 
     /// Whether to ignore nulls.
@@ -65,7 +71,9 @@ pub struct AccumulatorArgs<'a> {
     /// ```
     pub is_distinct: bool,
 
-    /// The physical expression of arguments the aggregate function takes.
+    /// The physical expressions for the aggregate function's arguments.
+    /// Use these expressions together with [`Self::schema`] to obtain the
+    /// [`FieldRef`] of each input via `expr.return_field(schema)`.
     pub exprs: &'a [Arc<dyn PhysicalExpr>],
 }
 

datafusion/physical-expr/src/aggregate.rs

@@ -34,6 +34,7 @@ pub mod utils {
     };
 }
 
+use std::borrow::Cow;
 use std::fmt::Debug;
 use std::sync::Arc;
 
@@ -376,13 +377,35 @@ impl AggregateFunctionExpr {
             .into()
     }
 
+    /// Returns a schema containing the fields corresponding to this
+    /// aggregate's input expressions.
+    ///
+    /// When an aggregate is invoked with only literal values, the
+    /// physical input schema is empty. In this case a schema is
+    /// synthesized from the literal expressions so that field metadata
+    /// (such as Arrow extension types) remains available to the
+    /// [`Accumulator`].
+    fn args_schema(&self) -> Cow<'_, Schema> {
+        if self.schema.fields().is_empty() {
+            Cow::Owned(Schema::new(
+                self.input_fields
+                    .iter()
+                    .map(|f| f.as_ref().clone())
+                    .collect::<Vec<_>>(),
+            ))
+        } else {
+            Cow::Borrowed(&self.schema)
+        }
+    }
+
     /// the accumulator used to accumulate values from the expressions.
     /// the accumulator expects the same number of arguments as `expressions` and must
     /// return states with the same description as `state_fields`
     pub fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+        let schema = self.args_schema();
         let acc_args = AccumulatorArgs {
             return_field: Arc::clone(&self.return_field),
-            schema: &self.schema,
+            schema: schema.as_ref(),
             ignore_nulls: self.ignore_nulls,
             order_bys: self.order_bys.as_ref(),
             is_distinct: self.is_distinct,
@@ -464,9 +487,10 @@ impl AggregateFunctionExpr {
 
     /// Creates accumulator implementation that supports retract
     pub fn create_sliding_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+        let schema = self.args_schema();
         let args = AccumulatorArgs {
             return_field: Arc::clone(&self.return_field),
-            schema: &self.schema,
+            schema: schema.as_ref(),
             ignore_nulls: self.ignore_nulls,
             order_bys: self.order_bys.as_ref(),
             is_distinct: self.is_distinct,
@@ -533,9 +557,10 @@ impl AggregateFunctionExpr {
     /// [`GroupsAccumulator`] implementation. If this returns true,
     /// `[Self::create_groups_accumulator`] will be called.
     pub fn groups_accumulator_supported(&self) -> bool {
+        let schema = self.args_schema();
         let args = AccumulatorArgs {
             return_field: Arc::clone(&self.return_field),
-            schema: &self.schema,
+            schema: schema.as_ref(),
             ignore_nulls: self.ignore_nulls,
             order_bys: self.order_bys.as_ref(),
             is_distinct: self.is_distinct,
@@ -552,9 +577,10 @@ impl AggregateFunctionExpr {
     /// For maximum performance, a [`GroupsAccumulator`] should be
     /// implemented in addition to [`Accumulator`].
     pub fn create_groups_accumulator(&self) -> Result<Box<dyn GroupsAccumulator>> {
+        let schema = self.args_schema();
         let args = AccumulatorArgs {
             return_field: Arc::clone(&self.return_field),
-            schema: &self.schema,
+            schema: schema.as_ref(),
             ignore_nulls: self.ignore_nulls,
             order_bys: self.order_bys.as_ref(),
             is_distinct: self.is_distinct,