Remove petgraph from logical plan dependency (#207)

Author: am357

partiql-eval/src/lib.rs

@@ -41,7 +41,7 @@ mod tests {
     }
 
     // TODO: rename once we move to DAG model completely
-    fn evaluate_dag(logical: LogicalPlan, bindings: MapBindings<Value>) -> Value {
+    fn evaluate_dag(logical: LogicalPlan<BindingsExpr>, bindings: MapBindings<Value>) -> Value {
         // TODO remove once we agree on using evaluate output in the following PR:
         // https://github.com/partiql/partiql-lang-rust/pull/202
         let output = Rc::new(RefCell::new(EvalOutputAccumulator::default()));
@@ -256,15 +256,15 @@ mod tests {
 
     #[test]
     fn select_dag() {
-        // Plan for `select a as b from data`
-        let mut logical = LogicalPlan::new();
-        let from = logical.0.add_node(BindingsExpr::Scan(logical::Scan {
+        let mut lg = LogicalPlan::new();
+
+        let from = lg.add_operator(BindingsExpr::Scan(logical::Scan {
             expr: ValueExpr::VarRef(BindingsName::CaseInsensitive("data".into())),
             as_key: "data".to_string(),
             at_key: None,
         }));
 
-        let project = logical.0.add_node(BindingsExpr::Project(logical::Project {
+        let project = lg.add_operator(BindingsExpr::Project(logical::Project {
             exprs: HashMap::from([(
                 "b".to_string(),
                 ValueExpr::Path(
@@ -276,13 +276,12 @@ mod tests {
             )]),
         }));
 
-        let sink = logical.0.add_node(BindingsExpr::Output);
+        let sink = lg.add_operator(BindingsExpr::Output);
 
-        logical
-            .0
-            .extend_with_edges(&[(from, project), (project, sink)]);
+        lg.add_flow(from, project);
+        lg.add_flow(project, sink);
 
-        if let Value::Bag(b) = evaluate_dag(logical, data_3_tuple()) {
+        if let Value::Bag(b) = evaluate_dag(lg, data_3_tuple()) {
             assert_eq!(b.len(), 3);
         } else {
             panic!("Wrong output")

partiql-eval/src/plan.rs

@@ -39,15 +39,36 @@ impl EvaluatorPlanner {
         }
     }
 
-    pub fn compile_dag(&self, plan: LogicalPlan) -> EvalPlan {
+    pub fn compile_dag(&self, plan: LogicalPlan<BindingsExpr>) -> EvalPlan {
         self.plan_eval_dag(plan)
     }
 
     #[inline]
-    fn plan_eval_dag(&self, lg: LogicalPlan) -> EvalPlan {
-        let plan = lg.0;
+    fn plan_eval_dag(&self, lg: LogicalPlan<BindingsExpr>) -> EvalPlan {
         let mut eval_plan = EvalPlan::default();
-        eval_plan.0 = plan.map(|_, n| self.get_eval_node(n), |_, e| e.clone());
+        let ops = lg.operators();
+        let flows = lg.flows();
+
+        let mut seen = HashMap::new();
+
+        flows.into_iter().for_each(|r| {
+            let (s, d) = r;
+
+            let mut nodes = vec![];
+            for op_id in vec![s, d] {
+                let logical_op = &ops[op_id.index() - 1];
+                let eval_op = if let Some(op) = seen.get(op_id) {
+                    *op
+                } else {
+                    let id = eval_plan.0.add_node(self.get_eval_node(logical_op));
+                    seen.insert(op_id, id);
+                    id
+                };
+                nodes.push(eval_op)
+            }
+
+            eval_plan.0.add_edge(nodes[0], nodes[1], ());
+        });
 
         eval_plan
     }

partiql-logical/Cargo.toml

@@ -21,7 +21,6 @@ edition.workspace = true
 
 [dependencies]
 partiql-value = { path = "../partiql-value" }
-petgraph = "0.6.*"
 ordered-float = "3.*"
 itertools = "0.10.*"
 unicase = "2.*"

partiql-logical/src/lib.rs

@@ -1,14 +1,53 @@
 use partiql_value::{BindingsName, Value};
-use petgraph::prelude::StableGraph;
-use petgraph::Directed;
 use std::collections::HashMap;
 
+#[derive(Eq, PartialEq, Debug, Clone, Copy, Hash)]
+pub struct OpId(usize);
+
+impl OpId {
+    pub fn index(&self) -> usize {
+        self.0
+    }
+}
+
 #[derive(Debug)]
-pub struct LogicalPlan(pub StableGraph<BindingsExpr, (), Directed>);
+pub struct LogicalPlan<T> {
+    nodes: Vec<T>,
+    edges: Vec<(OpId, OpId)>,
+}
 
-impl LogicalPlan {
+impl<T> LogicalPlan<T> {
     pub fn new() -> Self {
-        LogicalPlan(StableGraph::<BindingsExpr, (), Directed>::new())
+        LogicalPlan {
+            nodes: vec![],
+            edges: vec![],
+        }
+    }
+
+    pub fn add_operator(&mut self, op: T) -> OpId {
+        self.nodes.push(op);
+        OpId(self.operator_count())
+    }
+
+    pub fn add_flow(&mut self, src: OpId, dst: OpId) {
+        let src_idx = src.index() - 1;
+        let dst_idx = dst.index() - 1;
+        assert!(src_idx <= self.operator_count());
+        assert!(dst_idx <= self.operator_count());
+
+        self.edges.push((src, dst));
+    }
+
+    pub fn operator_count(&self) -> usize {
+        self.nodes.len()
+    }
+
+    pub fn operators(&self) -> &Vec<T> {
+        &self.nodes
+    }
+
+    pub fn flows(&self) -> &Vec<(OpId, OpId)> {
+        &self.edges
     }
 }
 
@@ -133,3 +172,21 @@ pub struct SelectValue {
 pub struct Distinct {
     pub out: Box<BindingsExpr>,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_plan() {
+        let mut p: LogicalPlan<BindingsExpr> = LogicalPlan::new();
+        let a = p.add_operator(BindingsExpr::OrderBy);
+        let b = p.add_operator(BindingsExpr::Output);
+        let c = p.add_operator(BindingsExpr::Limit);
+        p.add_flow(a, b);
+        p.add_flow(a, c);
+        p.add_flow(b, c);
+        assert_eq!(3, p.operators().len());
+        assert_eq!(3, p.flows().len());
+    }
+}