Consolidate Expr::unalias_nested

datafusion/expr/src/expr.rs

@@ -33,7 +33,9 @@ use crate::{
 use crate::{window_frame, Volatility};
 
 use arrow::datatypes::{DataType, FieldRef};
-use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::tree_node::{
+    Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
+};
 use datafusion_common::{
     internal_err, plan_err, Column, DFSchema, Result, ScalarValue, TableReference,
 };
@@ -1193,9 +1195,8 @@ impl Expr {
 
     /// Recursively potentially multiple aliases from an expression.
     ///
-    /// If the expression is not an alias, the expression is returned unchanged.
-    /// This method removes directly nested aliases, but not other nested
-    /// aliases.
+    /// This method removes nested aliases and returns `Transformed`
+    /// to signal if the expression was changed.
     ///
     /// # Example
     /// ```
@@ -1212,11 +1213,22 @@ impl Expr {
     /// let expr = col("foo").alias("bar").alias("baz");
     /// assert_eq!(expr.unalias_nested(), col("foo"));
     /// ```
-    pub fn unalias_nested(self) -> Expr {
-        match self {
-            Expr::Alias(alias) => alias.expr.unalias_nested(),
-            _ => self,
-        }
+    pub fn unalias_nested(self) -> Transformed<Expr> {
+        self.transform_down(|expr| {
+            match expr {
+                Expr::Exists { .. } | Expr::ScalarSubquery(_) | Expr::InSubquery(_) => {
+                    // subqueries could contain aliases so we don't recurse into those
+                    Ok(Transformed::new(expr, false, TreeNodeRecursion::Jump))
+                }
+                Expr::Alias(Alias { expr, .. }) => {
+                    // directly recurse into the alias expression
+                    // as we just modified the tree structure
+                    Ok(Transformed::yes(expr.unalias_nested().data))
+                }
+                _ => Ok(Transformed::no(expr)),
+            }
+        })
+        .expect("unalias_nested doesn't return err")
     }
 
     /// Return `self IN <list>` if `negated` is false, otherwise

datafusion/expr/src/logical_plan/plan.rs

@@ -869,8 +869,7 @@ impl LogicalPlan {
             }
             LogicalPlan::Filter { .. } => {
                 assert_eq!(1, expr.len());
-                let predicate = expr.pop().unwrap();
-                let predicate = Filter::remove_aliases(predicate)?.data;
+                let predicate = expr.pop().unwrap().unalias_nested().data;
 
                 Filter::try_new(predicate, Arc::new(inputs.swap_remove(0)))
                     .map(LogicalPlan::Filter)
@@ -2200,38 +2199,6 @@ impl Filter {
         }
         false
     }
-
-    /// Remove aliases from a predicate for use in a `Filter`
-    ///
-    /// filter predicates should not contain aliased expressions so we remove
-    /// any aliases.
-    ///
-    /// before this logic was added we would have aliases within filters such as
-    /// for benchmark q6:
-    ///
-    /// ```sql
-    /// lineitem.l_shipdate >= Date32(\"8766\")
-    /// AND lineitem.l_shipdate < Date32(\"9131\")
-    /// AND CAST(lineitem.l_discount AS Decimal128(30, 15)) AS lineitem.l_discount >=
-    /// Decimal128(Some(49999999999999),30,15)
-    /// AND CAST(lineitem.l_discount AS Decimal128(30, 15)) AS lineitem.l_discount <=
-    /// Decimal128(Some(69999999999999),30,15)
-    /// AND lineitem.l_quantity < Decimal128(Some(2400),15,2)
-    /// ```
-    pub fn remove_aliases(predicate: Expr) -> Result<Transformed<Expr>> {
-        predicate.transform_down(|expr| {
-            match expr {
-                Expr::Exists { .. } | Expr::ScalarSubquery(_) | Expr::InSubquery(_) => {
-                    // subqueries could contain aliases so we don't recurse into those
-                    Ok(Transformed::new(expr, false, TreeNodeRecursion::Jump))
-                }
-                Expr::Alias(Alias { expr, .. }) => {
-                    Ok(Transformed::new(*expr, true, TreeNodeRecursion::Jump))
-                }
-                _ => Ok(Transformed::no(expr)),
-            }
-        })
-    }
 }
 
 /// Window its input based on a set of window spec and window function (e.g. SUM or RANK)

datafusion/optimizer/src/common_subexpr_eliminate.rs

@@ -259,9 +259,12 @@ impl CommonSubexprEliminate {
         self.try_unary_plan(expr, input, config)?
             .transform_data(|(mut new_expr, new_input)| {
                 assert_eq!(new_expr.len(), 1); // passed in vec![predicate]
-                let new_predicate = new_expr.pop().unwrap();
-                Ok(Filter::remove_aliases(new_predicate)?
-                    .update_data(|new_predicate| (new_predicate, new_input)))
+                let new_predicate = new_expr
+                    .pop()
+                    .unwrap()
+                    .unalias_nested()
+                    .update_data(|new_predicate| (new_predicate, new_input));
+                Ok(new_predicate)
             })?
             .map_data(|(new_predicate, new_input)| {
                 Filter::try_new(new_predicate, Arc::new(new_input))

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -608,7 +608,7 @@ fn rewrite_expr(expr: Expr, input: &Projection) -> Result<Transformed<Expr>> {
                 // * the current column is an expression "f"
                 //
                 // return the expression `d + e` (not `d + e` as f)
-                let input_expr = input.expr[idx].clone().unalias_nested();
+                let input_expr = input.expr[idx].clone().unalias_nested().data;
                 Ok(Transformed::yes(input_expr))
             }
             // Unsupported type for consecutive projection merge analysis.