Skip to content

Support DISTINCT for SetOperator #689

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Nov 7, 2022
4 changes: 2 additions & 2 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand Down
35 changes: 31 additions & 4 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
},
Expand All @@ -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(())
}
}
}
Expand All @@ -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`
Comment on lines +135 to +137
Copy link
Contributor Author

@unvalley unvalley Nov 4, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added some comments at ae6fb81 , but I wasn't quite sure how to describe this enum🥲

#[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`.
Expand Down
27 changes: 26 additions & 1 deletion src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4173,10 +4173,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)?),
};
}
Expand All @@ -4193,6 +4194,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> {
Expand Down
6 changes: 5 additions & 1 deletion tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4078,14 +4078,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)");
Expand All @@ -4094,6 +4097,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]
Expand Down
2 changes: 1 addition & 1 deletion tests/sqlparser_postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1310,7 +1310,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,
Expand Down