Support DISTINCT for SetOperator

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, SetOperatorOption,
+    TableAlias, TableFactor, TableWithJoins, Top, Values, With,
 };
 pub use self::value::{DateTimeField, TrimWhereField, Value};
 

src/ast/query.rs

@@ -64,6 +64,23 @@ impl fmt::Display for Query {
     }
 }
 
+/// Options for SetOperator
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum SetOperatorOption {
+    All,
+    Distinct,
+}
+
+impl fmt::Display for SetOperatorOption {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            SetOperatorOption::All => write!(f, " ALL"),
+            SetOperatorOption::Distinct => write!(f, " DISTINCT"),
+        }
+    }
+}
+
 /// A node in a tree, representing a "query body" expression, roughly:
 /// `SELECT ... [ {UNION|EXCEPT|INTERSECT} SELECT ...]`
 #[allow(clippy::large_enum_variant)]
@@ -78,7 +95,7 @@ pub enum SetExpr {
     /// UNION/EXCEPT/INTERSECT of two queries
     SetOperation {
         op: SetOperator,
-        all: bool,
+        op_option: Option<SetOperatorOption>,
         left: Box<SetExpr>,
         right: Box<SetExpr>,
     },
@@ -98,10 +115,14 @@ impl fmt::Display for SetExpr {
                 left,
                 right,
                 op,
-                all,
+                op_option,
             } => {
-                let all_str = if *all { " ALL" } else { "" };
-                write!(f, "{} {}{} {}", left, op, all_str, right)
+                write!(f, "{} {}", left, op)?;
+                if let Some(ref o) = op_option {
+                    write!(f, "{}", o)?;
+                }
+                write!(f, " {}", right)?;
+                Ok(())
             }
         }
     }

src/parser.rs

@@ -4173,10 +4173,11 @@ impl<'a> Parser<'a> {
                 break;
             }
             self.next_token(); // skip past the set operator
+            let op_option = self.parse_set_operator_option(&op);
             expr = SetExpr::SetOperation {
                 left: Box::new(expr),
                 op: op.unwrap(),
-                all: self.parse_keyword(Keyword::ALL),
+                op_option,
                 right: Box::new(self.parse_query_body(next_precedence)?),
             };
         }
@@ -4193,6 +4194,40 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_set_operator_option(
+        &mut self,
+        op: &Option<SetOperator>,
+    ) -> Option<SetOperatorOption> {
+        match op {
+            Some(SetOperator::Union) => {
+                if self.parse_keyword(Keyword::ALL) {
+                    Some(SetOperatorOption::All)
+                } else if self.parse_keyword(Keyword::DISTINCT)
+                    && dialect_of!(self is MySqlDialect | BigQueryDialect)
+                {
+                    Some(SetOperatorOption::Distinct)
+                } else {
+                    None
+                }
+            }
+            Some(SetOperator::Except) => {
+                if self.parse_keyword(Keyword::ALL) {
+                    Some(SetOperatorOption::All)
+                } else {
+                    None
+                }
+            }
+            Some(SetOperator::Intersect) => {
+                if self.parse_keyword(Keyword::ALL) {
+                    Some(SetOperatorOption::All)
+                } else {
+                    None
+                }
+            }
+            _ => 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

@@ -26,7 +26,7 @@ use sqlparser::ast::SelectItem::UnnamedExpr;
 use sqlparser::ast::*;
 use sqlparser::dialect::{
     AnsiDialect, BigQueryDialect, ClickHouseDialect, GenericDialect, HiveDialect, MsSqlDialect,
-    PostgreSqlDialect, SQLiteDialect, SnowflakeDialect,
+    MySqlDialect, PostgreSqlDialect, SQLiteDialect, SnowflakeDialect,
 };
 use sqlparser::keywords::ALL_KEYWORDS;
 use sqlparser::parser::{Parser, ParserError};
@@ -4096,6 +4096,16 @@ fn parse_union() {
     verified_stmt("(SELECT * FROM new EXCEPT SELECT * FROM old) UNION ALL (SELECT * FROM old EXCEPT SELECT * FROM new) ORDER BY 1");
 }
 
+#[test]
+fn parse_union_distinct() {
+    let dialects = TestedDialects {
+        dialects: vec![Box::new(BigQueryDialect {}), Box::new(MySqlDialect {})],
+    };
+    dialects.verified_stmt("SELECT 1 UNION DISTINCT SELECT 2");
+    dialects.verified_stmt("SELECT 1 UNION DISTINCT SELECT 2 INTERSECT ALL SELECT 3");
+    dialects.verified_stmt("SELECT 1 EXCEPT (SELECT 2 UNION DISTINCT SELECT 3)");
+}
+
 #[test]
 fn parse_values() {
     verified_stmt("SELECT * FROM (VALUES (1), (2), (3))");

tests/sqlparser_postgres.rs

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