WIP / Don't recurse in apply_subqueries

datafusion/expr/src/logical_plan/plan.rs

@@ -634,7 +634,7 @@ impl LogicalPlan {
     pub fn using_columns(&self) -> Result<Vec<HashSet<Column>>, DataFusionError> {
         let mut using_columns: Vec<HashSet<Column>> = vec![];
 
-        self.apply_with_subqueries(&mut |plan| {
+        self.apply(&mut |plan| {
             if let LogicalPlan::Join(Join {
                 join_constraint: JoinConstraint::Using,
                 on,
@@ -1347,17 +1347,6 @@ impl LogicalPlan {
         )
     }
 
-    pub fn apply_with_subqueries<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
-        &self,
-        f: &mut F,
-    ) -> Result<TreeNodeRecursion> {
-        self.apply(&mut |n| {
-            f(n)?
-                .visit_children(|| self.apply_subqueries(|c| c.apply_with_subqueries(f)))?
-                .visit_sibling(|| self.apply_children(|c| c.apply_with_subqueries(f)))
-        })
-    }
-
     pub fn transform_with_subqueries<F: Fn(Self) -> Result<Transformed<Self>>>(
         self,
         f: &F,
@@ -1499,7 +1488,7 @@ impl LogicalPlan {
     ) -> Result<HashMap<String, Option<DataType>>, DataFusionError> {
         let mut param_types: HashMap<String, Option<DataType>> = HashMap::new();
 
-        self.apply_with_subqueries(&mut |plan| {
+        self.apply(&mut |plan| {
             plan.apply_expressions(|expr| {
                 expr.apply(&mut |expr| {
                     if let Expr::Placeholder(Placeholder { id, data_type }) = expr {

datafusion/optimizer/src/analyzer/mod.rs

@@ -21,12 +21,11 @@ use log::debug;
 
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::instant::Instant;
-use datafusion_common::tree_node::TreeNodeRecursion;
+use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion_common::{DataFusionError, Result};
 use datafusion_expr::expr::Exists;
 use datafusion_expr::expr::InSubquery;
 use datafusion_expr::expr_rewriter::FunctionRewrite;
-use datafusion_expr::utils::inspect_expr_pre;
 use datafusion_expr::{Expr, LogicalPlan};
 
 use crate::analyzer::count_wildcard_rule::CountWildcardRule;
@@ -155,19 +154,22 @@ impl Analyzer {
 
 /// Do necessary check and fail the invalid plan
 fn check_plan(plan: &LogicalPlan) -> Result<()> {
-    plan.apply_with_subqueries(&mut |plan: &LogicalPlan| {
-        for expr in plan.expressions().iter() {
+    plan.apply(&mut |plan: &LogicalPlan| {
+        plan.apply_expressions(|expr| {
             // recursively look for subqueries
-            inspect_expr_pre(expr, |expr| match expr {
-                Expr::Exists(Exists { subquery, .. })
-                | Expr::InSubquery(InSubquery { subquery, .. })
-                | Expr::ScalarSubquery(subquery) => {
-                    check_subquery_expr(plan, &subquery.subquery, expr)
-                }
-                _ => Ok(()),
+            expr.apply(&mut |expr| {
+                match expr {
+                    Expr::Exists(Exists { subquery, .. })
+                    | Expr::InSubquery(InSubquery { subquery, .. })
+                    | Expr::ScalarSubquery(subquery) => {
+                        check_subquery_expr(plan, &subquery.subquery, expr)?;
+                    }
+                    _ => {}
+                };
+                Ok(TreeNodeRecursion::Continue)
             })?;
-        }
-
+            Ok(TreeNodeRecursion::Continue)
+        })?;
         Ok(TreeNodeRecursion::Continue)
     })?;
 

datafusion/optimizer/src/analyzer/subquery.rs

@@ -283,7 +283,7 @@ fn strip_inner_query(inner_plan: &LogicalPlan) -> &LogicalPlan {
 
 fn get_correlated_expressions(inner_plan: &LogicalPlan) -> Result<Vec<Expr>> {
     let mut exprs = vec![];
-    inner_plan.apply_with_subqueries(&mut |plan| {
+    inner_plan.apply(&mut |plan| {
         if let LogicalPlan::Filter(Filter { predicate, .. }) = plan {
             let (correlated, _): (Vec<_>, Vec<_>) = split_conjunction(predicate)
                 .into_iter()

datafusion/optimizer/src/plan_signature.rs

@@ -21,7 +21,7 @@ use std::{
     num::NonZeroUsize,
 };
 
-use datafusion_common::tree_node::TreeNodeRecursion;
+use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion_expr::LogicalPlan;
 
 /// Non-unique identifier of a [`LogicalPlan`].
@@ -73,7 +73,7 @@ impl LogicalPlanSignature {
 /// Get total number of [`LogicalPlan`]s in the plan.
 fn get_node_number(plan: &LogicalPlan) -> NonZeroUsize {
     let mut node_number = 0;
-    plan.apply_with_subqueries(&mut |_plan| {
+    plan.apply(&mut |_plan| {
         node_number += 1;
         Ok(TreeNodeRecursion::Continue)
     })