Enable subquery in From (#560)

Author: jpschorr

partiql-logical-planner/src/lower.rs

@@ -23,7 +23,7 @@ use std::borrow::Cow;
 
 use partiql_value::BindingsName;
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 
 use crate::builtins::{FnSymTab, FN_SYM_TAB};
 use itertools::Itertools;
@@ -57,12 +57,21 @@ macro_rules! eq_or_fault {
 
 #[macro_export]
 macro_rules! true_or_fault {
+    ($self:ident, $expr:expr, $msg:expr) => {
+        if !$expr {
+            true_or_fault_err!($self, $expr, $msg);
+            return partiql_ast::visit::Traverse::Stop;
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! true_or_fault_err {
     ($self:ident, $expr:expr, $msg:expr) => {
         if !$expr {
             $self
                 .errors
                 .push(AstTransformError::IllegalState($msg.to_string()));
-            return partiql_ast::visit::Traverse::Stop;
         }
     };
 }
@@ -161,9 +170,6 @@ pub struct AstToLogical<'a> {
     path_stack: Vec<Vec<PathComponent>>,
     sort_stack: Vec<Vec<logical::SortSpec>>,
     aggregate_exprs: Vec<AggregateExpression>,
-
-    from_lets: HashSet<NodeId>,
-
     projection_renames: Vec<FnvIndexMap<String, BindingsName<'a>>>,
 
     aliases: FnvIndexMap<NodeId, SymbolPrimitive>,
@@ -173,7 +179,7 @@ pub struct AstToLogical<'a> {
     agg_id: IdGenerator,
 
     // output
-    plan: LogicalPlan<BindingsOp>,
+    plan_stack: Vec<LogicalPlan<BindingsOp>>,
 
     // catalog & data flow data
     key_registry: name_resolver::KeyRegistry,
@@ -233,8 +239,6 @@ impl<'a> AstToLogical<'a> {
             sort_stack: Default::default(),
             aggregate_exprs: Default::default(),
 
-            from_lets: Default::default(),
-
             projection_renames: Default::default(),
 
             aliases: Default::default(),
@@ -244,7 +248,7 @@ impl<'a> AstToLogical<'a> {
             agg_id: Default::default(),
 
             // output
-            plan: Default::default(),
+            plan_stack: Default::default(),
 
             key_registry: registry,
             fnsym_tab,
@@ -258,13 +262,19 @@ impl<'a> AstToLogical<'a> {
         mut self,
         query: &ast::AstNode<ast::TopLevelQuery>,
     ) -> Result<logical::LogicalPlan<logical::BindingsOp>, AstTransformationError> {
+        self.enter_plan();
         query.visit(&mut self);
+        true_or_fault_err!(
+            self,
+            self.plan_stack.len() == 1,
+            "self.plan_stack.len() != 1"
+        );
         if !self.errors.is_empty() {
             return Err(AstTransformationError {
                 errors: self.errors,
             });
         }
-        Ok(self.plan)
+        Ok(self.plan_stack.pop().unwrap())
     }
 
     #[inline]
@@ -509,6 +519,21 @@ impl<'a> AstToLogical<'a> {
         self.q_stack.last_mut().unwrap()
     }
 
+    #[inline]
+    fn enter_plan(&mut self) {
+        self.plan_stack.push(LogicalPlan::default());
+    }
+
+    #[inline]
+    fn exit_plan(&mut self) -> LogicalPlan<BindingsOp> {
+        self.plan_stack.pop().expect("environment level")
+    }
+
+    #[inline]
+    fn curr_plan(&mut self) -> &mut LogicalPlan<BindingsOp> {
+        self.plan_stack.last_mut().expect("plan")
+    }
+
     #[inline]
     fn enter_benv(&mut self) {
         self.bexpr_stack.push(vec![]);
@@ -691,8 +716,8 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
         let mut benv = self.exit_benv();
         eq_or_fault!(self, benv.len(), 1, "Expect benv.len() == 1");
         let out = benv.pop().unwrap();
-        let sink_id = self.plan.add_operator(BindingsOp::Sink);
-        self.plan.add_flow(out, sink_id);
+        let sink_id = self.curr_plan().add_operator(BindingsOp::Sink);
+        self.curr_plan().add_flow(out, sink_id);
         Traverse::Continue
     }
 
@@ -712,7 +737,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
                 let mut clauses = clauses.evaluation_order().into_iter();
                 if let Some(mut src_id) = clauses.next() {
                     for dst_id in clauses {
-                        self.plan.add_flow(src_id, dst_id);
+                        self.curr_plan().add_flow(src_id, dst_id);
                         src_id = dst_id;
                     }
                     self.push_bexpr(src_id);
@@ -726,7 +751,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
                 );
                 let mut out = *benv.first().unwrap();
                 benv.into_iter().skip(1).for_each(|op| {
-                    self.plan.add_flow(out, op);
+                    self.curr_plan().add_flow(out, op);
                     out = op;
                 });
                 self.push_bexpr(out);
@@ -759,7 +784,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
 
         match query_set {
             QuerySet::BagOp(bag_op) => {
-                eq_or_fault!(self, benv.len(), 2, "benv.len() != 2");
+                eq_or_fault!(self, benv.len(), 2, "qs benv.len() != 2");
                 let rid = benv.pop().unwrap();
                 let lid = benv.pop().unwrap();
 
@@ -777,20 +802,20 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
                     None => logical::SetQuantifier::Distinct,
                 };
 
-                let id = self.plan.add_operator(BindingsOp::BagOp(BagOp {
+                let id = self.curr_plan().add_operator(BindingsOp::BagOp(BagOp {
                     bag_op: bag_operator,
                     setq,
                 }));
-                self.plan.add_flow_with_branch_num(lid, id, 0);
-                self.plan.add_flow_with_branch_num(rid, id, 1);
+                self.curr_plan().add_flow_with_branch_num(lid, id, 0);
+                self.curr_plan().add_flow_with_branch_num(rid, id, 1);
                 self.push_bexpr(id);
             }
             QuerySet::Select(_) => {}
             QuerySet::Expr(_) => {
                 eq_or_fault!(self, env.len(), 1, "env.len() != 1");
                 let expr = env.into_iter().next().unwrap();
                 let op = BindingsOp::ExprQuery(logical::ExprQuery { expr });
-                let id = self.plan.add_operator(op);
+                let id = self.curr_plan().add_operator(op);
                 self.push_bexpr(id);
             }
             QuerySet::Values(_) => {
@@ -839,7 +864,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
                 aggregate_exprs: self.aggregate_exprs.clone(),
                 group_as_alias: None,
             });
-            let id = self.plan.add_operator(group_by);
+            let id = self.curr_plan().add_operator(group_by);
             self.current_clauses_mut().group_by_clause.replace(id);
         }
         Traverse::Continue
@@ -859,7 +884,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
         eq_or_fault!(self, env.len(), 0, "env.len() != 0");
 
         if let Some(SetQuantifier::Distinct) = projection.setq {
-            let id = self.plan.add_operator(BindingsOp::Distinct);
+            let id = self.curr_plan().add_operator(BindingsOp::Distinct);
             self.current_clauses_mut().distinct.replace(id);
         }
         Traverse::Continue
@@ -934,7 +959,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
                 logical::BindingsOp::ProjectValue(logical::ProjectValue { expr })
             }
         };
-        let id = self.plan.add_operator(select);
+        let id = self.curr_plan().add_operator(select);
         self.current_clauses_mut().select_clause.replace(id);
         Traverse::Continue
     }
@@ -1468,8 +1493,9 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
     }
 
     fn enter_from_let(&mut self, from_let: &'ast FromLet) -> Traverse {
-        self.from_lets.insert(*self.current_node());
         *self.current_ctx_mut() = QueryContext::FromLet;
+        self.enter_plan();
+        self.enter_benv();
         self.enter_env();
 
         let id = *self.current_node();
@@ -1485,10 +1511,20 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
 
     fn exit_from_let(&mut self, from_let: &'ast FromLet) -> Traverse {
         *self.current_ctx_mut() = QueryContext::Query;
+        let subplan = self.exit_plan();
+        let benv = self.exit_benv();
         let mut env = self.exit_env();
-        eq_or_fault!(self, env.len(), 1, "env.len() != 1");
+        eq_or_fault!(self, env.len() + benv.len(), 1, "env.len()+benv.len() != 1");
 
-        let expr = env.pop().unwrap();
+        let expr = if !benv.is_empty() {
+            // Subquery in From Let
+            let subq = logical::SubQueryExpr { plan: subplan };
+            ValueExpr::SubQueryExpr(subq)
+        } else {
+            // Expression in From Let
+            self.curr_plan().merge_plan(subplan); // merge in subplan, as there is no subquery
+            env.pop().unwrap()
+        };
 
         let FromLet {
             kind,
@@ -1527,10 +1563,13 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
                 ProjectAllMode::Unwrap,
             ),
         };
-        let id = self.plan.add_operator(bexpr);
+
+        let id = self.curr_plan().add_operator(bexpr);
         self.push_bexpr(id);
+
         if let Some(select_id) = self.current_clauses_mut().select_clause {
-            if let Some(BindingsOp::ProjectAll(mode)) = self.plan.operator_as_mut(select_id) {
+            if let Some(BindingsOp::ProjectAll(mode)) = self.curr_plan().operator_as_mut(select_id)
+            {
                 *mode = project_all_mode
             }
         }
@@ -1546,7 +1585,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
 
     fn exit_join(&mut self, join: &'ast Join) -> Traverse {
         let mut benv = self.exit_benv();
-        eq_or_fault!(self, benv.len(), 2, "benv.len() != 2");
+        eq_or_fault!(self, benv.len(), 2, "j benv.len() != 2");
 
         let mut env = self.exit_env();
         true_or_fault!(
@@ -1569,17 +1608,17 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
 
         let rid = benv.pop().unwrap();
         let lid = benv.pop().unwrap();
-        let left = Box::new(self.plan.operator(lid).unwrap().clone());
-        let right = Box::new(self.plan.operator(rid).unwrap().clone());
+        let left = Box::new(self.curr_plan().operator(lid).unwrap().clone());
+        let right = Box::new(self.curr_plan().operator(rid).unwrap().clone());
         let join = logical::BindingsOp::Join(logical::Join {
             kind,
             left,
             right,
             on,
         });
-        let join = self.plan.add_operator(join);
-        self.plan.add_flow_with_branch_num(lid, join, 0);
-        self.plan.add_flow_with_branch_num(rid, join, 1);
+        let join = self.curr_plan().add_operator(join);
+        self.curr_plan().add_flow_with_branch_num(lid, join, 0);
+        self.curr_plan().add_flow_with_branch_num(rid, join, 1);
         self.push_bexpr(join);
         Traverse::Continue
     }
@@ -1611,7 +1650,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
         let filter = logical::BindingsOp::Filter(logical::Filter {
             expr: env.pop().unwrap(),
         });
-        let id = self.plan.add_operator(filter);
+        let id = self.curr_plan().add_operator(filter);
 
         self.current_clauses_mut().where_clause.replace(id);
         Traverse::Continue
@@ -1629,7 +1668,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
         let having = BindingsOp::Having(logical::Having {
             expr: env.pop().unwrap(),
         });
-        let id = self.plan.add_operator(having);
+        let id = self.curr_plan().add_operator(having);
 
         self.current_clauses_mut().having_clause.replace(id);
         Traverse::Continue
@@ -1679,8 +1718,9 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
             ));
             return Traverse::Stop;
         }
+        let mut errors = Vec::default();
         let select_clause = self
-            .plan
+            .curr_plan()
             .operator_as_mut(select_clause_op_id.expect("select_clause_op_id not None"))
             .unwrap();
         let mut binding = Vec::new();
@@ -1705,14 +1745,14 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
                     logical::Lit::String(s) => s.clone(),
                     _ => {
                         // Report error but allow visitor to continue
-                        self.errors.push(AstTransformError::IllegalState(
+                        errors.push(AstTransformError::IllegalState(
                             "Unexpected literal type".to_string(),
                         ));
                         String::new()
                     }
                 },
                 _ => {
-                    self.errors.push(AstTransformError::IllegalState(
+                    errors.push(AstTransformError::IllegalState(
                         "Unexpected alias type".to_string(),
                     ));
                     return Traverse::Stop;
@@ -1728,14 +1768,16 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
             exprs.insert(alias, value);
         }
 
+        self.errors.extend(errors);
+
         let group_by: BindingsOp = BindingsOp::GroupBy(logical::GroupBy {
             strategy,
             exprs,
             aggregate_exprs,
             group_as_alias,
         });
 
-        let id = self.plan.add_operator(group_by);
+        let id = self.curr_plan().add_operator(group_by);
         self.current_clauses_mut().group_by_clause.replace(id);
         Traverse::Continue
     }
@@ -1763,7 +1805,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
     fn exit_order_by_expr(&mut self, _order_by_expr: &'ast OrderByExpr) -> Traverse {
         let specs = self.exit_sort();
         let order_by = logical::BindingsOp::OrderBy(logical::OrderBy { specs });
-        let id = self.plan.add_operator(order_by);
+        let id = self.curr_plan().add_operator(order_by);
         if matches!(self.current_ctx(), Some(QueryContext::Query)) {
             self.current_clauses_mut().order_by_clause.replace(id);
         } else {
@@ -1836,7 +1878,7 @@ impl<'ast> Visitor<'ast> for AstToLogical<'_> {
         };
 
         let limit_offset = logical::BindingsOp::LimitOffset(logical::LimitOffset { limit, offset });
-        let id = self.plan.add_operator(limit_offset);
+        let id = self.curr_plan().add_operator(limit_offset);
         if matches!(self.current_ctx(), Some(QueryContext::Query)) {
             self.current_clauses_mut().limit_offset_clause.replace(id);
         } else {