Improve documentation for AggregateUdfImpl::simplify

Author: alamb

datafusion/expr/src/function.rs

@@ -27,6 +27,8 @@ pub use datafusion_functions_aggregate_common::accumulator::{
     AccumulatorArgs, AccumulatorFactoryFunction, StateFieldsArgs,
 };
 
+use crate::expr::{AggregateFunction, WindowFunction};
+use crate::simplify::SimplifyContext;
 pub use datafusion_functions_window_common::expr::ExpressionArgs;
 pub use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 pub use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
@@ -64,28 +66,22 @@ pub type PartitionEvaluatorFactory =
 pub type StateTypeFunction =
     Arc<dyn Fn(&DataType) -> Result<Arc<Vec<DataType>>> + Send + Sync>;
 
-/// [crate::udaf::AggregateUDFImpl::simplify] simplifier closure
-/// A closure with two arguments:
-/// * 'aggregate_function': [crate::expr::AggregateFunction] for which simplified has been invoked
-/// * 'info': [crate::simplify::SimplifyContext]
+/// Return type for [crate::udaf::AggregateUDFImpl::simplify]
+///
+/// This closure is invoked with two arguments:
+/// * 'aggregate_function': [AggregateFunction] with already simplified arguments
+/// * 'info': [SimplifyContext]
 ///
 /// Closure returns simplified [Expr] or an error.
-pub type AggregateFunctionSimplification = Box<
-    dyn Fn(
-        crate::expr::AggregateFunction,
-        &crate::simplify::SimplifyContext,
-    ) -> Result<Expr>,
->;
+pub type AggregateFunctionSimplification =
+    Box<dyn Fn(AggregateFunction, &SimplifyContext) -> Result<Expr>>;
 
-/// [crate::udwf::WindowUDFImpl::simplify] simplifier closure
-/// A closure with two arguments:
-/// * 'window_function': [crate::expr::WindowFunction] for which simplified has been invoked
-/// * 'info': [crate::simplify::SimplifyContext]
+/// Return type for [crate::udwf::WindowUDFImpl::simplify]
+///
+/// This closure is invoked with two arguments:
+/// * 'window_function': [WindowFunction] for which simplified has been invoked
+/// * 'info': [SimplifyContext]
 ///
 /// Closure returns simplified [Expr] or an error.
-pub type WindowFunctionSimplification = Box<
-    dyn Fn(
-        crate::expr::WindowFunction,
-        &crate::simplify::SimplifyContext,
-    ) -> Result<Expr>,
->;
+pub type WindowFunctionSimplification =
+    Box<dyn Fn(WindowFunction, &SimplifyContext) -> Result<Expr>>;

datafusion/expr/src/udaf.rs

@@ -651,10 +651,10 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync {
         AggregateOrderSensitivity::HardRequirement
     }
 
-    /// Optionally apply per-UDaF simplification / rewrite rules.
+    /// Return a closure for simplifying a user defined aggregate.
     ///
-    /// This can be used to apply function specific simplification rules during
-    /// optimization (e.g. `arrow_cast` --> `Expr::Cast`). The default
+    /// Use this to apply function specific simplification rules during
+    /// optimization (e.g. `percentile_cont()` --> `Min`). The default
     /// implementation does nothing.
     ///
     /// Note that DataFusion handles simplifying arguments and  "constant
@@ -664,10 +664,11 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync {
     ///
     /// # Returns
     ///
-    /// [None] if simplify is not defined or,
+    /// [None] if simplify is not defined
     ///
     /// Or, a closure with two arguments:
-    /// * 'aggregate_function': [AggregateFunction] for which simplified has been invoked
+    /// * 'aggregate_function': [AggregateFunction], which includes already simplified
+    ///   arguments
     /// * 'info': [crate::simplify::SimplifyContext]
     ///
     /// closure returns simplified [Expr] or an error.