Feature/alternate function extension

datafusion-examples/examples/expr_api.rs

@@ -33,7 +33,7 @@ use datafusion_expr::execution_props::ExecutionProps;
 use datafusion_expr::expr::BinaryExpr;
 use datafusion_expr::interval_arithmetic::Interval;
 use datafusion_expr::simplify::SimplifyContext;
-use datafusion_expr::{AggregateExt, ColumnarValue, ExprSchemable, Operator};
+use datafusion_expr::{ColumnarValue, ExprSchemable, Operator};
 
 /// This example demonstrates the DataFusion [`Expr`] API.
 ///
@@ -95,13 +95,12 @@ fn expr_fn_demo() -> Result<()> {
     let agg = first_value.call(vec![col("price")]);
     assert_eq!(agg.to_string(), "first_value(price)");
 
-    // You can use the AggregateExt trait to create more complex aggregates
+    // You can use the ExprFunctionExt trait to create more complex aggregates
     // such as `FIRST_VALUE(price FILTER quantity > 100 ORDER BY ts )
     let agg = first_value
         .call(vec![col("price")])
-        .order_by(vec![col("ts").sort(false, false)])
-        .filter(col("quantity").gt(lit(100)))
-        .build()?; // build the aggregate
+        .order_by(vec![col("ts").sort(false, false)])?
+        .filter(col("quantity").gt(lit(100)))?;
     assert_eq!(
         agg.to_string(),
         "first_value(price) FILTER (WHERE quantity > Int32(100)) ORDER BY [ts DESC NULLS LAST]"

datafusion/core/src/dataframe/mod.rs

@@ -1697,7 +1697,7 @@ mod tests {
     use datafusion_common_runtime::SpawnedTask;
     use datafusion_expr::{
         array_agg, cast, create_udf, expr, lit, BuiltInWindowFunction,
-        ScalarFunctionImplementation, Volatility, WindowFrame, WindowFunctionDefinition,
+        ScalarFunctionImplementation, Volatility, WindowFunctionDefinition,
     };
     use datafusion_functions_aggregate::expr_fn::count_distinct;
     use datafusion_physical_expr::expressions::Column;
@@ -1867,11 +1867,8 @@ mod tests {
                 BuiltInWindowFunction::FirstValue,
             ),
             vec![col("aggregate_test_100.c1")],
-            vec![col("aggregate_test_100.c2")],
-            vec![],
-            WindowFrame::new(None),
-            None,
-        ));
+        ))
+        .partition_by(vec![col("aggregate_test_100.c2")])?;
         let t2 = t.select(vec![col("c1"), first_row])?;
         let plan = t2.plan.clone();
 

datafusion/core/src/physical_planner.rs

@@ -101,7 +101,7 @@ fn create_function_physical_name(
     fun: &str,
     distinct: bool,
     args: &[Expr],
-    order_by: Option<&Vec<Expr>>,
+    order_by: &Option<Vec<Expr>>,
 ) -> Result<String> {
     let names: Vec<String> = args
         .iter()
@@ -115,7 +115,7 @@ fn create_function_physical_name(
 
     let phys_name = format!("{}({}{})", fun, distinct_str, names.join(","));
 
-    Ok(order_by
+    Ok(order_by.as_ref()
         .map(|order_by| format!("{} ORDER BY [{}]", phys_name, expr_vec_fmt!(order_by)))
         .unwrap_or(phys_name))
 }
@@ -220,7 +220,7 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result<String> {
             order_by,
             ..
         }) => {
-            create_function_physical_name(&fun.to_string(), false, args, Some(order_by))
+            create_function_physical_name(&fun.to_string(), false, args, order_by)
         }
         Expr::AggregateFunction(AggregateFunction {
             func_def,
@@ -233,7 +233,7 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result<String> {
             func_def.name(),
             *distinct,
             args,
-            order_by.as_ref(),
+            order_by,
         ),
         Expr::GroupingSet(grouping_set) => match grouping_set {
             GroupingSet::Rollup(exprs) => Ok(format!(
@@ -1744,22 +1744,25 @@ pub fn create_window_expr_with_name(
     let name = name.into();
     let physical_schema: &Schema = &logical_schema.into();
     match e {
-        Expr::WindowFunction(WindowFunction {
-            fun,
-            args,
-            partition_by,
-            order_by,
-            window_frame,
-            null_treatment,
-        }) => {
+        Expr::WindowFunction(window_fun) => {
+            let window_frame = window_fun.get_frame_or_default();
+            let WindowFunction {
+                fun,
+                args,
+                partition_by,
+                order_by,
+                null_treatment,
+                ..
+            } = window_fun;
+
             let physical_args =
                 create_physical_exprs(args, logical_schema, execution_props)?;
             let partition_by =
                 create_physical_exprs(partition_by, logical_schema, execution_props)?;
             let order_by =
-                create_physical_sort_exprs(order_by, logical_schema, execution_props)?;
+                create_physical_sort_exprs(order_by.as_ref().unwrap_or(&vec![]), logical_schema, execution_props)?;
 
-            if !is_window_frame_bound_valid(window_frame) {
+            if !is_window_frame_bound_valid(&window_frame) {
                 return plan_err!(
                         "Invalid window frame: start bound ({}) cannot be larger than end bound ({})",
                         window_frame.start_bound, window_frame.end_bound

datafusion/core/tests/dataframe/mod.rs

@@ -55,8 +55,8 @@ use datafusion_expr::expr::{GroupingSet, Sort};
 use datafusion_expr::var_provider::{VarProvider, VarType};
 use datafusion_expr::{
     array_agg, cast, col, exists, expr, in_subquery, lit, max, out_ref_col, placeholder,
-    scalar_subquery, when, wildcard, Expr, ExprSchemable, WindowFrame, WindowFrameBound,
-    WindowFrameUnits, WindowFunctionDefinition,
+    scalar_subquery, when, wildcard, Expr, ExprSchemable, WindowFrame,
+    WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
 };
 use datafusion_functions_aggregate::expr_fn::{avg, count, sum};
 
@@ -183,15 +183,13 @@ async fn test_count_wildcard_on_window() -> Result<()> {
         .select(vec![Expr::WindowFunction(expr::WindowFunction::new(
             WindowFunctionDefinition::AggregateUDF(count_udaf()),
             vec![wildcard()],
-            vec![],
-            vec![Expr::Sort(Sort::new(Box::new(col("a")), false, true))],
-            WindowFrame::new_bounds(
-                WindowFrameUnits::Range,
-                WindowFrameBound::Preceding(ScalarValue::UInt32(Some(6))),
-                WindowFrameBound::Following(ScalarValue::UInt32(Some(2))),
-            ),
-            None,
-        ))])?
+        ))
+        .order_by(vec![Expr::Sort(Sort::new(Box::new(col("a")), false, true))])?
+        .window_frame(WindowFrame::new_bounds(
+            WindowFrameUnits::Range,
+            WindowFrameBound::Preceding(ScalarValue::UInt32(Some(6))),
+            WindowFrameBound::Following(ScalarValue::UInt32(Some(2))),
+        ))?])?
         .explain(false, false)?
         .collect()
         .await?;

datafusion/core/tests/expr_api/mod.rs

@@ -20,8 +20,7 @@ use arrow_array::builder::{ListBuilder, StringBuilder};
 use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringArray, StructArray};
 use arrow_schema::{DataType, Field};
 use datafusion::prelude::*;
-use datafusion_common::{assert_contains, DFSchema, ScalarValue};
-use datafusion_expr::AggregateExt;
+use datafusion_common::{assert_contains, DFSchema, Result, ScalarValue};
 use datafusion_functions::core::expr_ext::FieldAccessor;
 use datafusion_functions_aggregate::first_last::first_value_udaf;
 use datafusion_functions_aggregate::sum::sum_udaf;
@@ -173,7 +172,6 @@ async fn test_aggregate_error() {
         .call(vec![col("props")])
         // not a sort column
         .order_by(vec![col("id")])
-        .build()
         .unwrap_err()
         .to_string();
     assert_contains!(
@@ -183,22 +181,18 @@ async fn test_aggregate_error() {
 }
 
 #[tokio::test]
-async fn test_aggregate_ext_order_by() {
+async fn test_aggregate_ext_order_by() -> Result<()> {
     let agg = first_value_udaf().call(vec![col("props")]);
 
     // ORDER BY id ASC
     let agg_asc = agg
         .clone()
-        .order_by(vec![col("id").sort(true, true)])
-        .build()
-        .unwrap()
+        .order_by(vec![col("id").sort(true, true)])?
         .alias("asc");
 
     // ORDER BY id DESC
     let agg_desc = agg
-        .order_by(vec![col("id").sort(false, true)])
-        .build()
-        .unwrap()
+        .order_by(vec![col("id").sort(false, true)])?
         .alias("desc");
 
     evaluate_agg_test(
@@ -224,16 +218,15 @@ async fn test_aggregate_ext_order_by() {
         ],
     )
     .await;
+    Ok(())
 }
 
 #[tokio::test]
-async fn test_aggregate_ext_filter() {
+async fn test_aggregate_ext_filter() -> Result<()> {
     let agg = first_value_udaf()
         .call(vec![col("i")])
-        .order_by(vec![col("i").sort(true, true)])
-        .filter(col("i").is_not_null())
-        .build()
-        .unwrap()
+        .order_by(vec![col("i").sort(true, true)])?
+        .filter(col("i").is_not_null())?
         .alias("val");
 
     #[rustfmt::skip]
@@ -248,16 +241,15 @@ async fn test_aggregate_ext_filter() {
         ],
     )
         .await;
+    Ok(())
 }
 
 #[tokio::test]
-async fn test_aggregate_ext_distinct() {
+async fn test_aggregate_ext_distinct() -> Result<()> {
     let agg = sum_udaf()
         .call(vec![lit(5)])
         // distinct sum should be 5, not 15
-        .distinct()
-        .build()
-        .unwrap()
+        .distinct()?
         .alias("distinct");
 
     evaluate_agg_test(
@@ -271,25 +263,22 @@ async fn test_aggregate_ext_distinct() {
         ],
     )
     .await;
+    Ok(())
 }
 
 #[tokio::test]
-async fn test_aggregate_ext_null_treatment() {
+async fn test_aggregate_ext_null_treatment() -> Result<()> {
     let agg = first_value_udaf()
         .call(vec![col("i")])
-        .order_by(vec![col("i").sort(true, true)]);
+        .order_by(vec![col("i").sort(true, true)])?;
 
     let agg_respect = agg
         .clone()
-        .null_treatment(NullTreatment::RespectNulls)
-        .build()
-        .unwrap()
+        .null_treatment(NullTreatment::RespectNulls)?
         .alias("respect");
 
     let agg_ignore = agg
-        .null_treatment(NullTreatment::IgnoreNulls)
-        .build()
-        .unwrap()
+        .null_treatment(NullTreatment::IgnoreNulls)?
         .alias("ignore");
 
     evaluate_agg_test(
@@ -315,6 +304,7 @@ async fn test_aggregate_ext_null_treatment() {
         ],
     )
     .await;
+    Ok(())
 }
 
 /// Evaluates the specified expr as an aggregate and compares the result to the