Add INNER and CROSS JOINs to planner and evaluator (#218)

Author: alancai98

partiql-eval/src/eval.rs

@@ -6,7 +6,7 @@ use thiserror::Error;
 
 use petgraph::algo::toposort;
 use petgraph::prelude::StableGraph;
-use petgraph::{Directed, Incoming, Outgoing};
+use petgraph::{Directed, Outgoing};
 
 use partiql_value::Value::{Boolean, Missing, Null};
 use partiql_value::{
@@ -18,7 +18,7 @@ use crate::env::basic::MapBindings;
 use crate::env::Bindings;
 
 #[derive(Debug)]
-pub struct EvalPlan(pub StableGraph<Box<dyn Evaluable>, (), Directed>);
+pub struct EvalPlan(pub StableGraph<Box<dyn Evaluable>, u8, Directed>);
 
 impl Default for EvalPlan {
     fn default() -> Self {
@@ -28,7 +28,7 @@ impl Default for EvalPlan {
 
 impl EvalPlan {
     fn new() -> Self {
-        EvalPlan(StableGraph::<Box<dyn Evaluable>, (), Directed>::new())
+        EvalPlan(StableGraph::<Box<dyn Evaluable>, u8, Directed>::new())
     }
 }
 
@@ -50,7 +50,7 @@ pub enum EvaluationError {
 
 pub trait Evaluable: Debug {
     fn evaluate(&mut self, ctx: &dyn EvalContext) -> Option<Value>;
-    fn update_input(&mut self, input: &Value);
+    fn update_input(&mut self, input: &Value, branch_num: u8);
 }
 
 #[derive(Debug)]
@@ -108,11 +108,103 @@ impl Evaluable for EvalScan {
         self.output.clone()
     }
 
-    fn update_input(&mut self, _input: &Value) {
+    fn update_input(&mut self, _input: &Value, _branch_num: u8) {
         todo!("update_input for Scan")
     }
 }
 
+#[derive(Debug)]
+pub enum EvalJoinKind {
+    Inner,
+    Left,
+    Right,
+    Full,
+    Cross,
+}
+
+#[derive(Debug)]
+pub struct EvalJoin {
+    pub kind: EvalJoinKind,
+    pub on: Option<Box<dyn EvalExpr>>,
+    pub input_l: Option<Value>,
+    pub input_r: Option<Value>,
+    pub output: Option<Value>,
+}
+
+impl EvalJoin {
+    pub fn new(kind: EvalJoinKind, on: Option<Box<dyn EvalExpr>>) -> Self {
+        EvalJoin {
+            kind,
+            on,
+            input_l: None,
+            input_r: None,
+            output: None,
+        }
+    }
+}
+
+impl Evaluable for EvalJoin {
+    fn evaluate(&mut self, ctx: &dyn EvalContext) -> Option<Value> {
+        // TODO: PartiQL defaults to lateral JOINs (RHS can reference binding tuples defined from the LHS)
+        //  https://partiql.org/assets/PartiQL-Specification.pdf#subsection.5.3. Adding this behavior
+        //  to be spec-compliant may result in changes to the DAG flows.
+        let output = match self.kind {
+            EvalJoinKind::Inner => {
+                let mut result = partiql_bag!();
+                for binding_tuple_l in self.input_l.clone().unwrap() {
+                    let binding_tuple_l = binding_tuple_l.coerce_to_tuple();
+                    for binding_tuple_r in self.input_r.clone().unwrap() {
+                        let binding_tuple_r = binding_tuple_r.coerce_to_tuple();
+                        let mut new_result = binding_tuple_l.clone();
+                        for pairs in binding_tuple_r.pairs() {
+                            new_result.insert(pairs.0, pairs.1.clone());
+                        }
+                        if let Some(on_condition) = &self.on {
+                            if on_condition.evaluate(&new_result, ctx) == Boolean(true) {
+                                result.push(new_result.into());
+                            }
+                        } else {
+                            result.push(new_result.into());
+                        }
+                    }
+                }
+                Some(result.into())
+            }
+            EvalJoinKind::Left => {
+                todo!("Left JOINs")
+            }
+            EvalJoinKind::Cross => {
+                let mut result = partiql_bag!();
+                for binding_tuple_l in self.input_l.clone().unwrap() {
+                    let binding_tuple_l = binding_tuple_l.coerce_to_tuple();
+                    for binding_tuple_r in self.input_r.clone().unwrap() {
+                        let binding_tuple_r = binding_tuple_r.coerce_to_tuple();
+                        let mut new_result = binding_tuple_l.clone();
+                        for pairs in binding_tuple_r.pairs() {
+                            new_result.insert(pairs.0, pairs.1.clone());
+                        }
+                        result.push(new_result.into());
+                    }
+                }
+                Some(result.into())
+            }
+            EvalJoinKind::Full | EvalJoinKind::Right => {
+                todo!("Full and Right Joins are not yet implemented for `partiql-lang-rust`")
+            }
+        };
+        self.output = output;
+        self.output.clone()
+    }
+
+    fn update_input(&mut self, input: &Value, branch_num: u8) {
+        match branch_num {
+            0 => self.input_l = Some(input.clone()),
+            1 => self.input_r = Some(input.clone()),
+            _ => panic!("EvalJoin nodes only support `0` and `1` for the `branch_num`"),
+        };
+    }
+}
+
 #[derive(Debug)]
 pub struct EvalUnpivot {
     pub expr: Box<dyn EvalExpr>,
@@ -154,7 +246,7 @@ impl Evaluable for EvalUnpivot {
         self.output.clone()
     }
 
-    fn update_input(&mut self, _input: &Value) {
+    fn update_input(&mut self, _input: &Value, _branch_num: u8) {
         todo!()
     }
 }
@@ -206,7 +298,7 @@ impl Evaluable for EvalFilter {
         self.output = Some(Value::Bag(Box::new(out)));
         self.output.clone()
     }
-    fn update_input(&mut self, input: &Value) {
+    fn update_input(&mut self, input: &Value, _branch_num: u8) {
         self.input = Some(input.clone())
     }
 }
@@ -257,7 +349,7 @@ impl Evaluable for EvalProject {
         self.output.clone()
     }
 
-    fn update_input(&mut self, input: &Value) {
+    fn update_input(&mut self, input: &Value, _branch_num: u8) {
         self.input = Some(input.clone());
     }
 }
@@ -304,7 +396,7 @@ impl Evaluable for EvalProjectValue {
         self.output.clone()
     }
 
-    fn update_input(&mut self, input: &Value) {
+    fn update_input(&mut self, input: &Value, _branch_num: u8) {
         self.input = Some(input.clone());
     }
 }
@@ -436,7 +528,7 @@ impl Evaluable for EvalDistinct {
         self.output.clone()
     }
 
-    fn update_input(&mut self, input: &Value) {
+    fn update_input(&mut self, input: &Value, _branch_num: u8) {
         self.input = Some(input.clone());
     }
 }
@@ -451,7 +543,7 @@ impl Evaluable for EvalSink {
     fn evaluate(&mut self, _ctx: &dyn EvalContext) -> Option<Value> {
         self.input.clone()
     }
-    fn update_input(&mut self, input: &Value) {
+    fn update_input(&mut self, input: &Value, _branch_num: u8) {
         self.input = Some(input.clone());
     }
 }
@@ -627,39 +719,33 @@ impl Evaluator {
         // that all v ∈ V \{v0} are reachable from v0. Note that this is the definition of trees
         // without the condition |E| = |V | − 1. Hence, all trees are DAGs.
         // Reference: https://link.springer.com/article/10.1007/s00450-009-0061-0
-        match graph.externals(Incoming).exactly_one() {
-            Ok(_) => {
-                let sorted_ops = toposort(&graph, None);
-                match sorted_ops {
-                    Ok(ops) => {
-                        let mut result = None;
-                        for idx in ops.into_iter() {
-                            let src = graph
-                                .node_weight_mut(idx)
-                                .expect("Error in retrieving node");
-                            result = src.evaluate(&*self.ctx);
-
-                            let mut ne = graph.neighbors_directed(idx, Outgoing).detach();
-                            while let Some(n) = ne.next_node(&graph) {
-                                let dst =
-                                    graph.node_weight_mut(n).expect("Error in retrieving node");
-                                dst.update_input(
-                                    &result.clone().expect("Error in retrieving source value"),
-                                );
-                            }
-                        }
-                        let evaluated = Evaluated {
-                            result: result.expect("Error in retrieving eval output"),
-                        };
-                        Ok(evaluated)
+        let sorted_ops = toposort(&graph, None);
+        match sorted_ops {
+            Ok(ops) => {
+                let mut result = None;
+                for idx in ops.into_iter() {
+                    let src = graph
+                        .node_weight_mut(idx)
+                        .expect("Error in retrieving node");
+                    result = src.evaluate(&*self.ctx);
+
+                    let mut ne = graph.neighbors_directed(idx, Outgoing).detach();
+                    while let Some((e, n)) = ne.next(&graph) {
+                        // use the edge weight to store the `branch_num`
+                        let branch_num = *graph
+                            .edge_weight(e)
+                            .expect("Error in retrieving weight for edge");
+                        let dst = graph.node_weight_mut(n).expect("Error in retrieving node");
+                        dst.update_input(
+                            &result.clone().expect("Error in retrieving source value"),
+                            branch_num,
+                        );
                     }
-                    Err(e) => Err(EvalErr {
-                        errors: vec![EvaluationError::InvalidEvaluationPlan(format!(
-                            "Malformed evaluation plan detected: {:?}",
-                            e
-                        ))],
-                    }),
                 }
+                let evaluated = Evaluated {
+                    result: result.expect("Error in retrieving eval output"),
+                };
+                Ok(evaluated)
             }
             Err(e) => Err(EvalErr {
                 errors: vec![EvaluationError::InvalidEvaluationPlan(format!(

partiql-eval/src/lib.rs

@@ -17,9 +17,9 @@ mod tests {
     use partiql_logical as logical;
     use partiql_logical::BindingsExpr::{Distinct, Project, ProjectValue};
     use partiql_logical::{
-        BagExpr, BinaryOp, BindingsExpr, ListExpr, LogicalPlan, PathComponent, TupleExpr, ValueExpr,
+        BagExpr, BinaryOp, BindingsExpr, JoinKind, ListExpr, LogicalPlan, PathComponent, TupleExpr,
+        ValueExpr,
     };
-
     use partiql_value as value;
     use partiql_value::{
         partiql_bag, partiql_list, partiql_tuple, Bag, BindingsName, List, Tuple, Value,
@@ -85,6 +85,23 @@ mod tests {
         )
     }
 
+    fn join_data() -> MapBindings<Value> {
+        let customers = partiql_list![
+            partiql_tuple![("id", 5), ("name", "Joe")],
+            partiql_tuple![("id", 7), ("name", "Mary")],
+        ];
+
+        let orders = partiql_list![
+            partiql_tuple![("custId", 7), ("productId", 101)],
+            partiql_tuple![("custId", 7), ("productId", 523)],
+        ];
+
+        let mut bindings = MapBindings::default();
+        bindings.insert("customers", customers.into());
+        bindings.insert("orders", orders.into());
+        bindings
+    }
+
     // Creates the plan: `SELECT <lhs> <op> <rhs> AS result FROM data` where <lhs> comes from data
     // Evaluates the plan and asserts the result is a bag of the tuple mapping to `expected_first_elem`
     // (i.e. <<{'result': <expected_first_elem>}>>)
@@ -452,6 +469,94 @@ mod tests {
         );
     }
 
+    #[test]
+    fn select_with_cross_join() {
+        let mut lg = LogicalPlan::new();
+
+        // Example 9 from spec with projected columns from different tables demonstrates a cross join:
+        // SELECT c.id, c.name, o.custId, o.productId FROM customers AS c, orders AS o
+        let from_lhs = lg.add_operator(scan("customers", "c"));
+        let from_rhs = lg.add_operator(scan("orders", "o"));
+
+        let project = lg.add_operator(Project(logical::Project {
+            exprs: HashMap::from([
+                ("id".to_string(), path_var("c", "id")),
+                ("name".to_string(), path_var("c", "name")),
+                ("custId".to_string(), path_var("o", "custId")),
+                ("productId".to_string(), path_var("o", "productId")),
+            ]),
+        }));
+
+        let join = lg.add_operator(BindingsExpr::Join(logical::Join {
+            kind: JoinKind::Cross,
+            on: None,
+        }));
+
+        let sink = lg.add_operator(BindingsExpr::Sink);
+        lg.add_flow_with_branch_num(from_lhs, join, 0);
+        lg.add_flow_with_branch_num(from_rhs, join, 1);
+        lg.add_flow_with_branch_num(join, project, 0);
+        lg.add_flow_with_branch_num(project, sink, 0);
+
+        let out = evaluate(lg, join_data());
+        println!("{:?}", &out);
+
+        assert_matches!(out, Value::Bag(bag) => {
+            let expected = partiql_bag![
+                partiql_tuple![("custId", 7), ("name", "Joe"), ("id", 5), ("productId", 101)],
+                partiql_tuple![("custId", 7), ("name", "Joe"), ("id", 5), ("productId", 523)],
+                partiql_tuple![("custId", 7), ("name", "Mary"), ("id", 7), ("productId", 101)],
+                partiql_tuple![("custId", 7), ("name", "Mary"), ("id", 7), ("productId", 523)],
+            ];
+            assert_eq!(*bag, expected);
+        });
+    }
+
+    #[test]
+    fn select_with_join_and_on() {
+        let mut lg = LogicalPlan::new();
+
+        // Similar to ex 9 from spec with projected columns from different tables with an inner JOIN and ON condition
+        // SELECT c.id, c.name, o.custId, o.productId FROM customers AS c, orders AS o ON c.id = o.custId
+        let from_lhs = lg.add_operator(scan("customers", "c"));
+        let from_rhs = lg.add_operator(scan("orders", "o"));
+
+        let project = lg.add_operator(Project(logical::Project {
+            exprs: HashMap::from([
+                ("id".to_string(), path_var("c", "id")),
+                ("name".to_string(), path_var("c", "name")),
+                ("custId".to_string(), path_var("o", "custId")),
+                ("productId".to_string(), path_var("o", "productId")),
+            ]),
+        }));
+
+        let join = lg.add_operator(BindingsExpr::Join(logical::Join {
+            kind: JoinKind::Inner,
+            on: Some(ValueExpr::BinaryExpr(
+                BinaryOp::Eq,
+                Box::new(path_var("c", "id")),
+                Box::new(path_var("o", "custId")),
+            )),
+        }));
+
+        let sink = lg.add_operator(BindingsExpr::Sink);
+        lg.add_flow_with_branch_num(from_lhs, join, 0);
+        lg.add_flow_with_branch_num(from_rhs, join, 1);
+        lg.add_flow_with_branch_num(join, project, 0);
+        lg.add_flow_with_branch_num(project, sink, 0);
+
+        let out = evaluate(lg, join_data());
+        println!("{:?}", &out);
+
+        assert_matches!(out, Value::Bag(bag) => {
+            let expected = partiql_bag![
+                partiql_tuple![("custId", 7), ("name", "Mary"), ("id", 7), ("productId", 101)],
+                partiql_tuple![("custId", 7), ("name", "Mary"), ("id", 7), ("productId", 523)],
+            ];
+            assert_eq!(*bag, expected);
+        });
+    }
+
     #[test]
     fn select() {
         let mut lg = LogicalPlan::new();