Spark, Databricks: parse the exists higher-order function

fuzz/fuzz_seeds/spark_2772_exists_function.sql

@@ -0,0 +1 @@
+SELECT exists(array(1, 2), x -> x > 1), EXISTS (VALUES (1))

src/dialect/databricks.rs

@@ -90,6 +90,11 @@ impl Dialect for DatabricksDialect {
         true
     }
 
+    /// See <https://docs.databricks.com/aws/en/sql/language-manual/functions/exists>
+    fn supports_exists_function(&self) -> bool {
+        true
+    }
+
     // https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-qry-select.html#syntax
     fn supports_select_wildcard_except(&self) -> bool {
         true

src/dialect/mod.rs

@@ -656,6 +656,13 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if `exists(...)` is a call to a function named `exists`,
+    /// as in `exists(array(1, 2), x -> x > 1)`, unless a query follows the
+    /// parenthesis, as in `EXISTS (SELECT 1)`.
+    fn supports_exists_function(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports the `LAMBDA` keyword spelling of
     /// lambda functions, for example:
     ///

src/dialect/spark.rs

@@ -81,6 +81,11 @@ impl Dialect for SparkSqlDialect {
         true
     }
 
+    /// See <https://spark.apache.org/docs/latest/sql-ref-functions-builtin.html#collection-functions>
+    fn supports_exists_function(&self) -> bool {
+        true
+    }
+
     /// See <https://spark.apache.org/docs/latest/sql-ref-syntax-qry-select.html>
     fn supports_select_wildcard_except(&self) -> bool {
         true

src/parser/mod.rs

@@ -1620,18 +1620,8 @@ impl<'a> Parser<'a> {
             Keyword::TRY_CAST => Ok(Some(self.parse_cast_expr(CastKind::TryCast)?)),
             Keyword::SAFE_CAST => Ok(Some(self.parse_cast_expr(CastKind::SafeCast)?)),
             Keyword::EXISTS
-            // Support parsing Databricks has a function named `exists`.
-            if !dialect_of!(self is DatabricksDialect)
-                || matches!(
-                        self.peek_nth_token_ref(1).token,
-                        Token::Word(Word {
-                            keyword: Keyword::SELECT | Keyword::WITH,
-                            ..
-                        })
-                    ) =>
-                {
-                    Ok(Some(self.parse_exists_expr(false)?))
-                }
+            // `None` lets a dialect's `exists` function parse as a function call.
+            => self.maybe_parse_exists_expr(false),
             Keyword::EXTRACT => Ok(Some(self.parse_extract_expr()?)),
             Keyword::CEIL => Ok(Some(self.parse_ceil_floor_expr(true)?)),
             Keyword::FLOOR => Ok(Some(self.parse_ceil_floor_expr(false)?)),
@@ -3024,6 +3014,16 @@ impl<'a> Parser<'a> {
         Ok(exists_node)
     }
 
+    /// Parse `EXISTS (<query>)` after the `EXISTS` keyword, or return `None` when
+    /// the dialect reads the parentheses as the arguments of an `exists` function.
+    fn maybe_parse_exists_expr(&mut self, negated: bool) -> Result<Option<Expr>, ParserError> {
+        if self.dialect.supports_exists_function() {
+            self.maybe_parse(|parser| parser.parse_exists_expr(negated))
+        } else {
+            self.parse_exists_expr(negated).map(Some)
+        }
+    }
+
     /// Parse a SQL `EXTRACT` expression e.g. `EXTRACT(YEAR FROM date)`.
     pub fn parse_extract_expr(&mut self) -> Result<Expr, ParserError> {
         self.expect_token(&Token::LParen)?;
@@ -3382,25 +3382,16 @@ impl<'a> Parser<'a> {
     ///
     /// Represented in the AST as `Expr::UnaryOp` with `UnaryOperator::Not`.
     pub fn parse_not(&mut self) -> Result<Expr, ParserError> {
-        match &self.peek_token_ref().token {
-            Token::Word(w) => match w.keyword {
-                Keyword::EXISTS => {
-                    let negated = true;
-                    let _ = self.parse_keyword(Keyword::EXISTS);
-                    self.parse_exists_expr(negated)
-                }
-                _ => Ok(Expr::UnaryOp {
-                    op: UnaryOperator::Not,
-                    expr: Box::new(
-                        self.parse_subexpr(self.dialect.prec_value(Precedence::UnaryNot))?,
-                    ),
-                }),
-            },
-            _ => Ok(Expr::UnaryOp {
-                op: UnaryOperator::Not,
-                expr: Box::new(self.parse_subexpr(self.dialect.prec_value(Precedence::UnaryNot))?),
-            }),
+        if self.parse_keyword(Keyword::EXISTS) {
+            if let Some(expr) = self.maybe_parse_exists_expr(true)? {
+                return Ok(expr);
+            }
+            self.prev_token();
         }
+        Ok(Expr::UnaryOp {
+            op: UnaryOperator::Not,
+            expr: Box::new(self.parse_subexpr(self.dialect.prec_value(Precedence::UnaryNot))?),
+        })
     }
 
     /// Parse expression types that start with a left brace '{'.

tests/sqlparser_common.rs

@@ -9058,7 +9058,7 @@ fn parse_exists_subquery() {
     verified_stmt("SELECT * FROM t WHERE EXISTS (WITH u AS (SELECT 1) SELECT * FROM u)");
     verified_stmt("SELECT EXISTS (SELECT 1)");
 
-    let res = all_dialects_except(|d| d.is::<DatabricksDialect>())
+    let res = all_dialects_except(|d| d.supports_exists_function())
         .parse_sql_statements("SELECT EXISTS (");
     assert_eq!(
         ParserError::ParserError(
@@ -9067,7 +9067,7 @@ fn parse_exists_subquery() {
         res.unwrap_err(),
     );
 
-    let res = all_dialects_except(|d| d.is::<DatabricksDialect>())
+    let res = all_dialects_except(|d| d.supports_exists_function())
         .parse_sql_statements("SELECT EXISTS (NULL)");
     assert_eq!(
         ParserError::ParserError(
@@ -9077,6 +9077,53 @@ fn parse_exists_subquery() {
     );
 }
 
+#[test]
+fn parse_exists_function() {
+    let dialects = all_dialects_where(|d| d.supports_exists_function());
+    match dialects.verified_expr("exists(array(1, 2), x -> x > 1)") {
+        Expr::Function(f) => assert_eq!(f.name.to_string(), "exists"),
+        other => panic!("expected a function, got {other:?}"),
+    }
+    for sql in [
+        "EXISTS (SELECT 1)",
+        "EXISTS (WITH t AS (SELECT 1) SELECT * FROM t)",
+        "EXISTS (VALUES (1))",
+        "EXISTS (TABLE t)",
+        "EXISTS ((SELECT 1))",
+        "EXISTS ((SELECT 1) UNION (SELECT 2))",
+    ] {
+        assert!(
+            matches!(dialects.verified_expr(sql), Expr::Exists { .. }),
+            "{sql}"
+        );
+    }
+    assert!(matches!(
+        dialects.verified_expr("NOT EXISTS (SELECT 1)"),
+        Expr::Exists { negated: true, .. }
+    ));
+    match dialects.verified_expr("NOT exists(array(1, 2), x -> x > 1)") {
+        Expr::UnaryOp {
+            op: UnaryOperator::Not,
+            expr,
+        } => assert!(matches!(*expr, Expr::Function(_))),
+        other => panic!("expected NOT of a function, got {other:?}"),
+    }
+    let dialects =
+        all_dialects_where(|d| d.supports_exists_function() && d.supports_from_first_select());
+    assert!(matches!(
+        dialects.verified_expr("EXISTS (FROM t SELECT a)"),
+        Expr::Exists { .. }
+    ));
+
+    assert_eq!(
+        all_dialects_except(|d| d.supports_exists_function())
+            .parse_sql_statements("SELECT exists(array(1, 2), x -> x > 1)"),
+        Err(ParserError::ParserError(
+            "Expected: SELECT, VALUES, or a subquery in the query body, found: array".to_string()
+        ))
+    );
+}
+
 #[test]
 fn parse_create_database() {
     let sql = "CREATE DATABASE mydb";

tests/sqlparser_spark.rs

@@ -157,6 +157,12 @@ fn test_lambda_transform() {
     spark().verified_stmt("SELECT transform(arr, x -> x * 2) FROM t");
 }
 
+#[test]
+fn test_exists_function() {
+    spark().verified_stmt("SELECT exists(arr, x -> x > 2), NOT exists(arr, x -> x < 0) FROM t");
+    spark().verified_stmt("SELECT * FROM t WHERE EXISTS (SELECT 1) AND NOT EXISTS ((SELECT 2))");
+}
+
 // --------------------------------
 // DIV integer division
 // --------------------------------