Add support for quantified comparison predicates (ALL/ANY/SOME) (#1459)

This commit is contained in:
Yoav Cohen 2024-10-09 23:47:14 +02:00 committed by GitHub
parent 0a941b87dc
commit a4fa9e08b7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 64 additions and 22 deletions

View file

@ -646,12 +646,16 @@ pub enum Expr {
regexp: bool, regexp: bool,
}, },
/// `ANY` operation e.g. `foo > ANY(bar)`, comparison operator is one of `[=, >, <, =>, =<, !=]` /// `ANY` operation e.g. `foo > ANY(bar)`, comparison operator is one of `[=, >, <, =>, =<, !=]`
/// <https://docs.snowflake.com/en/sql-reference/operators-subquery#all-any>
AnyOp { AnyOp {
left: Box<Expr>, left: Box<Expr>,
compare_op: BinaryOperator, compare_op: BinaryOperator,
right: Box<Expr>, 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 `[=, >, <, =>, =<, !=]` /// `ALL` operation e.g. `foo > ALL(bar)`, comparison operator is one of `[=, >, <, =>, =<, !=]`
/// <https://docs.snowflake.com/en/sql-reference/operators-subquery#all-any>
AllOp { AllOp {
left: Box<Expr>, left: Box<Expr>,
compare_op: BinaryOperator, compare_op: BinaryOperator,
@ -1332,12 +1336,30 @@ impl fmt::Display for Expr {
left, left,
compare_op, compare_op,
right, 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 { Expr::AllOp {
left, left,
compare_op, compare_op,
right, 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 } => { Expr::UnaryOp { op, expr } => {
if op == &UnaryOperator::PGPostfixFactorial { if op == &UnaryOperator::PGPostfixFactorial {
write!(f, "{expr}{op}") write!(f, "{expr}{op}")

View file

@ -1302,13 +1302,9 @@ impl<'a> Parser<'a> {
} }
fn try_parse_expr_sub_query(&mut self) -> Result<Option<Expr>, ParserError> { fn try_parse_expr_sub_query(&mut self) -> Result<Option<Expr>, ParserError> {
if self if !self.peek_sub_query() {
.parse_one_of_keywords(&[Keyword::SELECT, Keyword::WITH])
.is_none()
{
return Ok(None); return Ok(None);
} }
self.prev_token();
Ok(Some(Expr::Subquery(self.parse_boxed_query()?))) 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 // Snowflake permits a subquery to be passed as an argument without
// an enclosing set of parens if it's the only argument. // an enclosing set of parens if it's the only argument.
if dialect_of!(self is SnowflakeDialect) if dialect_of!(self is SnowflakeDialect) && self.peek_sub_query() {
&& self
.parse_one_of_keywords(&[Keyword::WITH, Keyword::SELECT])
.is_some()
{
self.prev_token();
let subquery = self.parse_boxed_query()?; let subquery = self.parse_boxed_query()?;
self.expect_token(&Token::RParen)?; self.expect_token(&Token::RParen)?;
return Ok(Expr::Function(Function { return Ok(Expr::Function(Function {
@ -2639,10 +2630,21 @@ impl<'a> Parser<'a> {
}; };
if let Some(op) = regular_binary_operator { 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)?; self.expect_token(&Token::LParen)?;
let right = self.parse_subexpr(precedence)?; let right = if self.peek_sub_query() {
self.expect_token(&Token::RParen)?; // 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!( if !matches!(
op, op,
@ -2667,10 +2669,11 @@ impl<'a> Parser<'a> {
compare_op: op, compare_op: op,
right: Box::new(right), right: Box::new(right),
}, },
Keyword::ANY => Expr::AnyOp { Keyword::ANY | Keyword::SOME => Expr::AnyOp {
left: Box::new(expr), left: Box::new(expr),
compare_op: op, compare_op: op,
right: Box::new(right), right: Box::new(right),
is_some: keyword == Keyword::SOME,
}, },
_ => unreachable!(), _ => unreachable!(),
}) })
@ -10507,11 +10510,7 @@ impl<'a> Parser<'a> {
vec![] vec![]
}; };
PivotValueSource::Any(order_by) PivotValueSource::Any(order_by)
} else if self } else if self.peek_sub_query() {
.parse_one_of_keywords(&[Keyword::SELECT, Keyword::WITH])
.is_some()
{
self.prev_token();
PivotValueSource::Subquery(self.parse_query()?) PivotValueSource::Subquery(self.parse_query()?)
} else { } else {
PivotValueSource::List(self.parse_comma_separated(Self::parse_expr_with_alias)?) 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<TokenWithLocation> { pub fn into_tokens(self) -> Vec<TokenWithLocation> {
self.tokens 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 { impl Word {

View file

@ -1907,6 +1907,7 @@ fn parse_binary_any() {
left: Box::new(Expr::Identifier(Ident::new("a"))), left: Box::new(Expr::Identifier(Ident::new("a"))),
compare_op: BinaryOperator::Eq, compare_op: BinaryOperator::Eq,
right: Box::new(Expr::Identifier(Ident::new("b"))), right: Box::new(Expr::Identifier(Ident::new("b"))),
is_some: false,
}), }),
select.projection[0] 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 ILIKE ANY ('%Jo%oe%', 'T%e')"#);
verified_stmt(r#"SELECT * FROM x WHERE a LIKE 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)");
}