Deprecate builtin first/last aggregate function and use UDAF

datafusion/core/Cargo.toml

@@ -102,6 +102,7 @@ datafusion-functions-aggregate = { workspace = true }
 datafusion-functions-array = { workspace = true, optional = true }
 datafusion-optimizer = { workspace = true }
 datafusion-physical-expr = { workspace = true }
+datafusion-physical-expr-common = { workspace = true }
 datafusion-physical-plan = { workspace = true }
 datafusion-sql = { workspace = true }
 flate2 = { version = "1.0.24", optional = true }

datafusion/core/src/physical_optimizer/convert_first_last.rs

@@ -20,11 +20,12 @@ use datafusion_common::{
     config::ConfigOptions,
     tree_node::{Transformed, TransformedResult, TreeNode},
 };
-use datafusion_physical_expr::expressions::{FirstValue, LastValue};
-use datafusion_physical_expr::{
-    equivalence::ProjectionMapping, reverse_order_bys, AggregateExpr,
-    EquivalenceProperties, PhysicalSortRequirement,
-};
+use datafusion_expr::Expr;
+use datafusion_functions_aggregate::first_last::{FirstValue, LastValue};
+use datafusion_physical_expr::{equivalence::ProjectionMapping, EquivalenceProperties};
+use datafusion_physical_expr_common::aggregate::{AggregateExpr, AggregateFunctionExpr};
+use datafusion_physical_expr_common::sort_expr::PhysicalSortRequirement;
+use datafusion_physical_expr_common::utils::reverse_order_bys;
 use datafusion_physical_plan::aggregates::concat_slices;
 use datafusion_physical_plan::{
     aggregates::{AggregateExec, AggregateMode},
@@ -199,15 +200,15 @@ fn try_convert_first_last_if_better(
     eq_properties: &EquivalenceProperties,
 ) -> Result<()> {
     for aggr_expr in aggr_exprs.iter_mut() {
-        let aggr_req = aggr_expr.order_bys().unwrap_or(&[]);
-        let reverse_aggr_req = reverse_order_bys(aggr_req);
-        let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_req);
+        let aggr_fun_expr = aggr_expr.as_any().downcast_ref::<AggregateFunctionExpr>();
+        let aggr_sort_expr = aggr_expr.order_bys().unwrap_or(&[]);
+        let reversed_ordering = reverse_order_bys(aggr_sort_expr);
+        let aggr_req = PhysicalSortRequirement::from_sort_exprs(aggr_sort_expr);
         let reverse_aggr_req =
-            PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req);
-
-        if let Some(first_value) = aggr_expr.as_any().downcast_ref::<FirstValue>() {
-            let mut first_value = first_value.clone();
+            PhysicalSortRequirement::from_sort_exprs(&reversed_ordering);
 
+        if aggr_fun_expr.is_some_and(|e| e.fun().name() == "FIRST_VALUE") {
+            let mut first_value = aggr_fun_expr.unwrap().clone();
             if eq_properties.ordering_satisfy_requirement(&concat_slices(
                 prefix_requirement,
                 &aggr_req,
@@ -218,10 +219,26 @@ fn try_convert_first_last_if_better(
                 prefix_requirement,
                 &reverse_aggr_req,
             )) {
+                let name = first_value.name();
+
+                let name = if let Some(n) = name.strip_prefix("FIRST") {
+                    format!("LAST{}", n)
+                } else {
+                    let expr = first_value.expressions().swap_remove(0);
+                    format!("LAST_VALUE({})", expr)
+                };
+
+                let reversed_sort_exprs = reverse_sort_exprs(&first_value);
+
                 // Converting to LAST_VALUE enables more efficient execution
                 // given the existing ordering:
-                let mut last_value = first_value.convert_to_last();
-                last_value = last_value.with_requirement_satisfied(true);
+                let last_value = first_value
+                    .with_name(name)
+                    .with_fun(LastValue::new().into())
+                    .with_sort_exprs(reversed_sort_exprs)
+                    .with_ordering_req(reversed_ordering)
+                    .with_requirement_satisfied(true);
+
                 *aggr_expr = Arc::new(last_value) as _;
             } else {
                 // Requirement is not satisfied with existing ordering.
@@ -230,8 +247,8 @@ fn try_convert_first_last_if_better(
             }
             continue;
         }
-        if let Some(last_value) = aggr_expr.as_any().downcast_ref::<LastValue>() {
-            let mut last_value = last_value.clone();
+        if aggr_fun_expr.is_some_and(|e| e.fun().name() == "LAST_VALUE") {
+            let mut last_value = aggr_fun_expr.unwrap().clone();
             if eq_properties.ordering_satisfy_requirement(&concat_slices(
                 prefix_requirement,
                 &aggr_req,
@@ -244,8 +261,25 @@ fn try_convert_first_last_if_better(
             )) {
                 // Converting to FIRST_VALUE enables more efficient execution
                 // given the existing ordering:
-                let mut first_value = last_value.convert_to_first();
-                first_value = first_value.with_requirement_satisfied(true);
+
+                let name = last_value.name();
+
+                let name = if let Some(n) = name.strip_prefix("LAST") {
+                    format!("FIRST{}", n)
+                } else {
+                    let expr = last_value.expressions().swap_remove(0);
+                    format!("FIRST_VALUE({})", expr)
+                };
+
+                let reversed_sort_exprs = reverse_sort_exprs(&last_value);
+
+                let first_value = last_value
+                    .with_name(name)
+                    .with_fun(FirstValue::new().into())
+                    .with_sort_exprs(reversed_sort_exprs)
+                    .with_ordering_req(reversed_ordering)
+                    .with_requirement_satisfied(true);
+
                 *aggr_expr = Arc::new(first_value) as _;
             } else {
                 // Requirement is not satisfied with existing ordering.
@@ -258,3 +292,16 @@ fn try_convert_first_last_if_better(
 
     Ok(())
 }
+
+fn reverse_sort_exprs(e: &AggregateFunctionExpr) -> Vec<Expr> {
+    e.sort_exprs()
+        .iter()
+        .map(|e| {
+            if let Expr::Sort(s) = e {
+                Expr::Sort(s.reverse())
+            } else {
+                e.clone()
+            }
+        })
+        .collect::<Vec<_>>()
+}

datafusion/core/src/prelude.rs

@@ -39,6 +39,7 @@ pub use datafusion_expr::{
     Expr,
 };
 pub use datafusion_functions::expr_fn::*;
+pub use datafusion_functions_aggregate::expr_fn::*;
 #[cfg(feature = "array_expressions")]
 pub use datafusion_functions_array::expr_fn::*;
 

datafusion/expr/src/aggregate_function.rs

@@ -49,10 +49,6 @@ pub enum AggregateFunction {
     ApproxDistinct,
     /// Aggregation into an array
     ArrayAgg,
-    /// First value in a group according to some ordering
-    FirstValue,
-    /// Last value in a group according to some ordering
-    LastValue,
     /// N'th value in a group according to some ordering
     NthValue,
     /// Variance (Sample)
@@ -121,8 +117,6 @@ impl AggregateFunction {
             Median => "MEDIAN",
             ApproxDistinct => "APPROX_DISTINCT",
             ArrayAgg => "ARRAY_AGG",
-            FirstValue => "FIRST_VALUE",
-            LastValue => "LAST_VALUE",
             NthValue => "NTH_VALUE",
             Variance => "VAR",
             VariancePop => "VAR_POP",
@@ -178,8 +172,6 @@ impl FromStr for AggregateFunction {
             "min" => AggregateFunction::Min,
             "sum" => AggregateFunction::Sum,
             "array_agg" => AggregateFunction::ArrayAgg,
-            "first_value" => AggregateFunction::FirstValue,
-            "last_value" => AggregateFunction::LastValue,
             "nth_value" => AggregateFunction::NthValue,
             "string_agg" => AggregateFunction::StringAgg,
             // statistical
@@ -294,9 +286,7 @@ impl AggregateFunction {
                 Ok(coerced_data_types[0].clone())
             }
             AggregateFunction::Grouping => Ok(DataType::Int32),
-            AggregateFunction::FirstValue
-            | AggregateFunction::LastValue
-            | AggregateFunction::NthValue => Ok(coerced_data_types[0].clone()),
+            AggregateFunction::NthValue => Ok(coerced_data_types[0].clone()),
             AggregateFunction::StringAgg => Ok(DataType::LargeUtf8),
         }
     }
@@ -351,9 +341,7 @@ impl AggregateFunction {
             | AggregateFunction::Stddev
             | AggregateFunction::StddevPop
             | AggregateFunction::Median
-            | AggregateFunction::ApproxMedian
-            | AggregateFunction::FirstValue
-            | AggregateFunction::LastValue => {
+            | AggregateFunction::ApproxMedian => {
                 Signature::uniform(1, NUMERICS.to_vec(), Volatility::Immutable)
             }
             AggregateFunction::NthValue => Signature::any(2, Volatility::Immutable),

datafusion/expr/src/expr.rs

@@ -520,6 +520,15 @@ impl Sort {
             nulls_first,
         }
     }
+
+    /// Create a new Sort expression with the opposite sort direction
+    pub fn reverse(&self) -> Self {
+        Self {
+            expr: self.expr.clone(),
+            asc: !self.asc,
+            nulls_first: self.nulls_first,
+        }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]

datafusion/expr/src/lib.rs

@@ -79,7 +79,8 @@ pub use logical_plan::*;
 pub use operator::Operator;
 pub use partition_evaluator::PartitionEvaluator;
 pub use signature::{
-    FuncMonotonicity, Signature, TypeSignature, Volatility, TIMEZONE_WILDCARD,
+    ArrayFunctionSignature, FuncMonotonicity, Signature, TypeSignature, Volatility,
+    TIMEZONE_WILDCARD,
 };
 pub use table_source::{TableProviderFilterPushDown, TableSource, TableType};
 pub use udaf::{AggregateUDF, AggregateUDFImpl};

datafusion/expr/src/type_coercion/aggregates.rs

@@ -93,7 +93,7 @@ pub fn coerce_types(
 ) -> Result<Vec<DataType>> {
     use DataType::*;
     // Validate input_types matches (at least one of) the func signature.
-    check_arg_count(agg_fun, input_types, &signature.type_signature)?;
+    check_arg_count(agg_fun.name(), input_types, &signature.type_signature)?;
 
     match agg_fun {
         AggregateFunction::Count | AggregateFunction::ApproxDistinct => {
@@ -293,10 +293,9 @@ pub fn coerce_types(
             }
             Ok(input_types.to_vec())
         }
-        AggregateFunction::Median
-        | AggregateFunction::FirstValue
-        | AggregateFunction::LastValue => Ok(input_types.to_vec()),
-        AggregateFunction::NthValue => Ok(input_types.to_vec()),
+        AggregateFunction::Median | AggregateFunction::NthValue => {
+            Ok(input_types.to_vec())
+        }
         AggregateFunction::Grouping => Ok(vec![input_types[0].clone()]),
         AggregateFunction::StringAgg => {
             if !is_string_agg_supported_arg_type(&input_types[0]) {
@@ -323,17 +322,16 @@ pub fn coerce_types(
 /// This method DOES NOT validate the argument types - only that (at least one,
 /// in the case of [`TypeSignature::OneOf`]) signature matches the desired
 /// number of input types.
-fn check_arg_count(
-    agg_fun: &AggregateFunction,
+pub fn check_arg_count(
+    func_name: &str,
     input_types: &[DataType],
     signature: &TypeSignature,
 ) -> Result<()> {
     match signature {
         TypeSignature::Uniform(agg_count, _) | TypeSignature::Any(agg_count) => {
             if input_types.len() != *agg_count {
                 return plan_err!(
-                    "The function {:?} expects {:?} arguments, but {:?} were provided",
-                    agg_fun,
+                    "The function {func_name} expects {:?} arguments, but {:?} were provided",
                     agg_count,
                     input_types.len()
                 );
@@ -342,8 +340,7 @@ fn check_arg_count(
         TypeSignature::Exact(types) => {
             if types.len() != input_types.len() {
                 return plan_err!(
-                    "The function {:?} expects {:?} arguments, but {:?} were provided",
-                    agg_fun,
+                    "The function {func_name} expects {:?} arguments, but {:?} were provided",
                     types.len(),
                     input_types.len()
                 );
@@ -352,19 +349,18 @@ fn check_arg_count(
         TypeSignature::OneOf(variants) => {
             let ok = variants
                 .iter()
-                .any(|v| check_arg_count(agg_fun, input_types, v).is_ok());
+                .any(|v| check_arg_count(func_name, input_types, v).is_ok());
             if !ok {
                 return plan_err!(
-                    "The function {:?} does not accept {:?} function arguments.",
-                    agg_fun,
+                    "The function {func_name} does not accept {:?} function arguments.",
                     input_types.len()
                 );
             }
         }
         TypeSignature::VariadicAny => {
             if input_types.is_empty() {
                 return plan_err!(
-                    "The function {agg_fun:?} expects at least one argument"
+                    "The function {func_name} expects at least one argument"
                 );
             }
         }
@@ -594,7 +590,7 @@ mod tests {
         let input_types = vec![DataType::Int64, DataType::Int32];
         let signature = fun.signature();
         let result = coerce_types(&fun, &input_types, &signature);
-        assert_eq!("Error during planning: The function Min expects 1 arguments, but 2 were provided", result.unwrap_err().strip_backtrace());
+        assert_eq!("Error during planning: The function MIN expects 1 arguments, but 2 were provided", result.unwrap_err().strip_backtrace());
 
         // test input args is invalid data type for sum or avg
         let fun = AggregateFunction::Sum;

datafusion/functions-aggregate/Cargo.toml

@@ -42,3 +42,4 @@ datafusion-expr = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 log = { workspace = true }
 paste = "1.0.14"
+sqlparser = { workspace = true }