Skip to content

Commit

Permalink
Add support for quantified comparison predicates (ALL/ANY/SOME) (#1459)
Browse files Browse the repository at this point in the history
  • Loading branch information
yoavcloud authored Oct 9, 2024
1 parent 0a941b8 commit a4fa9e0
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 22 deletions.
26 changes: 24 additions & 2 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -646,12 +646,16 @@ pub enum Expr {
regexp: bool,
},
/// `ANY` operation e.g. `foo > ANY(bar)`, comparison operator is one of `[=, >, <, =>, =<, !=]`
/// <https://docs.snowflake.com/en/sql-reference/operators-subquery#all-any>
AnyOp {
left: Box<Expr>,
compare_op: BinaryOperator,
right: Box<Expr>,
// ANY and SOME are synonymous: https://docs.cloudera.com/cdw-runtime/cloud/using-hiveql/topics/hive_comparison_predicates.html
is_some: bool,
},
/// `ALL` operation e.g. `foo > ALL(bar)`, comparison operator is one of `[=, >, <, =>, =<, !=]`
/// <https://docs.snowflake.com/en/sql-reference/operators-subquery#all-any>
AllOp {
left: Box<Expr>,
compare_op: BinaryOperator,
Expand Down Expand Up @@ -1332,12 +1336,30 @@ impl fmt::Display for Expr {
left,
compare_op,
right,
} => write!(f, "{left} {compare_op} ANY({right})"),
is_some,
} => {
let add_parens = !matches!(right.as_ref(), Expr::Subquery(_));
write!(
f,
"{left} {compare_op} {}{}{right}{}",
if *is_some { "SOME" } else { "ANY" },
if add_parens { "(" } else { "" },
if add_parens { ")" } else { "" },
)
}
Expr::AllOp {
left,
compare_op,
right,
} => write!(f, "{left} {compare_op} ALL({right})"),
} => {
let add_parens = !matches!(right.as_ref(), Expr::Subquery(_));
write!(
f,
"{left} {compare_op} ALL{}{right}{}",
if add_parens { "(" } else { "" },
if add_parens { ")" } else { "" },
)
}
Expr::UnaryOp { op, expr } => {
if op == &UnaryOperator::PGPostfixFactorial {
write!(f, "{expr}{op}")
Expand Down
51 changes: 31 additions & 20 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1302,13 +1302,9 @@ impl<'a> Parser<'a> {
}

fn try_parse_expr_sub_query(&mut self) -> Result<Option<Expr>, ParserError> {
if self
.parse_one_of_keywords(&[Keyword::SELECT, Keyword::WITH])
.is_none()
{
if !self.peek_sub_query() {
return Ok(None);
}
self.prev_token();

Ok(Some(Expr::Subquery(self.parse_boxed_query()?)))
}
Expand All @@ -1334,12 +1330,7 @@ impl<'a> Parser<'a> {

// Snowflake permits a subquery to be passed as an argument without
// an enclosing set of parens if it's the only argument.
if dialect_of!(self is SnowflakeDialect)
&& self
.parse_one_of_keywords(&[Keyword::WITH, Keyword::SELECT])
.is_some()
{
self.prev_token();
if dialect_of!(self is SnowflakeDialect) && self.peek_sub_query() {
let subquery = self.parse_boxed_query()?;
self.expect_token(&Token::RParen)?;
return Ok(Expr::Function(Function {
Expand Down Expand Up @@ -2639,10 +2630,21 @@ impl<'a> Parser<'a> {
};

if let Some(op) = regular_binary_operator {
if let Some(keyword) = self.parse_one_of_keywords(&[Keyword::ANY, Keyword::ALL]) {
if let Some(keyword) =
self.parse_one_of_keywords(&[Keyword::ANY, Keyword::ALL, Keyword::SOME])
{
self.expect_token(&Token::LParen)?;
let right = self.parse_subexpr(precedence)?;
self.expect_token(&Token::RParen)?;
let right = if self.peek_sub_query() {
// We have a subquery ahead (SELECT\WITH ...) need to rewind and
// use the parenthesis for parsing the subquery as an expression.
self.prev_token(); // LParen
self.parse_subexpr(precedence)?
} else {
// Non-subquery expression
let right = self.parse_subexpr(precedence)?;
self.expect_token(&Token::RParen)?;
right
};

if !matches!(
op,
Expand All @@ -2667,10 +2669,11 @@ impl<'a> Parser<'a> {
compare_op: op,
right: Box::new(right),
},
Keyword::ANY => Expr::AnyOp {
Keyword::ANY | Keyword::SOME => Expr::AnyOp {
left: Box::new(expr),
compare_op: op,
right: Box::new(right),
is_some: keyword == Keyword::SOME,
},
_ => unreachable!(),
})
Expand Down Expand Up @@ -10507,11 +10510,7 @@ impl<'a> Parser<'a> {
vec![]
};
PivotValueSource::Any(order_by)
} else if self
.parse_one_of_keywords(&[Keyword::SELECT, Keyword::WITH])
.is_some()
{
self.prev_token();
} else if self.peek_sub_query() {
PivotValueSource::Subquery(self.parse_query()?)
} else {
PivotValueSource::List(self.parse_comma_separated(Self::parse_expr_with_alias)?)
Expand Down Expand Up @@ -12177,6 +12176,18 @@ impl<'a> Parser<'a> {
pub fn into_tokens(self) -> Vec<TokenWithLocation> {
self.tokens
}

/// Returns true if the next keyword indicates a sub query, i.e. SELECT or WITH
fn peek_sub_query(&mut self) -> bool {
if self
.parse_one_of_keywords(&[Keyword::SELECT, Keyword::WITH])
.is_some()
{
self.prev_token();
return true;
}
false
}
}

impl Word {
Expand Down
9 changes: 9 additions & 0 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1907,6 +1907,7 @@ fn parse_binary_any() {
left: Box::new(Expr::Identifier(Ident::new("a"))),
compare_op: BinaryOperator::Eq,
right: Box::new(Expr::Identifier(Ident::new("b"))),
is_some: false,
}),
select.projection[0]
);
Expand Down Expand Up @@ -11395,3 +11396,11 @@ fn test_select_where_with_like_or_ilike_any() {
verified_stmt(r#"SELECT * FROM x WHERE a ILIKE ANY ('%Jo%oe%', 'T%e')"#);
verified_stmt(r#"SELECT * FROM x WHERE a LIKE ANY ('%Jo%oe%', 'T%e')"#);
}

#[test]
fn test_any_some_all_comparison() {
verified_stmt("SELECT c1 FROM tbl WHERE c1 = ANY(SELECT c2 FROM tbl)");
verified_stmt("SELECT c1 FROM tbl WHERE c1 >= ALL(SELECT c2 FROM tbl)");
verified_stmt("SELECT c1 FROM tbl WHERE c1 <> SOME(SELECT c2 FROM tbl)");
verified_stmt("SELECT 1 = ANY(WITH x AS (SELECT 1) SELECT * FROM x)");
}

0 comments on commit a4fa9e0

Please sign in to comment.