Make FirstValue an UDAF, Change AggregateUDFImpl::accumulator signature, support ORDER BY for UDAFs

datafusion-examples/examples/advanced_udaf.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow_schema::Schema;
 use datafusion::{arrow::datatypes::DataType, logical_expr::Volatility};
 use datafusion_physical_expr::NullState;
 use std::{any::Any, sync::Arc};
@@ -85,7 +86,14 @@ impl AggregateUDFImpl for GeoMeanUdaf {
     /// is supported, DataFusion will use this row oriented
     /// accumulator when the aggregate function is used as a window function
     /// or when there are only aggregates (no GROUP BY columns) in the plan.
-    fn accumulator(&self, _arg: &DataType) -> Result<Box<dyn Accumulator>> {
+    fn accumulator(
+        &self,
+        _arg: &DataType,
+        _sort_exprs: &[Expr],
+        _schema: &Schema,
+        _ignore_nulls: bool,
+        _requirement_satisfied: bool,
+    ) -> Result<Box<dyn Accumulator>> {
         Ok(Box::new(GeometricMean::new()))
     }
 
@@ -191,7 +199,7 @@ impl Accumulator for GeometricMean {
 
 // create local session context with an in-memory table
 fn create_context() -> Result<SessionContext> {
-    use datafusion::arrow::datatypes::{Field, Schema};
+    use datafusion::arrow::datatypes::Field;
     use datafusion::datasource::MemTable;
     // define a schema.
     let schema = Arc::new(Schema::new(vec![

datafusion-examples/examples/simple_udaf.rs

@@ -150,7 +150,7 @@ async fn main() -> Result<()> {
         Arc::new(DataType::Float64),
         Volatility::Immutable,
         // This is the accumulator factory; DataFusion uses it to create new accumulators.
-        Arc::new(|_| Ok(Box::new(GeometricMean::new()))),
+        Arc::new(|_, _, _| Ok(Box::new(GeometricMean::new()))),
         // This is the description of the state. `state()` must match the types here.
         Arc::new(vec![DataType::Float64, DataType::UInt32]),
     );

datafusion/core/src/execution/context/mod.rs

@@ -69,11 +69,14 @@ use datafusion_common::{
     OwnedTableReference, SchemaReference,
 };
 use datafusion_execution::registry::SerializerRegistry;
+use datafusion_expr::type_coercion::aggregates::NUMERICS;
+use datafusion_expr::{create_first_value, Signature, Volatility};
 use datafusion_expr::{
     logical_plan::{DdlStatement, Statement},
     var_provider::is_system_variables,
     Expr, StringifiedPlan, UserDefinedLogicalNode, WindowUDF,
 };
+use datafusion_physical_expr::create_first_value_accumulator;
 use datafusion_sql::{
     parser::{CopyToSource, CopyToStatement, DFParser},
     planner::{object_name_to_table_reference, ContextProvider, ParserOptions, SqlToRel},
@@ -1457,6 +1460,13 @@ impl SessionState {
         datafusion_functions_array::register_all(&mut new_self)
             .expect("can not register array expressions");
 
+        let first_value = create_first_value(
+            "FIRST_VALUE",
+            Signature::uniform(1, NUMERICS.to_vec(), Volatility::Immutable),
+            Arc::new(create_first_value_accumulator),
+        );
+        let _ = new_self.register_udaf(Arc::new(first_value));
+
         new_self
     }
     /// Returns new [`SessionState`] using the provided

datafusion/core/src/physical_planner.rs

@@ -247,24 +247,20 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result<String> {
             distinct,
             args,
             filter,
-            order_by,
+            order_by: _,
             null_treatment: _,
         }) => match func_def {
             AggregateFunctionDefinition::BuiltIn(..) => {
                 create_function_physical_name(func_def.name(), *distinct, args)
             }
             AggregateFunctionDefinition::UDF(fun) => {
-                // TODO: Add support for filter and order by in AggregateUDF
+                // TODO: Add support for filter by in AggregateUDF
                 if filter.is_some() {
                     return exec_err!(
                         "aggregate expression with filter is not supported"
                     );
                 }
-                if order_by.is_some() {
-                    return exec_err!(
-                        "aggregate expression with order_by is not supported"
-                    );
-                }
+
                 let names = args
                     .iter()
                     .map(|e| create_physical_name(e, false))
@@ -1657,8 +1653,11 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
             order_by,
             null_treatment,
         }) => {
-            let args =
-                create_physical_exprs(args, logical_input_schema, execution_props)?;
+            let args = args
+                .iter()
+                .map(|e| create_physical_expr(e, logical_input_schema, execution_props))
+                .collect::<Result<Vec<_>>>()?;
+
             let filter = match filter {
                 Some(e) => Some(create_physical_expr(
                     e,
@@ -1667,20 +1666,28 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
                 )?),
                 None => None,
             };
-            let order_by = match order_by {
-                Some(e) => Some(create_physical_sort_exprs(
-                    e,
-                    logical_input_schema,
-                    execution_props,
-                )?),
-                None => None,
-            };
             let ignore_nulls = null_treatment
                 .unwrap_or(sqlparser::ast::NullTreatment::RespectNulls)
                 == NullTreatment::IgnoreNulls;
             let (agg_expr, filter, order_by) = match func_def {
                 AggregateFunctionDefinition::BuiltIn(fun) => {
-                    let ordering_reqs = order_by.clone().unwrap_or(vec![]);
+                    let physical_sort_exprs = match order_by {
+                        Some(e) => Some(
+                            e.iter()
+                                .map(|expr| {
+                                    create_physical_sort_expr(
+                                        expr,
+                                        logical_input_schema,
+                                        execution_props,
+                                    )
+                                })
+                                .collect::<Result<Vec<_>>>()?,
+                        ),
+                        None => None,
+                    };
+
+                    let ordering_reqs: Vec<PhysicalSortExpr> =
+                        physical_sort_exprs.clone().unwrap_or(vec![]);
                     let agg_expr = aggregates::create_aggregate_expr(
                         fun,
                         *distinct,
@@ -1690,16 +1697,37 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
                         name,
                         ignore_nulls,
                     )?;
-                    (agg_expr, filter, order_by)
+                    (agg_expr, filter, physical_sort_exprs)
                 }
                 AggregateFunctionDefinition::UDF(fun) => {
+                    let sort_exprs = order_by.clone().unwrap_or(vec![]);
+                    let physical_sort_exprs = match order_by {
+                        Some(e) => Some(
+                            e.iter()
+                                .map(|expr| {
+                                    create_physical_sort_expr(
+                                        expr,
+                                        logical_input_schema,
+                                        execution_props,
+                                    )
+                                })
+                                .collect::<Result<Vec<_>>>()?,
+                        ),
+                        None => None,
+                    };
+
+                    let ordering_reqs: Vec<PhysicalSortExpr> =
+                        physical_sort_exprs.clone().unwrap_or(vec![]);
                     let agg_expr = udaf::create_aggregate_expr(
                         fun,
                         &args,
+                        &sort_exprs,
+                        &ordering_reqs,
                         physical_input_schema,
                         name,
-                    );
-                    (agg_expr?, filter, order_by)
+                        ignore_nulls,
+                    )?;
+                    (agg_expr, filter, physical_sort_exprs)
                 }
                 AggregateFunctionDefinition::Name(_) => {
                     return internal_err!(

datafusion/core/tests/user_defined/user_defined_aggregates.rs

@@ -45,7 +45,7 @@ use datafusion::{
 };
 use datafusion_common::{assert_contains, cast::as_primitive_array, exec_err};
 use datafusion_expr::{
-    create_udaf, AggregateUDFImpl, GroupsAccumulator, SimpleAggregateUDF,
+    create_udaf, AggregateUDFImpl, Expr, GroupsAccumulator, SimpleAggregateUDF,
 };
 use datafusion_physical_expr::expressions::AvgAccumulator;
 
@@ -234,7 +234,7 @@ async fn simple_udaf() -> Result<()> {
         vec![DataType::Float64],
         Arc::new(DataType::Float64),
         Volatility::Immutable,
-        Arc::new(|_| Ok(Box::<AvgAccumulator>::default())),
+        Arc::new(|_, _, _| Ok(Box::<AvgAccumulator>::default())),
         Arc::new(vec![DataType::UInt64, DataType::Float64]),
     );
 
@@ -262,7 +262,7 @@ async fn deregister_udaf() -> Result<()> {
         vec![DataType::Float64],
         Arc::new(DataType::Float64),
         Volatility::Immutable,
-        Arc::new(|_| Ok(Box::<AvgAccumulator>::default())),
+        Arc::new(|_, _, _| Ok(Box::<AvgAccumulator>::default())),
         Arc::new(vec![DataType::UInt64, DataType::Float64]),
     );
 
@@ -290,7 +290,7 @@ async fn case_sensitive_identifiers_user_defined_aggregates() -> Result<()> {
         vec![DataType::Float64],
         Arc::new(DataType::Float64),
         Volatility::Immutable,
-        Arc::new(|_| Ok(Box::<AvgAccumulator>::default())),
+        Arc::new(|_, _, _| Ok(Box::<AvgAccumulator>::default())),
         Arc::new(vec![DataType::UInt64, DataType::Float64]),
     );
 
@@ -333,7 +333,7 @@ async fn test_user_defined_functions_with_alias() -> Result<()> {
         vec![DataType::Float64],
         Arc::new(DataType::Float64),
         Volatility::Immutable,
-        Arc::new(|_| Ok(Box::<AvgAccumulator>::default())),
+        Arc::new(|_, _, _| Ok(Box::<AvgAccumulator>::default())),
         Arc::new(vec![DataType::UInt64, DataType::Float64]),
     )
     .with_aliases(vec!["dummy_alias"]);
@@ -497,7 +497,7 @@ impl TimeSum {
 
         let captured_state = Arc::clone(&test_state);
         let accumulator: AccumulatorFactoryFunction =
-            Arc::new(move |_| Ok(Box::new(Self::new(Arc::clone(&captured_state)))));
+            Arc::new(move |_, _, _| Ok(Box::new(Self::new(Arc::clone(&captured_state)))));
 
         let time_sum = AggregateUDF::from(SimpleAggregateUDF::new(
             name,
@@ -596,7 +596,7 @@ impl FirstSelector {
         let signatures = vec![TypeSignature::Exact(Self::input_datatypes())];
 
         let accumulator: AccumulatorFactoryFunction =
-            Arc::new(|_| Ok(Box::new(Self::new())));
+            Arc::new(|_, _, _| Ok(Box::new(Self::new())));
 
         let volatility = Volatility::Immutable;
 
@@ -717,7 +717,14 @@ impl AggregateUDFImpl for TestGroupsAccumulator {
         Ok(DataType::UInt64)
     }
 
-    fn accumulator(&self, _arg: &DataType) -> Result<Box<dyn Accumulator>> {
+    fn accumulator(
+        &self,
+        _arg: &DataType,
+        _sort_exprs: &[Expr],
+        _schema: &Schema,
+        _ignore_nulls: bool,
+        _requirement_satisfied: bool,
+    ) -> Result<Box<dyn Accumulator>> {
         // should use groups accumulator
         panic!("accumulator shouldn't invoke");
     }

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -291,7 +291,7 @@ async fn udaf_as_window_func() -> Result<()> {
         vec![DataType::Int32],
         Arc::new(DataType::Int32),
         Volatility::Immutable,
-        Arc::new(|_| Ok(Box::new(MyAccumulator))),
+        Arc::new(|_, _, _| Ok(Box::new(MyAccumulator))),
         Arc::new(vec![DataType::Int32]),
     );
 

datafusion/expr/src/expr.rs

@@ -577,14 +577,15 @@ impl AggregateFunction {
         distinct: bool,
         filter: Option<Box<Expr>>,
         order_by: Option<Vec<Expr>>,
+        null_treatment: Option<NullTreatment>,
     ) -> Self {
         Self {
             func_def: AggregateFunctionDefinition::UDF(udf),
             args,
             distinct,
             filter,
             order_by,
-            null_treatment: None,
+            null_treatment,
         }
     }
 }