More work on moving to Field from DataType for aggregates

Author: timsaucer

datafusion-examples/examples/advanced_udaf.rs

@@ -94,7 +94,7 @@ impl AggregateUDFImpl for GeoMeanUdaf {
     /// This is the description of the state. accumulator's state() must match the types here.
     fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
         Ok(vec![
-            Field::new("prod", args.return_type.clone(), true),
+            Field::new("prod", args.return_field.data_type().clone(), true),
             Field::new("n", DataType::UInt32, true),
         ])
     }

datafusion/core/tests/user_defined/user_defined_aggregates.rs

@@ -18,6 +18,8 @@
 //! This module contains end to end demonstrations of creating
 //! user defined aggregate functions
 
+use std::any::Any;
+use std::collections::HashMap;
 use std::hash::{DefaultHasher, Hash, Hasher};
 use std::mem::{size_of, size_of_val};
 use std::sync::{
@@ -26,10 +28,10 @@ use std::sync::{
 };
 
 use arrow::array::{
-    types::UInt64Type, AsArray, Int32Array, PrimitiveArray, StringArray, StructArray,
+    types::UInt64Type, Array, AsArray, Int32Array, PrimitiveArray, StringArray,
+    StructArray, UInt64Array,
 };
 use arrow::datatypes::{Fields, Schema};
-
 use datafusion::common::test_util::batches_to_string;
 use datafusion::dataframe::DataFrame;
 use datafusion::datasource::MemTable;
@@ -48,7 +50,7 @@ use datafusion::{
     prelude::SessionContext,
     scalar::ScalarValue,
 };
-use datafusion_common::assert_contains;
+use datafusion_common::{assert_contains, exec_datafusion_err};
 use datafusion_common::{cast::as_primitive_array, exec_err};
 use datafusion_expr::{
     col, create_udaf, function::AccumulatorArgs, AggregateUDFImpl, GroupsAccumulator,
@@ -781,7 +783,7 @@ struct TestGroupsAccumulator {
 }
 
 impl AggregateUDFImpl for TestGroupsAccumulator {
-    fn as_any(&self) -> &dyn std::any::Any {
+    fn as_any(&self) -> &dyn Any {
         self
     }
 
@@ -890,3 +892,134 @@ impl GroupsAccumulator for TestGroupsAccumulator {
         size_of::<u64>()
     }
 }
+
+#[derive(Debug)]
+struct MetadataBasedAggregateUdf {
+    name: String,
+    signature: Signature,
+    metadata: HashMap<String, String>,
+}
+
+impl MetadataBasedAggregateUdf {
+    fn new(metadata: HashMap<String, String>) -> Self {
+        // The name we return must be unique. Otherwise we will not call distinct
+        // instances of this UDF. This is a small hack for the unit tests to get unique
+        // names, but you could do something more elegant with the metadata.
+        let name = format!("metadata_based_udf_{}", metadata.len());
+        Self {
+            name,
+            signature: Signature::exact(vec![DataType::UInt64], Volatility::Immutable),
+            metadata,
+        }
+    }
+}
+
+impl AggregateUDFImpl for MetadataBasedAggregateUdf {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        &self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        unimplemented!("this should never be called since return_field is implemented");
+    }
+
+    fn return_field(&self, _arg_fields: &[Field]) -> Result<Field> {
+        Ok(Field::new(self.name(), DataType::UInt64, true)
+            .with_metadata(self.metadata.clone()))
+    }
+
+    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+        let input_expr = acc_args
+            .exprs
+            .first()
+            .ok_or(exec_datafusion_err!("Expected one argument"))?;
+        let input_field = input_expr.return_field(acc_args.schema)?;
+
+        let double_output = input_field
+            .metadata()
+            .get("modify_values")
+            .map(|v| v == "double_output")
+            .unwrap_or(false);
+
+        Ok(Box::new(MetadataBasedAccumulator {
+            double_output,
+            curr_sum: 0,
+        }))
+    }
+}
+
+#[derive(Debug)]
+struct MetadataBasedAccumulator {
+    double_output: bool,
+    curr_sum: u64,
+}
+
+impl Accumulator for MetadataBasedAccumulator {
+    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+        let arr = values[0]
+            .as_any()
+            .downcast_ref::<UInt64Array>()
+            .ok_or(exec_datafusion_err!("Expected UInt64Array"))?;
+
+        self.curr_sum = arr.iter().fold(self.curr_sum, |a, b| a + b.unwrap_or(0));
+
+        Ok(())
+    }
+
+    fn evaluate(&mut self) -> Result<ScalarValue> {
+        let v = match self.double_output {
+            true => self.curr_sum * 2,
+            false => self.curr_sum,
+        };
+
+        Ok(ScalarValue::from(v))
+    }
+
+    fn size(&self) -> usize {
+        9
+    }
+
+    fn state(&mut self) -> Result<Vec<ScalarValue>> {
+        Ok(vec![ScalarValue::from(self.curr_sum)])
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
+        self.update_batch(states)
+    }
+}
+
+#[tokio::test]
+async fn test_metadata_based_accumulator() -> Result<()> {
+    let ctx = SessionContext::new();
+    let arr = UInt64Array::from(vec![1]);
+    let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(arr) as _)])?;
+    ctx.register_batch("t", batch).unwrap();
+
+    let udaf_no_metadata =
+        AggregateUDF::from(MetadataBasedAggregateUdf::new(HashMap::new()));
+    let udaf_with_metadata = AggregateUDF::from(MetadataBasedAggregateUdf::new(
+        HashMap::from_iter([("modify_values".to_string(), "double_output".to_string())]),
+    ));
+    ctx.register_udaf(udaf_no_metadata);
+    ctx.register_udaf(udaf_with_metadata);
+
+    let sql_df = ctx
+        .sql("SELECT metadata_based_udf_0(a) FROM t group by a")
+        .await?;
+    sql_df.show().await?;
+
+    let sql_df = ctx
+        .sql("SELECT metadata_based_udf_1(a) FROM t group by a")
+        .await?;
+    sql_df.show().await?;
+
+    Ok(())
+}

datafusion/expr-common/src/type_coercion/aggregates.rs

@@ -17,7 +17,7 @@
 
 use crate::signature::TypeSignature;
 use arrow::datatypes::{
-    DataType, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE,
+    DataType, Field, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE,
     DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE,
 };
 
@@ -89,7 +89,7 @@ pub static TIMES: &[DataType] = &[
 /// number of input types.
 pub fn check_arg_count(
     func_name: &str,
-    input_types: &[DataType],
+    input_types: &[Field],
     signature: &TypeSignature,
 ) -> Result<()> {
     match signature {

datafusion/expr/src/udaf.rs

@@ -410,7 +410,7 @@ where
 ///    fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { unimplemented!() }
 ///    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
 ///        Ok(vec![
-///             Field::new("value", args.return_type.clone(), true),
+///             args.return_field.clone().with_name("value"),
 ///             Field::new("ordering", DataType::UInt32, true)
 ///        ])
 ///    }
@@ -745,11 +745,10 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
     /// be derived from `name`. See [`format_state_name`] for a utility function
     /// to generate a unique name.
     fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Field>> {
-        let fields = vec![Field::new(
-            format_state_name(args.name, "value"),
-            args.return_type.clone(),
-            true,
-        )];
+        let fields = vec![args
+            .return_field
+            .clone()
+            .with_name(format_state_name(args.name, "value"))];
 
         Ok(fields
             .into_iter()

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

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::datatypes::{DataType, Field, Schema};
+use arrow::datatypes::{Field, Schema};
 use datafusion_common::Result;
 use datafusion_expr_common::accumulator::Accumulator;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
@@ -28,7 +28,7 @@ use std::sync::Arc;
 #[derive(Debug)]
 pub struct AccumulatorArgs<'a> {
     /// The return type of the aggregate function.
-    pub return_type: &'a DataType,
+    pub return_field: &'a Field,
 
     /// The schema of the input arguments
     pub schema: &'a Schema,
@@ -81,11 +81,11 @@ pub struct StateFieldsArgs<'a> {
     /// The name of the aggregate function.
     pub name: &'a str,
 
-    /// The input types of the aggregate function.
-    pub input_types: &'a [DataType],
+    /// The input fields of the aggregate function.
+    pub input_fields: &'a [Field],
 
-    /// The return type of the aggregate function.
-    pub return_type: &'a DataType,
+    /// The return fields of the aggregate function.
+    pub return_field: &'a Field,
 
     /// The ordering fields of the aggregate function.
     pub ordering_fields: &'a [Field],

datafusion/functions-aggregate/benches/count.rs

@@ -28,7 +28,7 @@ use std::sync::Arc;
 fn prepare_accumulator() -> Box<dyn GroupsAccumulator> {
     let schema = Arc::new(Schema::new(vec![Field::new("f", DataType::Int32, true)]));
     let accumulator_args = AccumulatorArgs {
-        return_type: &DataType::Int64,
+        return_field: &Field::new("f", DataType::Int64, true),
         schema: &schema,
         ignore_nulls: false,
         ordering_req: &LexOrdering::default(),

datafusion/functions-aggregate/benches/sum.rs

@@ -26,9 +26,10 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use std::sync::Arc;
 
 fn prepare_accumulator(data_type: &DataType) -> Box<dyn GroupsAccumulator> {
-    let schema = Arc::new(Schema::new(vec![Field::new("f", data_type.clone(), true)]));
+    let field = Field::new("f", data_type.clone(), true);
+    let schema = Arc::new(Schema::new(vec![field.clone()]));
     let accumulator_args = AccumulatorArgs {
-        return_type: data_type,
+        return_field: &field,
         schema: &schema,
         ignore_nulls: false,
         ordering_req: &LexOrdering::default(),

datafusion/functions-aggregate/src/array_agg.rs

@@ -114,15 +114,15 @@ impl AggregateUDFImpl for ArrayAgg {
             return Ok(vec![Field::new_list(
                 format_state_name(args.name, "distinct_array_agg"),
                 // See COMMENTS.md to understand why nullable is set to true
-                Field::new_list_field(args.input_types[0].clone(), true),
+                Field::new_list_field(args.input_fields[0].data_type().clone(), true),
                 true,
             )]);
         }
 
         let mut fields = vec![Field::new_list(
             format_state_name(args.name, "array_agg"),
             // See COMMENTS.md to understand why nullable is set to true
-            Field::new_list_field(args.input_types[0].clone(), true),
+            Field::new_list_field(args.input_fields[0].data_type().clone(), true),
             true,
         )];
 
@@ -984,7 +984,7 @@ mod tests {
     }
 
     struct ArrayAggAccumulatorBuilder {
-        data_type: DataType,
+        return_field: Field,
         distinct: bool,
         ordering: LexOrdering,
         schema: Schema,
@@ -997,7 +997,7 @@ mod tests {
 
         fn new(data_type: DataType) -> Self {
             Self {
-                data_type: data_type.clone(),
+                return_field: Field::new("f", data_type.clone(), true),
                 distinct: false,
                 ordering: Default::default(),
                 schema: Schema {
@@ -1029,7 +1029,7 @@ mod tests {
 
         fn build(&self) -> Result<Box<dyn Accumulator>> {
             ArrayAgg::default().accumulator(AccumulatorArgs {
-                return_type: &self.data_type,
+                return_field: &self.return_field,
                 schema: &self.schema,
                 ignore_nulls: false,
                 ordering_req: &self.ordering,