Support DISTINCT for SetOperator (#689)

* Update SetOperation field all to op_option

* Implement parse_set_operator_option

cargo fmt

Fix parse_set_operator_option after next_token

* Add test for parsing union distinct

* Rename to SetQualifier and fix fmt space

* Add None to SetQualifier

* Update parse method

* Rename to SetQuantifier

* Rename parse_set_operator_option parse_set_operator

* Fix test to parse union, except, intersect

* Add some comments to SetQuantifier

* Fix comment

Author: unvalley

src/ast/mod.rs

@@ -32,8 +32,8 @@ pub use self::ddl::{
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     Cte, Fetch, Join, JoinConstraint, JoinOperator, LateralView, LockType, Offset, OffsetRows,
-    OrderByExpr, Query, Select, SelectInto, SelectItem, SetExpr, SetOperator, TableAlias,
-    TableFactor, TableWithJoins, Top, Values, With,
+    OrderByExpr, Query, Select, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier,
+    TableAlias, TableFactor, TableWithJoins, Top, Values, With,
 };
 pub use self::value::{DateTimeField, TrimWhereField, Value};
 

src/ast/query.rs

@@ -78,7 +78,7 @@ pub enum SetExpr {
     /// UNION/EXCEPT/INTERSECT of two queries
     SetOperation {
         op: SetOperator,
-        all: bool,
+        set_quantifier: SetQuantifier,
         left: Box<SetExpr>,
         right: Box<SetExpr>,
     },
@@ -98,10 +98,17 @@ impl fmt::Display for SetExpr {
                 left,
                 right,
                 op,
-                all,
+                set_quantifier,
             } => {
-                let all_str = if *all { " ALL" } else { "" };
-                write!(f, "{} {}{} {}", left, op, all_str, right)
+                write!(f, "{} {}", left, op)?;
+                match set_quantifier {
+                    SetQuantifier::All | SetQuantifier::Distinct => {
+                        write!(f, " {}", set_quantifier)?
+                    }
+                    SetQuantifier::None => write!(f, "{}", set_quantifier)?,
+                }
+                write!(f, " {}", right)?;
+                Ok(())
             }
         }
     }
@@ -125,6 +132,26 @@ impl fmt::Display for SetOperator {
     }
 }
 
+/// A quantifier for [SetOperator].
+// TODO: Restrict parsing specific SetQuantifier in some specific dialects.
+// For example, BigQuery does not support `DISTINCT` for `EXCEPT` and `INTERSECT`
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum SetQuantifier {
+    All,
+    Distinct,
+    None,
+}
+
+impl fmt::Display for SetQuantifier {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            SetQuantifier::All => write!(f, "ALL"),
+            SetQuantifier::Distinct => write!(f, "DISTINCT"),
+            SetQuantifier::None => write!(f, ""),
+        }
+    }
+}
 /// A restricted variant of `SELECT` (without CTEs/`ORDER BY`), which may
 /// appear either as the only body item of a `Query`, or as an operand
 /// to a set operation like `UNION`.

src/parser.rs

@@ -4199,10 +4199,11 @@ impl<'a> Parser<'a> {
                 break;
             }
             self.next_token(); // skip past the set operator
+            let set_quantifier = self.parse_set_quantifier(&op);
             expr = SetExpr::SetOperation {
                 left: Box::new(expr),
                 op: op.unwrap(),
-                all: self.parse_keyword(Keyword::ALL),
+                set_quantifier,
                 right: Box::new(self.parse_query_body(next_precedence)?),
             };
         }
@@ -4219,6 +4220,30 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_set_quantifier(&mut self, op: &Option<SetOperator>) -> SetQuantifier {
+        match op {
+            Some(SetOperator::Union) => {
+                if self.parse_keyword(Keyword::ALL) {
+                    SetQuantifier::All
+                } else if self.parse_keyword(Keyword::DISTINCT) {
+                    SetQuantifier::Distinct
+                } else {
+                    SetQuantifier::None
+                }
+            }
+            Some(SetOperator::Except) | Some(SetOperator::Intersect) => {
+                if self.parse_keyword(Keyword::ALL) {
+                    SetQuantifier::All
+                } else if self.parse_keyword(Keyword::DISTINCT) {
+                    SetQuantifier::Distinct
+                } else {
+                    SetQuantifier::None
+                }
+            }
+            _ => SetQuantifier::None,
+        }
+    }
+
     /// Parse a restricted `SELECT` statement (no CTEs / `UNION` / `ORDER BY`),
     /// assuming the initial `SELECT` was already consumed
     pub fn parse_select(&mut self) -> Result<Select, ParserError> {

tests/sqlparser_common.rs

@@ -4072,14 +4072,17 @@ fn parse_derived_tables() {
 }
 
 #[test]
-fn parse_union() {
+fn parse_union_except_intersect() {
     // TODO: add assertions
     verified_stmt("SELECT 1 UNION SELECT 2");
     verified_stmt("SELECT 1 UNION ALL SELECT 2");
+    verified_stmt("SELECT 1 UNION DISTINCT SELECT 1");
     verified_stmt("SELECT 1 EXCEPT SELECT 2");
     verified_stmt("SELECT 1 EXCEPT ALL SELECT 2");
+    verified_stmt("SELECT 1 EXCEPT DISTINCT SELECT 1");
     verified_stmt("SELECT 1 INTERSECT SELECT 2");
     verified_stmt("SELECT 1 INTERSECT ALL SELECT 2");
+    verified_stmt("SELECT 1 INTERSECT DISTINCT SELECT 1");
     verified_stmt("SELECT 1 UNION SELECT 2 UNION SELECT 3");
     verified_stmt("SELECT 1 EXCEPT SELECT 2 UNION SELECT 3"); // Union[Except[1,2], 3]
     verified_stmt("SELECT 1 INTERSECT (SELECT 2 EXCEPT SELECT 3)");
@@ -4088,6 +4091,7 @@ fn parse_union() {
     verified_stmt("SELECT 1 UNION SELECT 2 INTERSECT SELECT 3"); // Union[1, Intersect[2,3]]
     verified_stmt("SELECT foo FROM tab UNION SELECT bar FROM TAB");
     verified_stmt("(SELECT * FROM new EXCEPT SELECT * FROM old) UNION ALL (SELECT * FROM old EXCEPT SELECT * FROM new) ORDER BY 1");
+    verified_stmt("(SELECT * FROM new EXCEPT DISTINCT SELECT * FROM old) UNION DISTINCT (SELECT * FROM old EXCEPT DISTINCT SELECT * FROM new) ORDER BY 1");
 }
 
 #[test]

tests/sqlparser_postgres.rs

@@ -1353,7 +1353,7 @@ fn parse_array_subquery_expr() {
             with: None,
             body: Box::new(SetExpr::SetOperation {
                 op: SetOperator::Union,
-                all: false,
+                set_quantifier: SetQuantifier::None,
                 left: Box::new(SetExpr::Select(Box::new(Select {
                     distinct: false,
                     top: None,