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,12 +1620,16 @@ 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)
+            // A query after the parenthesis is still the EXISTS predicate.
+            if !self.dialect.supports_exists_function()
                 || matches!(
                         self.peek_nth_token_ref(1).token,
                         Token::Word(Word {
-                            keyword: Keyword::SELECT | Keyword::WITH,
+                            keyword: Keyword::SELECT
+                                | Keyword::WITH
+                                | Keyword::TABLE
+                                | Keyword::VALUES
+                                | Keyword::FROM,
                             ..
                         })
                     ) =>

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,40 @@ 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)",
+    ] {
+        assert!(
+            matches!(dialects.verified_expr(sql), Expr::Exists { .. }),
+            "{sql}"
+        );
+    }
+    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";