Consolidate TreeNode transform and rewrite APIs

datafusion-examples/examples/rewrite_expr.rs

@@ -91,33 +91,35 @@ impl AnalyzerRule for MyAnalyzerRule {
 
 impl MyAnalyzerRule {
     fn analyze_plan(plan: LogicalPlan) -> Result<LogicalPlan> {
-        plan.transform(&|plan| {
+        plan.transform_up(&|plan| {
             Ok(match plan {
                 LogicalPlan::Filter(filter) => {
                     let predicate = Self::analyze_expr(filter.predicate.clone())?;
-                    Transformed::Yes(LogicalPlan::Filter(Filter::try_new(
+                    Transformed::yes(LogicalPlan::Filter(Filter::try_new(
                         predicate,
                         filter.input,
                     )?))
                 }
-                _ => Transformed::No(plan),
+                _ => Transformed::no(plan),
             })
         })
+        .map(|t| t.data)
     }
 
     fn analyze_expr(expr: Expr) -> Result<Expr> {
-        expr.transform(&|expr| {
+        expr.transform_up(&|expr| {
             // closure is invoked for all sub expressions
             Ok(match expr {
                 Expr::Literal(ScalarValue::Int64(i)) => {
                     // transform to UInt64
-                    Transformed::Yes(Expr::Literal(ScalarValue::UInt64(
+                    Transformed::yes(Expr::Literal(ScalarValue::UInt64(
                         i.map(|i| i as u64),
                     )))
                 }
-                _ => Transformed::No(expr),
+                _ => Transformed::no(expr),
             })
         })
+        .map(|t| t.data)
     }
 }
 
@@ -161,7 +163,7 @@ impl OptimizerRule for MyOptimizerRule {
 
 /// use rewrite_expr to modify the expression tree.
 fn my_rewrite(expr: Expr) -> Result<Expr> {
-    expr.transform(&|expr| {
+    expr.transform_up(&|expr| {
         // closure is invoked for all sub expressions
         Ok(match expr {
             Expr::Between(Between {
@@ -175,14 +177,15 @@ fn my_rewrite(expr: Expr) -> Result<Expr> {
                 let low: Expr = *low;
                 let high: Expr = *high;
                 if negated {
-                    Transformed::Yes(expr.clone().lt(low).or(expr.gt(high)))
+                    Transformed::yes(expr.clone().lt(low).or(expr.gt(high)))
                 } else {
-                    Transformed::Yes(expr.clone().gt_eq(low).and(expr.lt_eq(high)))
+                    Transformed::yes(expr.clone().gt_eq(low).and(expr.lt_eq(high)))
                 }
             }
-            _ => Transformed::No(expr),
+            _ => Transformed::no(expr),
         })
     })
+    .map(|t| t.data)
 }
 
 #[derive(Default)]