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};
@@ -30,7 +31,8 @@ use datafusion::error::Result;
 use datafusion::prelude::*;
 use datafusion_common::{cast::as_float64_array, ScalarValue};
 use datafusion_expr::{
-    Accumulator, AggregateUDF, AggregateUDFImpl, GroupsAccumulator, Signature,
+    function::AccumulatorArgs, Accumulator, AggregateUDF, AggregateUDFImpl,
+    GroupsAccumulator, Signature,
 };
 
 /// This example shows how to use the full AggregateUDFImpl API to implement a user
@@ -85,7 +87,7 @@ 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, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
         Ok(Box::new(GeometricMean::new()))
     }
 
@@ -191,7 +193,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/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},
@@ -82,6 +85,7 @@ use datafusion_sql::{
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
+use log::debug;
 use parking_lot::RwLock;
 use sqlparser::dialect::dialect_from_str;
 use url::Url;
@@ -1457,6 +1461,22 @@ 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),
+        );
+
+        match new_self.register_udaf(Arc::new(first_value)) {
+            Ok(Some(existing_udaf)) => {
+                debug!("Overwrite existing UDF: {}", existing_udaf.name());
+            }
+            Ok(None) => {}
+            Err(err) => {
+                panic!("Failed to register UDF: {}", err);
+            }
+        }
+
         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))
@@ -1667,20 +1663,34 @@ 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 try_create_physical_sort_expr =
+                |order_by: &Option<Vec<Expr>>| -> Result<Option<Vec<PhysicalSortExpr>>> {
+                    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,
+                    };
+                    Ok(physical_sort_exprs)
+                };
+
             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 = try_create_physical_sort_expr(order_by)?;
+                    let ordering_reqs: Vec<PhysicalSortExpr> =
+                        physical_sort_exprs.clone().unwrap_or(vec![]);
                     let agg_expr = aggregates::create_aggregate_expr(
                         fun,
                         *distinct,
@@ -1690,16 +1700,23 @@ 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 = try_create_physical_sort_expr(order_by)?;
+                    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,8 @@ use datafusion::{
 };
 use datafusion_common::{assert_contains, cast::as_primitive_array, exec_err};
 use datafusion_expr::{
-    create_udaf, AggregateUDFImpl, GroupsAccumulator, SimpleAggregateUDF,
+    create_udaf, function::AccumulatorArgs, AggregateUDFImpl, GroupsAccumulator,
+    SimpleAggregateUDF,
 };
 use datafusion_physical_expr::expressions::AvgAccumulator;
 
@@ -717,7 +718,7 @@ impl AggregateUDFImpl for TestGroupsAccumulator {
         Ok(DataType::UInt64)
     }
 
-    fn accumulator(&self, _arg: &DataType) -> Result<Box<dyn Accumulator>> {
+    fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
         // should use groups accumulator
         panic!("accumulator shouldn't invoke");
     }

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,
         }
     }
 }

datafusion/expr/src/expr_fn.rs

@@ -21,16 +21,18 @@ use crate::expr::{
     AggregateFunction, BinaryExpr, Cast, Exists, GroupingSet, InList, InSubquery,
     Placeholder, ScalarFunction, TryCast,
 };
-use crate::function::PartitionEvaluatorFactory;
+use crate::function::{
+    AccumulatorArgs, AccumulatorFactoryFunction, PartitionEvaluatorFactory,
+};
+use crate::udaf::format_state_name;
 use crate::{
     aggregate_function, built_in_function, conditional_expressions::CaseBuilder,
-    logical_plan::Subquery, AccumulatorFactoryFunction, AggregateUDF,
-    BuiltinScalarFunction, Expr, LogicalPlan, Operator, ScalarFunctionImplementation,
-    ScalarUDF, Signature, Volatility,
+    logical_plan::Subquery, AggregateUDF, BuiltinScalarFunction, Expr, LogicalPlan,
+    Operator, ScalarFunctionImplementation, ScalarUDF, Signature, Volatility,
 };
 use crate::{AggregateUDFImpl, ColumnarValue, ScalarUDFImpl, WindowUDF, WindowUDFImpl};
-use arrow::datatypes::DataType;
-use datafusion_common::{Column, Result};
+use arrow::datatypes::{DataType, Field};
+use datafusion_common::{internal_err, Column, Result};
 use std::any::Any;
 use std::fmt::Debug;
 use std::ops::Not;
@@ -719,6 +721,16 @@ pub fn create_udaf(
     ))
 }
 
+/// Creates a new UDAF with a specific signature, state type and return type.
+/// The signature and state type must match the `Accumulator's implementation`.
+pub fn create_first_value(
+    name: &str,
+    signature: Signature,
+    accumulator: AccumulatorFactoryFunction,
+) -> AggregateUDF {
+    AggregateUDF::from(FirstValue::new(name, signature, accumulator))
+}
+
 /// Implements [`AggregateUDFImpl`] for functions that have a single signature and
 /// return type.
 pub struct SimpleAggregateUDF {
@@ -796,15 +808,94 @@ impl AggregateUDFImpl for SimpleAggregateUDF {
         Ok(self.return_type.clone())
     }
 
-    fn accumulator(&self, arg: &DataType) -> Result<Box<dyn crate::Accumulator>> {
-        (self.accumulator)(arg)
+    fn accumulator(
+        &self,
+        acc_args: AccumulatorArgs,
+    ) -> Result<Box<dyn crate::Accumulator>> {
+        (self.accumulator)(acc_args)
     }
 
     fn state_type(&self, _return_type: &DataType) -> Result<Vec<DataType>> {
         Ok(self.state_type.clone())
     }
 }
 
+pub struct FirstValue {
+    name: String,
+    signature: Signature,
+    accumulator: AccumulatorFactoryFunction,
+}
+
+impl Debug for FirstValue {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        f.debug_struct("FirstValue")
+            .field("name", &self.name)
+            .field("signature", &self.signature)
+            .field("fun", &"<FUNC>")
+            .finish()
+    }
+}
+
+impl FirstValue {
+    pub fn new(
+        name: impl Into<String>,
+        signature: Signature,
+        accumulator: AccumulatorFactoryFunction,
+    ) -> Self {
+        let name = name.into();
+        Self {
+            name,
+            signature,
+            accumulator,
+        }
+    }
+}
+
+impl AggregateUDFImpl for FirstValue {
+    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> {
+        Ok(arg_types[0].clone())
+    }
+
+    fn accumulator(
+        &self,
+        acc_args: AccumulatorArgs,
+    ) -> Result<Box<dyn crate::Accumulator>> {
+        (self.accumulator)(acc_args)
+    }
+
+    fn state_type(&self, _return_type: &DataType) -> Result<Vec<DataType>> {
+        internal_err!("FirstValue does not have a state type")
+    }
+
+    fn state_fields(
+        &self,
+        name: &str,
+        value_type: DataType,
+        ordering_fields: Vec<Field>,
+    ) -> Result<Vec<Field>> {
+        let mut fields = vec![Field::new(
+            format_state_name(name, "first_value"),
+            value_type,
+            true,
+        )];
+        fields.extend(ordering_fields);
+        fields.push(Field::new("is_set", DataType::Boolean, true));
+        Ok(fields)
+    }
+}
+
 /// Creates a new UDWF with a specific signature, state type and return type.
 ///
 /// The signature and state type must match the [`PartitionEvaluator`]'s implementation`.