From a4fa9e08b7a73c2d6efc360f9d481352e3521b7f Mon Sep 17 00:00:00 2001 From: Yoav Cohen <59807311+yoavcloud@users.noreply.github.com> Date: Wed, 9 Oct 2024 23:47:14 +0200 Subject: [PATCH] Add support for quantified comparison predicates (ALL/ANY/SOME) (#1459) --- src/ast/mod.rs | 26 ++++++++++++++++++-- src/parser/mod.rs | 51 ++++++++++++++++++++++++--------------- tests/sqlparser_common.rs | 9 +++++++ 3 files changed, 64 insertions(+), 22 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 39f9def8e..2fbe91afc 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -646,12 +646,16 @@ pub enum Expr { regexp: bool, }, /// `ANY` operation e.g. `foo > ANY(bar)`, comparison operator is one of `[=, >, <, =>, =<, !=]` + /// AnyOp { left: Box, compare_op: BinaryOperator, right: Box, + // 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 `[=, >, <, =>, =<, !=]` + /// AllOp { left: Box, compare_op: BinaryOperator, @@ -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}") diff --git a/src/parser/mod.rs b/src/parser/mod.rs index c9e66e1f3..cd9be1d8f 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -1302,13 +1302,9 @@ impl<'a> Parser<'a> { } fn try_parse_expr_sub_query(&mut self) -> Result, 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()?))) } @@ -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 { @@ -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, @@ -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!(), }) @@ -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)?) @@ -12177,6 +12176,18 @@ impl<'a> Parser<'a> { pub fn into_tokens(self) -> Vec { 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 { diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 15132903c..5327880a4 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -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] ); @@ -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)"); +}