todo fix issue with collect

alexandre-ricciardi · Jun 4, 2024 · 6462af4 · 6462af4
1 parent 09c68dc
commit 6462af4
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 9 deletions.
diff --git a/integration/cypher-tests/tests/with_collect_tests.rs b/integration/cypher-tests/tests/with_collect_tests.rs
@@ -35,7 +35,7 @@ async fn _test_with_collect(mut client: Client) {
                                                                 create (new)-[:IsFriendOf]->(new2:Person {weight: $weight})
                                                                 return s", p).await;
     let res = client.execute_cypher_request("match (s:Person)-[:IsFriendOf]->(new:Person)-[:IsFriendOf]->(end:Person)
-                                                                with collect(end) as ps
-                                                                return ps").await;
+                                                                return collect(end) 
+                                                                ").await;
     println!("{}", res.expect("persons").to_string());
 }
diff --git a/lib/zawgl-cypher-query-model/src/model/mod.rs b/lib/zawgl-cypher-query-model/src/model/mod.rs
@@ -104,7 +104,17 @@ impl BoolResult {
         BoolResult{name, value}
     }
 }
+#[derive(Debug, Clone)]
+pub struct ListResult {
+    pub name: String,
+    pub values: Vec<EvalResultItem>,
+}
 
+impl ListResult {
+    pub fn new(name: String, values: Vec<EvalResultItem>) -> Self {
+        ListResult{name, values}
+    }
+}
 #[derive(Debug, Clone)]
 pub struct RelationshipResult {
     pub name: String,
@@ -118,12 +128,27 @@ impl RelationshipResult {
         RelationshipResult{name, value, source_nid, target_nid}
     }
 }
+#[derive(Debug, Clone)]
 pub enum EvalResultItem {
     Node(NodeResult),
     Relationship(RelationshipResult),
     Scalar(ScalarResult),
     Bool(BoolResult),
     String(StringResult),
+    List(ListResult),
+}
+
+impl EvalResultItem {
+    pub fn get_name(&self) -> &str {
+        match self {
+            EvalResultItem::Node(i) => &i.name,
+            EvalResultItem::Relationship(i) => &i.name,
+            EvalResultItem::Scalar(i) => &i.name,
+            EvalResultItem::Bool(i) => &i.name,
+            EvalResultItem::String(i) => &i.name,
+            EvalResultItem::List(i) => &i.name,
+        }
+    }
 }
 
 #[derive(Debug, Clone)]

diff --git a/lib/zawgl-front/src/cypher/parser/mod.rs b/lib/zawgl-front/src/cypher/parser/mod.rs
@@ -238,6 +238,10 @@ mod test_parser {
         with collect(end) as ps
         return ps");
     }
-
+    #[test]
+    fn test_return_collect() {
+        run("match (s:Person)-[:IsFriendOf]->(new:Person)-[:IsFriendOf]->(end:Person)
+        return collect(end) as ps");
+    }
 }
 
diff --git a/lib/zawgl-front/src/lib.rs b/lib/zawgl-front/src/lib.rs
@@ -92,7 +92,14 @@ fn eval_item_to_bson(eval_item: &EvalResultItem) -> Result<Document, CypherError
         }),
         EvalResultItem::String(value) =>  Ok(doc! {
             value.name.to_string(): value.value.to_string()
-        })
+        }),
+        EvalResultItem::List(list) => {
+            let mut res = vec![];
+            for item in &list.values {
+                res.push(eval_item_to_bson(item)?)
+            }
+            Ok(doc! {list.name.to_string(): res})
+        }
     }
 }
 

diff --git a/lib/zawgl-front/src/planner/mod.rs b/lib/zawgl-front/src/planner/mod.rs
@@ -26,7 +26,7 @@ use zawgl_core::{graph::{EdgeData, EdgeIndex, NodeIndex}, graph_engine::GraphEng
 mod pattern_builder;
 
 use pattern_builder::{build_pattern, merge_patterns};
-use zawgl_cypher_query_model::{ast::AstVisitorError, model::{BoolResult, EvalResultItem, EvalScopeClause, EvalScopeExpression, NodeResult, RelationshipResult, ScalarResult, StringResult, ValueItem, WhereClause}, QueryResult, QueryStep, StepType};
+use zawgl_cypher_query_model::{ast::AstVisitorError, model::{BoolResult, EvalResultItem, EvalScopeClause, EvalScopeExpression, ListResult, NodeResult, RelationshipResult, ScalarResult, StringResult, ValueItem, WhereClause}, QueryResult, QueryStep, StepType};
 
 use crate::cypher::{parser, query_engine::{where_clause_filter::WhereClauseAstVisitor, CypherError}};
 
@@ -98,18 +98,18 @@ pub fn handle_query_steps(steps: Vec<QueryStep>, graph_engine: &mut GraphEngine)
                 }
             },
             StepType::WITH(eval_scope) => {
-                (_, eval_results) = handle_eval(&mut results, eval_scope)?;
+                (_, eval_results) = handle_eval(&mut results, eval_scope, &eval_results)?;
             },
             StepType::RETURN(eval_scope) => {
-                (result_graphs, return_eval_results) = handle_eval(&mut results, eval_scope)?;
+                (result_graphs, return_eval_results) = handle_eval(&mut results, eval_scope, &eval_results)?;
             },
         }
         first_step = false;
     }
     Ok(QueryResult::new(result_graphs, return_eval_results))
 }
 
-fn handle_eval(results: &mut Vec::<Vec<PropertyGraph>>, eval_scope: EvalScopeClause) -> Result<(Vec<PropertyGraph>, Vec<Vec<EvalResultItem>>), CypherError> {
+fn handle_eval(results: &mut Vec::<Vec<PropertyGraph>>, eval_scope: EvalScopeClause, eval_results: &Vec<Vec<EvalResultItem>>) -> Result<(Vec<PropertyGraph>, Vec<Vec<EvalResultItem>>), CypherError> {
     let matched_graphs = flatten_results(results);
     let mut grouping = Vec::new();
     for ret_exp in &eval_scope.expressions {
@@ -178,6 +178,13 @@ fn handle_eval(results: &mut Vec::<Vec<PropertyGraph>>, eval_scope: EvalScopeCla
                                         }
                                     }
                                 }
+                                for eval_row in eval_results {
+                                    for eval_item in eval_row {
+                                        if eval_item.get_name() == named_item {
+                                            row.push(eval_item.clone());
+                                        }                                            
+                                    }
+                                }
                             }
                         }
                     },
@@ -193,13 +200,16 @@ fn handle_eval(results: &mut Vec::<Vec<PropertyGraph>>, eval_scope: EvalScopeCla
                     let ret_name = if let Some(a) = &fun.alias {
                         a.to_string()
                     } else {
-                        "sum".to_string()
+                        fun.name.to_string()
                     };
                     match fun.name.as_str() {
                         "sum" => {
                             let sum = compute_sum(&fun.args, &graphs);
                             row.push(EvalResultItem::Scalar(ScalarResult::new(ret_name, sum)));
                         },
+                        "collect" => {
+                            row.push(EvalResultItem::List(ListResult::new(ret_name, build_item_list(&fun.args, &graphs)?)));
+                        }
                         _ => {}
                     }
                 },
@@ -323,6 +333,36 @@ fn compute_sum(args: &Vec<ValueItem>, graphs: &Vec<&PropertyGraph>) -> f64 {
     sum_value
 }
 
+fn build_item_list(args: &Vec<ValueItem>, graphs: &Vec<&PropertyGraph>) -> Result<Vec<EvalResultItem>, CypherError> {
+    let mut list = Vec::new();
+    for graph in graphs {
+        for node in graph.get_nodes() {
+            if let Some(var) = node.get_var() {
+                for arg in args {
+                    if let ValueItem::NamedItem(name) = arg {
+                        if name == var {
+                            list.push(EvalResultItem::Node(make_node(None, &name, node)));
+                        }
+                    }
+                }
+            }
+        }
+        for rel in graph.get_edges() {
+            if let Some(var) = rel.relationship.get_var() {
+                for arg in args {
+                    if let ValueItem::NamedItem(name) = arg {
+                        if name == var {
+                            list.push(EvalResultItem::Relationship(make_relationship(None, &name, rel, graph)?));
+                        }
+                    }
+                }
+
+            }
+        }
+    }
+    Ok(list)
+}
+
 fn make_node(alias: Option<&String>, name: &str, node: &Node) -> NodeResult {
     let ret_name = if let Some(a) = alias {
         a.to_string()