Minor: Support more args for udaf (#10146)

* support more args for udaf Signed-off-by: jayzhan211 <[email protected]> * fmt Signed-off-by: jayzhan211 <[email protected]> --------- Signed-off-by: jayzhan211 <[email protected]>
apache · Apr 21, 2024 · 0e21281 · 0e21281
1 parent 907475c
commit 0e21281
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 10 deletions.
diff --git a/datafusion-cli/Cargo.lock b/datafusion-cli/Cargo.lock
diff --git a/datafusion/functions-aggregate/Cargo.toml b/datafusion/functions-aggregate/Cargo.toml
@@ -45,3 +45,4 @@ datafusion-expr = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 log = { workspace = true }
 paste = "1.0.14"
+sqlparser = { workspace = true }
diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs
@@ -36,14 +36,14 @@ use datafusion_physical_expr_common::expressions;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
 use datafusion_physical_expr_common::utils::reverse_order_bys;
+use sqlparser::ast::NullTreatment;
 use std::any::Any;
 use std::fmt::Debug;
 use std::sync::Arc;
 
 make_udaf_function!(
     FirstValue,
     first_value,
-    value,
     "Returns the first value in a group of values.",
     first_value_udaf
 );

diff --git a/datafusion/functions-aggregate/src/macros.rs b/datafusion/functions-aggregate/src/macros.rs
@@ -16,19 +16,24 @@
 // under the License.
 
 macro_rules! make_udaf_function {
-    ($UDAF:ty, $EXPR_FN:ident, $($arg:ident)*, $DOC:expr, $AGGREGATE_UDF_FN:ident) => {
+    ($UDAF:ty, $EXPR_FN:ident, $DOC:expr, $AGGREGATE_UDF_FN:ident) => {
         paste::paste! {
             // "fluent expr_fn" style function
             #[doc = $DOC]
-            pub fn $EXPR_FN($($arg: Expr),*) -> Expr {
+            pub fn $EXPR_FN(
+                args: Vec<Expr>,
+                distinct: bool,
+                filter: Option<Box<Expr>>,
+                order_by: Option<Vec<Expr>>,
+                null_treatment: Option<NullTreatment>
+            ) -> Expr {
                 Expr::AggregateFunction(datafusion_expr::expr::AggregateFunction::new_udf(
                     $AGGREGATE_UDF_FN(),
-                    vec![$($arg),*],
-                    // TODO: Support arguments for `expr` API
-                    false,
-                    None,
-                    None,
-                    None,
+                    args,
+                    distinct,
+                    filter,
+                    order_by,
+                    null_treatment,
                 ))
             }
 

diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
@@ -613,7 +613,7 @@ async fn roundtrip_expr_api() -> Result<()> {
             lit(1),
         ),
         array_replace_all(make_array(vec![lit(1), lit(2), lit(3)]), lit(2), lit(4)),
-        first_value(lit(1)),
+        first_value(vec![lit(1)], false, None, None, None),
     ];
 
     // ensure expressions created with the expr api can be round tripped