Add plan flows for JOIN. Add 'self managed' graph nodes. (#385)

Author: jpschorr

partiql-eval/src/eval/evaluable.rs

@@ -27,13 +27,22 @@ macro_rules! take_input {
     };
 }
 
+/// Whether an [`Evaluable`] takes input from the plan graph or manages its own iteration.
+pub enum EvalType {
+    SelfManaged,
+    GraphManaged,
+}
+
 /// `Evaluable` represents each evaluation operator in the evaluation plan as an evaluable entity.
 pub trait Evaluable: Debug {
     fn evaluate(&mut self, ctx: &dyn EvalContext) -> Value;
     fn update_input(&mut self, input: Value, branch_num: u8);
     fn get_vars(&self) -> Option<&[String]> {
         None
     }
+    fn eval_type(&self) -> EvalType {
+        EvalType::GraphManaged
+    }
 }
 
 /// Represents an evaluation `Scan` operator; `Scan` operator scans the given bindings from its
@@ -305,6 +314,10 @@ impl Evaluable for EvalJoin {
     fn update_input(&mut self, input: Value, _branch_num: u8) {
         self.input = Some(input);
     }
+
+    fn eval_type(&self) -> EvalType {
+        EvalType::SelfManaged
+    }
 }
 
 /// An SQL aggregation function call that has been rewritten to be evaluated with the `GROUP BY`

partiql-eval/src/eval/mod.rs

@@ -18,7 +18,7 @@ use petgraph::graph::NodeIndex;
 use crate::error::{EvalErr, EvaluationError};
 use petgraph::visit::EdgeRef;
 
-use crate::eval::evaluable::Evaluable;
+use crate::eval::evaluable::{EvalType, Evaluable};
 
 pub mod evaluable;
 pub mod expr;
@@ -77,32 +77,44 @@ impl EvalPlan {
 
         let mut result = None;
         for idx in ops.into_iter() {
-            result = Some(self.get_node(idx)?.evaluate(&*ctx));
-
-            // return on first evaluation error
-            if ctx.has_errors() {
-                return Err(EvalErr {
-                    errors: ctx.errors(),
-                });
-            }
-
             let destinations: Vec<(usize, (u8, NodeIndex))> = self
                 .plan_graph()
                 .edges_directed(idx, Outgoing)
                 .map(|e| (*e.weight(), e.target()))
                 .enumerate()
                 .collect_vec();
-            let branches = destinations.len();
-            for (i, (branch_num, dst_id)) in destinations {
-                let res = if i == branches - 1 {
-                    result.take()
-                } else {
-                    result.clone()
-                };
-
-                let res =
-                    res.ok_or_else(|| err_illegal_state("Error in retrieving source value"))?;
-                self.get_node(dst_id)?.update_input(res, branch_num);
+
+            // Some evaluables (i.e., `JOIN`) manage their own inputs
+            let graph_managed = destinations.is_empty()
+                || destinations.iter().any(|(_, (_, dest_idx))| {
+                    matches!(
+                        self.get_node(*dest_idx).map(|d| d.eval_type()),
+                        Ok(EvalType::GraphManaged)
+                    )
+                });
+            if graph_managed {
+                let src = self.get_node(idx)?;
+                result = Some(src.evaluate(&*ctx));
+
+                // return on first evaluation error
+                if ctx.has_errors() {
+                    return Err(EvalErr {
+                        errors: ctx.errors(),
+                    });
+                }
+
+                let num_destinations = destinations.len();
+                for (i, (branch_num, dst_id)) in destinations {
+                    let res = if i == num_destinations - 1 {
+                        result.take()
+                    } else {
+                        result.clone()
+                    };
+
+                    let res =
+                        res.ok_or_else(|| err_illegal_state("Error in retrieving source value"))?;
+                    self.get_node(dst_id)?.update_input(res, branch_num);
+                }
             }
         }
 

partiql-logical-planner/src/lower.rs

@@ -1345,6 +1345,8 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> {
             right,
         });
         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);
         self.push_bexpr(join);
         Traverse::Continue
     }