Add Union/Intersect/Except to parser (#111)

Author: jpschorr

partiql-ast/src/ast.rs

@@ -64,13 +64,21 @@ impl<T> ToAstNode for T {}
 /// for creating the node. See [ToAstNode] for more details on the usage.
 ///
 /// [1]: https://crates.io/crates/derive_builder
-#[derive(Builder, Clone, Eq, PartialEq, Debug)]
+#[derive(Builder, Clone, Debug)]
 pub struct AstNode<T, Loc: Display> {
     pub node: T,
     #[builder(setter(strip_option), default)]
     pub location: Option<Location<Loc>>,
 }
 
+impl<T: PartialEq, Loc: Display> PartialEq for AstNode<T, Loc> {
+    fn eq(&self, other: &Self) -> bool {
+        self.node == other.node
+    }
+}
+
+impl<T: Eq, Loc: Display> Eq for AstNode<T, Loc> {}
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct Item {
     pub kind: ItemKind,
@@ -206,11 +214,6 @@ pub enum ConflictAction {
     DoNothing,
 }
 
-#[derive(Clone, Debug, PartialEq)]
-pub struct Query {
-    pub expr: Box<Expr>,
-}
-
 #[derive(Clone, Debug, PartialEq)]
 pub struct Expr {
     pub kind: ExprKind,
@@ -233,9 +236,7 @@ pub type BetweenAst = AstBytePos<Between>;
 pub type InAst = AstBytePos<In>;
 pub type SimpleCaseAst = AstBytePos<SimpleCase>;
 pub type SearchCaseAst = AstBytePos<SearchCase>;
-pub type UnionAst = AstBytePos<Union>;
-pub type ExceptAst = AstBytePos<Except>;
-pub type IntersectAst = AstBytePos<Intersect>;
+pub type SetExprAst = AstBytePos<SetExpr>;
 pub type PathAst = AstBytePos<Path>;
 pub type CallAst = AstBytePos<Call>;
 pub type CallAggAst = AstBytePos<CallAgg>;
@@ -251,6 +252,38 @@ pub type GroupByExprAst = AstBytePos<GroupByExpr>;
 pub type GroupKeyAst = AstBytePos<GroupKey>;
 pub type OrderByExprAst = AstBytePos<OrderByExpr>;
 pub type SortSpecAst = AstBytePos<SortSpec>;
+pub type QueryAst = AstBytePos<Query>;
+pub type QuerySetAst = AstBytePos<QuerySet>;
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct Query {
+    pub set: QuerySetAst,
+    pub order_by: Option<Box<OrderByExprAst>>,
+    pub limit: Option<Box<Expr>>,
+    pub offset: Option<Box<Expr>>,
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum QuerySet {
+    SetOp(Box<SetExprAst>),
+    Select(Box<SelectAst>),
+    Expr(Box<Expr>),
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct SetExpr {
+    pub setop: SetOperator,
+    pub setq: SetQuantifier,
+    pub lhs: Box<QuerySetAst>,
+    pub rhs: Box<QuerySetAst>,
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum SetOperator {
+    Union,
+    Except,
+    Intersect,
+}
 
 /// The expressions that can result in values.
 #[derive(Clone, Debug, PartialEq)]
@@ -277,17 +310,13 @@ pub enum ExprKind {
     Bag(BagAst),
     List(ListAst),
     Sexp(SexpAst),
-    /// Set operators
-    Union(UnionAst),
-    Except(ExceptAst),
-    Intersect(IntersectAst),
     /// Other expression types
     Path(PathAst),
     Call(CallAst),
     CallAgg(CallAggAst),
 
-    /// `SELECT` and its parts.
-    Select(SelectAst),
+    /// Query, e.g. `UNION` | `EXCEPT` | `INTERSECT` | `SELECT` and their parts.
+    Query(QueryAst),
 
     /// Indicates an error occurred during query processing; The exact error details are out of band of the AST
     Error,
@@ -454,24 +483,6 @@ pub struct LitTime {
     pub value: TimeValue,
 }
 
-#[derive(Clone, Debug, PartialEq)]
-pub struct Union {
-    pub setq: SetQuantifier,
-    pub operands: Vec<Box<Expr>>,
-}
-
-#[derive(Clone, Debug, PartialEq)]
-pub struct Except {
-    pub setq: SetQuantifier,
-    pub operands: Vec<Box<Expr>>,
-}
-
-#[derive(Clone, Debug, PartialEq)]
-pub struct Intersect {
-    pub setq: SetQuantifier,
-    pub operands: Vec<Box<Expr>>,
-}
-
 #[derive(Clone, Debug, PartialEq)]
 pub struct Path {
     pub root: Box<Expr>,
@@ -529,9 +540,6 @@ pub struct Select {
     pub where_clause: Option<Box<Expr>>,
     pub group_by: Option<Box<GroupByExprAst>>,
     pub having: Option<Box<Expr>>,
-    pub order_by: Option<Box<OrderByExprAst>>,
-    pub limit: Option<Box<Expr>>,
-    pub offset: Option<Box<Expr>>,
 }
 
 #[derive(Clone, Debug, PartialEq)]

partiql-ast/tests/test_ast.rs

@@ -10,9 +10,13 @@ fn test_ast_init() {
 
     let _i = Item {
         kind: ItemKind::Query(Query {
-            expr: Box::new(Expr {
+            set: QuerySet::Expr(Box::new(Expr {
                 kind: ExprKind::Lit(Lit::Int32Lit(23).to_ast(BytePosition::from(1)..12.into())),
-            }),
+            }))
+            .to_ast(BytePosition::from(1)..12.into()),
+            offset: None,
+            order_by: None,
+            limit: None,
         }),
     };
 

partiql-parser/src/parse/mod.rs

@@ -124,8 +124,8 @@ mod tests {
             let res = parse_partiql($q);
             println!("{:#?}", res);
             match res {
-                Ok(_) => (),
-                _ => assert!(false, "{:?}", res),
+                Ok(ast) => ast,
+                _ => panic!("{:?}", res),
             }
         }};
     }
@@ -265,32 +265,32 @@ mod tests {
 
         #[test]
         fn or_simple() {
-            parse!(r#"TRUE OR FALSE"#)
+            parse!(r#"TRUE OR FALSE"#);
         }
 
         #[test]
         fn or() {
-            parse!(r#"t1.super OR test(t2.name, t1.name)"#)
+            parse!(r#"t1.super OR test(t2.name, t1.name)"#);
         }
 
         #[test]
         fn and_simple() {
-            parse!(r#"TRUE and FALSE"#)
+            parse!(r#"TRUE and FALSE"#);
         }
 
         #[test]
         fn and() {
-            parse!(r#"test(t2.name, t1.name) AND t1.id = t2.id"#)
+            parse!(r#"test(t2.name, t1.name) AND t1.id = t2.id"#);
         }
 
         #[test]
         fn or_and() {
-            parse!(r#"t1.super OR test(t2.name, t1.name) AND t1.id = t2.id"#)
+            parse!(r#"t1.super OR test(t2.name, t1.name) AND t1.id = t2.id"#);
         }
 
         #[test]
         fn infix() {
-            parse!(r#"1 + -2 * +3 % 4^5 / 6 - 7  <= 3.14 AND 'foo' || 'bar' LIKE '%oba%'"#)
+            parse!(r#"1 + -2 * +3 % 4^5 / 6 - 7  <= 3.14 AND 'foo' || 'bar' LIKE '%oba%'"#);
         }
     }
 
@@ -299,37 +299,37 @@ mod tests {
 
         #[test]
         fn selectstar() {
-            parse!("SELECT *")
+            parse!("SELECT *");
         }
 
         #[test]
         fn select1() {
-            parse!("SELECT g")
+            parse!("SELECT g");
         }
 
         #[test]
         fn select_list() {
-            parse!("SELECT g, k as ck, h")
+            parse!("SELECT g, k as ck, h");
         }
 
         #[test]
         fn fun_call() {
-            parse!(r#"fun_call('bar', 1,2,3,4,5,'foo')"#)
+            parse!(r#"fun_call('bar', 1,2,3,4,5,'foo')"#);
         }
 
         #[test]
         fn select3() {
-            parse!("SELECT g, k, function('2') as fn_result")
+            parse!("SELECT g, k, function('2') as fn_result");
         }
 
         #[test]
         fn group() {
-            parse!("SELECT g FROM data GROUP BY a")
+            parse!("SELECT g FROM data GROUP BY a");
         }
 
         #[test]
         fn group_complex() {
-            parse!("SELECT g FROM data GROUP BY a AS x, b + c AS y, foo(d) AS z GROUP AS g")
+            parse!("SELECT g FROM data GROUP BY a AS x, b + c AS y, foo(d) AS z GROUP AS g");
         }
 
         #[test]
@@ -341,27 +341,29 @@ mod tests {
 
         #[test]
         fn where_simple() {
-            parse!(r#"SELECT a FROM tb WHERE hk = 1"#)
+            parse!(r#"SELECT a FROM tb WHERE hk = 1"#);
         }
 
         #[test]
         fn where_boolean() {
-            parse!(r#"SELECT a FROM tb WHERE t1.super OR test(t2.name, t1.name) AND t1.id = t2.id"#)
+            parse!(
+                r#"SELECT a FROM tb WHERE t1.super OR test(t2.name, t1.name) AND t1.id = t2.id"#
+            );
         }
 
         #[test]
         fn limit() {
-            parse!(r#"SELECT * FROM a LIMIT 10"#)
+            parse!(r#"SELECT * FROM a LIMIT 10"#);
         }
 
         #[test]
         fn offset() {
-            parse!(r#"SELECT * FROM a OFFSET 10"#)
+            parse!(r#"SELECT * FROM a OFFSET 10"#);
         }
 
         #[test]
         fn limit_offset() {
-            parse!(r#"SELECT * FROM a LIMIT 10 OFFSET 2"#)
+            parse!(r#"SELECT * FROM a LIMIT 10 OFFSET 2"#);
         }
 
         #[test]
@@ -378,7 +380,44 @@ mod tests {
             )
             AS deltas FROM SOURCE_VIEW_DELTA_FULL_TRANSACTIONS delta_full_transactions
             "#;
-            parse!(q)
+            parse!(q);
+        }
+    }
+
+    mod set_ops {
+        use super::*;
+
+        #[test]
+        fn set_ops() {
+            parse!(
+                r#"(SELECT * FROM a LIMIT 10 OFFSET 2) UNION SELECT * FROM b INTERSECT c EXCEPT SELECT * FROM d"#
+            );
+        }
+
+        #[test]
+        fn union_prec() {
+            let l = parse!(r#"a union b union c"#);
+            let r = parse!(r#"(a union b) union c"#);
+            assert_eq!(l, r);
+        }
+
+        #[test]
+        fn intersec_prec() {
+            let l = parse!(r#"a union b intersect c"#);
+            let r = parse!(r#"a union (b intersect c)"#);
+            assert_eq!(l, r);
+        }
+
+        #[test]
+        fn limit() {
+            let l = parse!(r#"SELECT a FROM b UNION SELECT x FROM y ORDER BY a LIMIT 10 OFFSET 5"#);
+            let r =
+                parse!(r#"(SELECT a FROM b UNION SELECT x FROM y) ORDER BY a LIMIT 10 OFFSET 5"#);
+            assert_eq!(l, r);
+            let r2 =
+                parse!(r#"SELECT a FROM b UNION (SELECT x FROM y ORDER BY a LIMIT 10 OFFSET 5)"#);
+            assert_ne!(l, r2);
+            assert_ne!(r, r2);
         }
     }