Handle subqueries

Author: alamb

datafusion/common/src/tree_node.rs

@@ -649,8 +649,12 @@ impl<T> Transformed<T> {
 }
 
 impl Transformed<()> {
-    /// Invoke the given function `f`  and combine the transformed state with the
-    /// current state
+    /// Invoke the given function `f`  and combine the transformed state with
+    /// the current state,
+    ///
+    /// if f() returns an Err, returns that err
+    /// If f() returns Ok, returns a true transformed flag if either self or
+    /// the result of f() was transformed
     pub fn and_then<F>(self, f: F) -> Result<Transformed<()>>
     where
         F: FnOnce() -> Result<Transformed<()>>,

datafusion/expr/src/logical_plan/mutate.rs

@@ -16,7 +16,8 @@
 // under the License.
 
 use super::plan::*;
-use crate::{lit, Expr};
+use crate::expr::{Exists, InSubquery};
+use crate::Expr;
 use datafusion_common::tree_node::Transformed;
 use datafusion_common::Result;
 use datafusion_common::{DFSchema, DFSchemaRef};
@@ -27,25 +28,13 @@ use std::sync::Arc;
 fn rewrite_expr_iter_mut<'a, F>(
     i: impl IntoIterator<Item = &'a mut Expr>,
     mut f: F,
-) -> Result<()>
+) -> Result<Transformed<()>>
 where
-    F: FnMut(Expr) -> Result<Expr>,
+    F: FnMut(&mut Expr) -> Result<Transformed<()>>,
 {
-    i.into_iter().try_for_each(|e| rewrite_expr(e, &mut f))
-}
-
-/// Rewrites an expression in place using a function that takes the expression by ownership
-fn rewrite_expr<'a, F>(e: &'a mut Expr, mut f: F) -> Result<()>
-where
-    F: FnMut(Expr) -> Result<Expr>,
-{
-    let mut t = lit(0);
-    std::mem::swap(e, &mut t);
-    // transform
-    let mut t = f(t)?;
-    // put it back
-    std::mem::swap(e, &mut t);
-    Ok(())
+    i.into_iter().fold(Ok(Transformed::no(())), |acc, expr| {
+        acc.and_then(|acc| acc.and_then(|| f(expr)))
+    })
 }
 
 impl LogicalPlan {
@@ -54,61 +43,58 @@ impl LogicalPlan {
     ///
     /// If the closure returns an error, the error is returned and the expressions
     /// are left in a partially modified state
-    pub fn rewrite_exprs<F>(mut self, mut f: F) -> Result<Self>
+    pub fn rewrite_exprs<F>(&mut self, mut f: F) -> Result<Transformed<()>>
     where
-        F: FnMut(Expr) -> Result<Expr>,
+        F: FnMut(&mut Expr) -> Result<Transformed<()>>,
     {
-        match &mut self {
+        match self {
             LogicalPlan::Projection(Projection { expr, .. }) => {
-                rewrite_expr_iter_mut(expr.iter_mut(), &mut f)?;
+                rewrite_expr_iter_mut(expr.iter_mut(), f)
             }
             LogicalPlan::Values(Values { values, .. }) => {
-                rewrite_expr_iter_mut(values.iter_mut().flatten(), &mut f)?;
-            }
-            LogicalPlan::Filter(Filter { predicate, .. }) => {
-                rewrite_expr(predicate, &mut f)?
+                rewrite_expr_iter_mut(values.iter_mut().flatten(), f)
             }
+            LogicalPlan::Filter(Filter { predicate, .. }) => f(predicate),
             LogicalPlan::Repartition(Repartition {
                 partitioning_scheme,
                 ..
             }) => match partitioning_scheme {
-                Partitioning::Hash(expr, _) => {
-                    rewrite_expr_iter_mut(expr.iter_mut(), &mut f)?
-                }
+                Partitioning::Hash(expr, _) => rewrite_expr_iter_mut(expr.iter_mut(), f),
                 Partitioning::DistributeBy(expr) => {
-                    rewrite_expr_iter_mut(expr.iter_mut(), &mut f)?
+                    rewrite_expr_iter_mut(expr.iter_mut(), f)
                 }
-                Partitioning::RoundRobinBatch(_) => {}
+                Partitioning::RoundRobinBatch(_) => Ok(Transformed::no(())),
             },
             LogicalPlan::Window(Window { window_expr, .. }) => {
-                rewrite_expr_iter_mut(window_expr.iter_mut(), &mut f)?;
+                rewrite_expr_iter_mut(window_expr.iter_mut(), f)
             }
             LogicalPlan::Aggregate(Aggregate {
                 group_expr,
                 aggr_expr,
                 ..
-            }) => rewrite_expr_iter_mut(
-                group_expr.iter_mut().chain(aggr_expr.iter_mut()),
-                &mut f,
-            )?,
+            }) => {
+                let exprs = group_expr.iter_mut().chain(aggr_expr.iter_mut());
+                rewrite_expr_iter_mut(exprs, f)
+            }
             // There are two part of expression for join, equijoin(on) and non-equijoin(filter).
             // 1. the first part is `on.len()` equijoin expressions, and the struct of each expr is `left-on = right-on`.
             // 2. the second part is non-equijoin(filter).
             LogicalPlan::Join(Join { on, filter, .. }) => {
                 // don't look at the equijoin expressions as a whole
-                rewrite_expr_iter_mut(
-                    on.iter_mut().flat_map(|(e1, e2)| {
-                        std::iter::once(e1).chain(std::iter::once(e2))
-                    }),
-                    &mut f,
-                )?;
+                let exprs = on
+                    .iter_mut()
+                    .flat_map(|(e1, e2)| std::iter::once(e1).chain(std::iter::once(e2)));
+
+                let result = rewrite_expr_iter_mut(exprs, &mut f)?;
 
                 if let Some(filter) = filter.as_mut() {
-                    rewrite_expr(filter, &mut f)?;
+                    result.and_then(|| f(filter))
+                } else {
+                    Ok(result)
                 }
             }
             LogicalPlan::Sort(Sort { expr, .. }) => {
-                rewrite_expr_iter_mut(expr.iter_mut(), &mut f)?
+                rewrite_expr_iter_mut(expr.iter_mut(), f)
             }
             LogicalPlan::Extension(extension) => {
                 // would be nice to avoid this copy -- maybe can
@@ -117,7 +103,7 @@ impl LogicalPlan {
                 todo!();
             }
             LogicalPlan::TableScan(TableScan { filters, .. }) => {
-                rewrite_expr_iter_mut(filters.iter_mut(), &mut f)?;
+                rewrite_expr_iter_mut(filters.iter_mut(), f)
             }
             LogicalPlan::Unnest(Unnest { column, .. }) => {
                 //f(&Expr::Column(column.clone()))
@@ -128,13 +114,14 @@ impl LogicalPlan {
                 select_expr,
                 sort_expr,
                 ..
-            })) => rewrite_expr_iter_mut(
-                on_expr
+            })) => {
+                let exprs = on_expr
                     .iter_mut()
                     .chain(select_expr.iter_mut())
-                    .chain(sort_expr.iter_mut().flat_map(|x| x.iter_mut())),
-                &mut f,
-            )?,
+                    .chain(sort_expr.iter_mut().flat_map(|x| x.iter_mut()));
+
+                rewrite_expr_iter_mut(exprs, f)
+            }
             // plans without expressions
             LogicalPlan::EmptyRelation(_)
             | LogicalPlan::RecursiveQuery(_)
@@ -151,10 +138,8 @@ impl LogicalPlan {
             | LogicalPlan::Ddl(_)
             | LogicalPlan::Copy(_)
             | LogicalPlan::DescribeTable(_)
-            | LogicalPlan::Prepare(_) => {}
+            | LogicalPlan::Prepare(_) => Ok(Transformed::no(())),
         }
-
-        Ok(self)
     }
 }
 
@@ -168,7 +153,7 @@ const PLACEHOLDER: OnceCell<Arc<LogicalPlan>> = OnceCell::new();
 /// of the Arc to avoid cloning the entire plan
 ///
 /// On error, the node will be partially rewritten (left with a placeholder logical plan)
-fn rewrite_arc<F>(node: &mut Arc<LogicalPlan>, f: &mut F) -> Result<Transformed<()>>
+fn rewrite_arc<F>(node: &mut Arc<LogicalPlan>, mut f: F) -> Result<Transformed<()>>
 where
     F: FnMut(&mut LogicalPlan) -> Result<Transformed<()>>,
 {
@@ -215,7 +200,7 @@ impl LogicalPlan {
     where
         F: FnMut(&mut LogicalPlan) -> Result<Transformed<()>>,
     {
-        match self {
+        let children_result = match self {
             LogicalPlan::Projection(Projection { input, .. }) => {
                 rewrite_arc(input, &mut f)
             }
@@ -269,6 +254,27 @@ impl LogicalPlan {
             | LogicalPlan::EmptyRelation { .. }
             | LogicalPlan::Values { .. }
             | LogicalPlan::DescribeTable(_) => Ok(Transformed::no(())),
-        }
+        }?;
+
+        // after visiting the actual children we we need to visit any subqueries
+        // that are inside the expressions
+        children_result.and_then(|| self.rewrite_subqueries(&mut f))
+    }
+
+    /// applies the closure `f` to LogicalPlans in any subquery expressions
+    ///
+    /// If Err is returned, the plan may be left in a partially modified state
+    fn rewrite_subqueries<F>(&mut self, mut f: F) -> Result<Transformed<()>>
+    where
+        F: FnMut(&mut LogicalPlan) -> Result<Transformed<()>>,
+    {
+        self.rewrite_exprs(|expr| match expr {
+            Expr::Exists(Exists { subquery, .. })
+            | Expr::InSubquery(InSubquery { subquery, .. })
+            | Expr::ScalarSubquery(subquery) => {
+                rewrite_arc(&mut subquery.subquery, &mut f)
+            }
+            _ => Ok(Transformed::no(())),
+        })
     }
 }