feature: Add a WindowUDFImpl::simplfy() API

Signed-off-by: guojidan <[email protected]>

Author: guojidan

datafusion-examples/examples/advanced_udwf.rs

@@ -26,7 +26,10 @@ use datafusion::error::Result;
 use datafusion::prelude::*;
 use datafusion_common::ScalarValue;
 use datafusion_expr::{
-    PartitionEvaluator, Signature, WindowFrame, WindowUDF, WindowUDFImpl,
+    expr::WindowFunction,
+    simplify::{ExprSimplifyResult, SimplifyInfo},
+    AggregateFunction, PartitionEvaluator, Signature, WindowFrame, WindowUDF,
+    WindowUDFImpl,
 };
 
 /// This example shows how to use the full WindowUDFImpl API to implement a user
@@ -75,6 +78,30 @@ impl WindowUDFImpl for SmoothItUdf {
         Ok(DataType::Float64)
     }
 
+    /// rewrite
+    fn simplify(
+        &self,
+        args: Vec<Expr>,
+        partition_by: &[Expr],
+        order_by: &[Expr],
+        window_frame: &WindowFrame,
+        null_treatment: &Option<datafusion_sql::sqlparser::ast::NullTreatment>,
+        _info: &dyn SimplifyInfo,
+    ) -> Result<ExprSimplifyResult> {
+        Ok(ExprSimplifyResult::Simplified(Expr::WindowFunction(
+            WindowFunction {
+                fun: datafusion_expr::WindowFunctionDefinition::AggregateFunction(
+                    AggregateFunction::Avg,
+                ),
+                args,
+                partition_by: partition_by.to_vec(),
+                order_by: order_by.to_vec(),
+                window_frame: window_frame.clone(),
+                null_treatment: *null_treatment,
+            },
+        )))
+    }
+
     /// Create a `PartitionEvalutor` to evaluate this function on a new
     /// partition.
     fn partition_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
@@ -144,7 +171,7 @@ async fn create_context() -> Result<SessionContext> {
 
     // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
     println!("pwd: {}", std::env::current_dir().unwrap().display());
-    let csv_path = "../../datafusion/core/tests/data/cars.csv".to_string();
+    let csv_path = "datafusion/core/tests/data/cars.csv".to_string();
     let read_options = CsvReadOptions::default().has_header(true);
 
     ctx.register_csv("cars", &csv_path, read_options).await?;

datafusion/expr/src/udwf.rs

@@ -18,11 +18,13 @@
 //! [`WindowUDF`]: User Defined Window Functions
 
 use crate::{
+    simplify::{ExprSimplifyResult, SimplifyInfo},
     Expr, PartitionEvaluator, PartitionEvaluatorFactory, ReturnTypeFunction, Signature,
     WindowFrame,
 };
 use arrow::datatypes::DataType;
 use datafusion_common::Result;
+use sqlparser::ast::NullTreatment;
 use std::{
     any::Any,
     fmt::{self, Debug, Display, Formatter},
@@ -170,6 +172,28 @@ impl WindowUDF {
         self.inner.return_type(args)
     }
 
+    /// Do the function rewrite
+    ///
+    /// See [`WindowUDFImpl::simplify`] for more details.
+    pub fn simplify(
+        &self,
+        args: Vec<Expr>,
+        partition_by: &[Expr],
+        order_by: &[Expr],
+        window_frame: &WindowFrame,
+        null_treatment: &Option<NullTreatment>,
+        info: &dyn SimplifyInfo,
+    ) -> Result<ExprSimplifyResult> {
+        self.inner.simplify(
+            args,
+            partition_by,
+            order_by,
+            window_frame,
+            null_treatment,
+            info,
+        )
+    }
+
     /// Return a `PartitionEvaluator` for evaluating this window function
     pub fn partition_evaluator_factory(&self) -> Result<Box<dyn PartitionEvaluator>> {
         self.inner.partition_evaluator()
@@ -266,6 +290,35 @@ pub trait WindowUDFImpl: Debug + Send + Sync {
     fn aliases(&self) -> &[String] {
         &[]
     }
+
+    /// Optionally apply per-UDWF simplification / rewrite rules.
+    ///
+    /// This can be used to apply function specific simplification rules during
+    /// optimization. The default implementation does nothing.
+    ///
+    /// Note that DataFusion handles simplifying arguments and  "constant
+    /// folding" (replacing a function call with constant arguments such as
+    /// `my_add(1,2) --> 3` ). Thus, there is no need to implement such
+    /// optimizations manually for specific UDFs.
+    ///
+    /// Example:
+    /// [`advanced_udwf.rs`]: https://github.com/apache/arrow-datafusion/blob/main/datafusion-examples/examples/advanced_udwf.rs
+    ///
+    /// # Returns
+    /// [`ExprSimplifyResult`] indicating the result of the simplification NOTE
+    /// if the function cannot be simplified, the arguments *MUST* be returned
+    /// unmodified
+    fn simplify(
+        &self,
+        args: Vec<Expr>,
+        _partition_by: &[Expr],
+        _order_by: &[Expr],
+        _window_frame: &WindowFrame,
+        _null_treatment: &Option<NullTreatment>,
+        _info: &dyn SimplifyInfo,
+    ) -> Result<ExprSimplifyResult> {
+        Ok(ExprSimplifyResult::Original(args))
+    }
 }
 
 /// WindowUDF that adds an alias to the underlying function. It is better to

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

@@ -40,13 +40,16 @@ use datafusion_common::{
 use datafusion_common::{
     internal_err, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue,
 };
-use datafusion_expr::expr::{InList, InSubquery};
 use datafusion_expr::simplify::ExprSimplifyResult;
 use datafusion_expr::{
     and, lit, or, BinaryExpr, BuiltinScalarFunction, Case, ColumnarValue, Expr, Like,
     Operator, ScalarFunctionDefinition, Volatility,
 };
 use datafusion_expr::{expr::ScalarFunction, interval_arithmetic::NullableInterval};
+use datafusion_expr::{
+    expr::{InList, InSubquery, WindowFunction},
+    WindowFunctionDefinition,
+};
 use datafusion_physical_expr::{create_physical_expr, execution_props::ExecutionProps};
 
 /// This structure handles API for expression simplification
@@ -1353,6 +1356,34 @@ impl<'a, S: SimplifyInfo> TreeNodeRewriter for Simplifier<'a, S> {
                 ))),
             },
 
+            Expr::WindowFunction(WindowFunction {
+                fun: WindowFunctionDefinition::WindowUDF(udwf),
+                args,
+                partition_by,
+                order_by,
+                window_frame,
+                null_treatment,
+            }) => match udwf.simplify(
+                args,
+                &partition_by,
+                &order_by,
+                &window_frame,
+                &null_treatment,
+                info,
+            )? {
+                ExprSimplifyResult::Original(args) => {
+                    Transformed::no(Expr::WindowFunction(WindowFunction {
+                        fun: WindowFunctionDefinition::WindowUDF(udwf),
+                        args,
+                        partition_by,
+                        order_by,
+                        window_frame,
+                        null_treatment,
+                    }))
+                }
+                ExprSimplifyResult::Simplified(expr) => Transformed::yes(expr),
+            },
+
             //
             // Rules for Between
             //