Simplify display format of AggregateFunctionExpr

Author: irenjj

datafusion/core/src/physical_planner.rs

@@ -1588,6 +1588,7 @@ type AggregateExprWithOptionalArgs = (
 pub fn create_aggregate_expr_with_name_and_maybe_filter(
     e: &Expr,
     name: Option<String>,
+    sql_name: String,
     logical_input_schema: &DFSchema,
     physical_input_schema: &Schema,
     execution_props: &ExecutionProps,
@@ -1642,6 +1643,7 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
                         .order_by(ordering_reqs)
                         .schema(Arc::new(physical_input_schema.to_owned()))
                         .alias(name)
+                        .sql_name(sql_name)
                         .with_ignore_nulls(ignore_nulls)
                         .with_distinct(*distinct)
                         .build()
@@ -1664,15 +1666,22 @@ pub fn create_aggregate_expr_and_maybe_filter(
     execution_props: &ExecutionProps,
 ) -> Result<AggregateExprWithOptionalArgs> {
     // unpack (nested) aliased logical expressions, e.g. "sum(col) as total"
-    let (name, e) = match e {
-        Expr::Alias(Alias { expr, name, .. }) => (Some(name.clone()), expr.as_ref()),
-        Expr::AggregateFunction(_) => (Some(e.schema_name().to_string()), e),
-        _ => (None, e),
+    let (name, sql_name, e) = match e {
+        Expr::Alias(Alias { expr, name, .. }) => {
+            (Some(name.clone()), String::default(), expr.as_ref())
+        }
+        Expr::AggregateFunction(_) => (
+            Some(e.schema_name().to_string()),
+            e.sql_name().to_string(),
+            e,
+        ),
+        _ => (None, String::default(), e),
     };
 
     create_aggregate_expr_with_name_and_maybe_filter(
         e,
         name,
+        sql_name,
         logical_input_schema,
         physical_input_schema,
         execution_props,

datafusion/expr/src/expr.rs

@@ -1147,6 +1147,10 @@ impl Expr {
         SchemaDisplay(self)
     }
 
+    pub fn sql_name(&self) -> impl Display + '_ {
+        SqlDisplay(self)
+    }
+
     /// Returns the qualifier and the schema name of this expression.
     ///
     /// Used when the expression forms the output field of a certain plan.
@@ -2596,6 +2600,43 @@ impl Display for SchemaDisplay<'_> {
     }
 }
 
+struct SqlDisplay<'a>(&'a Expr);
+impl Display for SqlDisplay<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self.0 {
+            Expr::Column(c) => {
+                write!(f, "{}", c.name)
+            }
+            Expr::Literal(_) => {
+                write!(f, "aa")
+            }
+            Expr::ScalarVariable(..) => {
+                write!(f, "bb")
+            }
+            Expr::OuterReferenceColumn(..) => {
+                write!(f, "cc")
+            }
+            Expr::Placeholder(_) => {
+                write!(f, "dd")
+            }
+            Expr::Wildcard { .. } => {
+                write!(f, "ee")
+            }
+            Expr::AggregateFunction(AggregateFunction { func, params }) => {
+                match func.sql_name(params) {
+                    Ok(name) => {
+                        write!(f, "{name}")
+                    }
+                    Err(e) => {
+                        write!(f, "got error from schema_name {}", e)
+                    }
+                }
+            }
+            _ => write!(f, "{}", self.0.schema_name()),
+        }
+    }
+}
+
 /// Get schema_name for Vector of expressions
 ///
 /// Internal usage. Please call `schema_name_from_exprs` instead
@@ -2607,11 +2648,21 @@ pub(crate) fn schema_name_from_exprs_comma_separated_without_space(
     schema_name_from_exprs_inner(exprs, ",")
 }
 
+pub(crate) fn sql_name_from_exprs_comma_separated_without_space(
+    exprs: &[Expr],
+) -> Result<String, fmt::Error> {
+    sql_name_from_exprs_inner(exprs, ",")
+}
+
 /// Get schema_name for Vector of expressions
 pub fn schema_name_from_exprs(exprs: &[Expr]) -> Result<String, fmt::Error> {
     schema_name_from_exprs_inner(exprs, ", ")
 }
 
+pub fn sql_name_from_exprs(exprs: &[Expr]) -> Result<String, fmt::Error> {
+    sql_name_from_exprs_inner(exprs, ", ")
+}
+
 fn schema_name_from_exprs_inner(exprs: &[Expr], sep: &str) -> Result<String, fmt::Error> {
     let mut s = String::new();
     for (i, e) in exprs.iter().enumerate() {
@@ -2624,6 +2675,18 @@ fn schema_name_from_exprs_inner(exprs: &[Expr], sep: &str) -> Result<String, fmt
     Ok(s)
 }
 
+fn sql_name_from_exprs_inner(exprs: &[Expr], sep: &str) -> Result<String, fmt::Error> {
+    let mut s = String::new();
+    for (i, e) in exprs.iter().enumerate() {
+        if i > 0 {
+            write!(&mut s, "{sep}")?;
+        }
+        write!(&mut s, "{}", SqlDisplay(e))?;
+    }
+
+    Ok(s)
+}
+
 pub fn schema_name_from_sorts(sorts: &[Sort]) -> Result<String, fmt::Error> {
     let mut s = String::new();
     for (i, e) in sorts.iter().enumerate() {

datafusion/expr/src/udaf.rs

@@ -31,8 +31,8 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 
 use crate::expr::{
     schema_name_from_exprs, schema_name_from_exprs_comma_separated_without_space,
-    schema_name_from_sorts, AggregateFunction, AggregateFunctionParams,
-    WindowFunctionParams,
+    schema_name_from_sorts, sql_name_from_exprs_comma_separated_without_space,
+    AggregateFunction, AggregateFunctionParams, WindowFunctionParams,
 };
 use crate::function::{
     AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs,
@@ -175,6 +175,10 @@ impl AggregateUDF {
         self.inner.schema_name(params)
     }
 
+    pub fn sql_name(&self, params: &AggregateFunctionParams) -> Result<String> {
+        self.inner.sql_name(params)
+    }
+
     pub fn window_function_schema_name(
         &self,
         params: &WindowFunctionParams,
@@ -452,6 +456,42 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
         Ok(schema_name)
     }
 
+    fn sql_name(&self, params: &AggregateFunctionParams) -> Result<String> {
+        let AggregateFunctionParams {
+            args,
+            distinct,
+            filter,
+            order_by,
+            null_treatment,
+        } = params;
+
+        let mut schema_name = String::new();
+
+        schema_name.write_fmt(format_args!(
+            "{}({}{})",
+            self.name(),
+            if *distinct { "DISTINCT " } else { "" },
+            sql_name_from_exprs_comma_separated_without_space(args)?
+        ))?;
+
+        if let Some(null_treatment) = null_treatment {
+            schema_name.write_fmt(format_args!(" {}", null_treatment))?;
+        }
+
+        if let Some(filter) = filter {
+            schema_name.write_fmt(format_args!(" FILTER (WHERE {filter})"))?;
+        };
+
+        if let Some(order_by) = order_by {
+            schema_name.write_fmt(format_args!(
+                " ORDER BY [{}]",
+                schema_name_from_sorts(order_by)?
+            ))?;
+        };
+
+        Ok(schema_name)
+    }
+
     /// Returns the name of the column this expression would create
     ///
     /// See [`Expr::schema_name`] for details

datafusion/physical-expr/src/aggregate.rs

@@ -65,6 +65,7 @@ pub struct AggregateExprBuilder {
     /// Physical expressions of the aggregate function
     args: Vec<Arc<dyn PhysicalExpr>>,
     alias: Option<String>,
+    sql_name: String,
     /// Arrow Schema for the aggregate function
     schema: SchemaRef,
     /// The physical order by expressions
@@ -83,6 +84,7 @@ impl AggregateExprBuilder {
             fun,
             args,
             alias: None,
+            sql_name: String::default(),
             schema: Arc::new(Schema::empty()),
             ordering_req: LexOrdering::default(),
             ignore_nulls: false,
@@ -99,6 +101,7 @@ impl AggregateExprBuilder {
             fun,
             args,
             alias,
+            sql_name,
             schema,
             ordering_req,
             ignore_nulls,
@@ -148,6 +151,7 @@ impl AggregateExprBuilder {
             args,
             data_type,
             name,
+            sql_name,
             schema: Arc::unwrap_or_clone(schema),
             ordering_req,
             ignore_nulls,
@@ -164,6 +168,11 @@ impl AggregateExprBuilder {
         self
     }
 
+    pub fn sql_name(mut self, name: String) -> Self {
+        self.sql_name = name;
+        self
+    }
+
     pub fn schema(mut self, schema: SchemaRef) -> Self {
         self.schema = schema;
         self
@@ -215,6 +224,7 @@ pub struct AggregateFunctionExpr {
     /// Output / return type of this aggregate
     data_type: DataType,
     name: String,
+    sql_name: String,
     schema: Schema,
     // The physical order by expressions
     ordering_req: LexOrdering,
@@ -245,6 +255,11 @@ impl AggregateFunctionExpr {
         &self.name
     }
 
+    /// Simplified name for `tree` explain.
+    pub fn sql_name(&self) -> &str {
+        &self.sql_name
+    }
+
     /// Return if the aggregation is distinct
     pub fn is_distinct(&self) -> bool {
         self.is_distinct