Support DISTINCT AS { STRUCT | VALUE } for BigQuery (apache#1880)

Author: bombsimon

src/ast/query.rs

@@ -3338,22 +3338,28 @@ impl fmt::Display for OpenJsonTableColumn {
 }
 
 /// BigQuery supports ValueTables which have 2 modes:
-/// `SELECT AS STRUCT`
-/// `SELECT AS VALUE`
+/// `SELECT [ALL | DISTINCT] AS STRUCT`
+/// `SELECT [ALL | DISTINCT] AS VALUE`
+///
 /// <https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#value_tables>
+/// <https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#select_list>
 #[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum ValueTableMode {
     AsStruct,
     AsValue,
+    DistinctAsStruct,
+    DistinctAsValue,
 }
 
 impl fmt::Display for ValueTableMode {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             ValueTableMode::AsStruct => write!(f, "AS STRUCT"),
             ValueTableMode::AsValue => write!(f, "AS VALUE"),
+            ValueTableMode::DistinctAsStruct => write!(f, "DISTINCT AS STRUCT"),
+            ValueTableMode::DistinctAsValue => write!(f, "DISTINCT AS VALUE"),
         }
     }
 }

src/parser/mod.rs

@@ -11514,18 +11514,7 @@ impl<'a> Parser<'a> {
         }
 
         let select_token = self.expect_keyword(Keyword::SELECT)?;
-        let value_table_mode =
-            if dialect_of!(self is BigQueryDialect) && self.parse_keyword(Keyword::AS) {
-                if self.parse_keyword(Keyword::VALUE) {
-                    Some(ValueTableMode::AsValue)
-                } else if self.parse_keyword(Keyword::STRUCT) {
-                    Some(ValueTableMode::AsStruct)
-                } else {
-                    self.expected("VALUE or STRUCT", self.peek_token())?
-                }
-            } else {
-                None
-            };
+        let value_table_mode = self.parse_value_table_mode()?;
 
         let mut top_before_distinct = false;
         let mut top = None;
@@ -11701,6 +11690,32 @@ impl<'a> Parser<'a> {
         })
     }
 
+    fn parse_value_table_mode(&mut self) -> Result<Option<ValueTableMode>, ParserError> {
+        if !dialect_of!(self is BigQueryDialect) {
+            return Ok(None);
+        }
+
+        let mode = if self.parse_keywords(&[Keyword::DISTINCT, Keyword::AS, Keyword::VALUE]) {
+            Some(ValueTableMode::DistinctAsValue)
+        } else if self.parse_keywords(&[Keyword::DISTINCT, Keyword::AS, Keyword::STRUCT]) {
+            Some(ValueTableMode::DistinctAsStruct)
+        } else if self.parse_keywords(&[Keyword::AS, Keyword::VALUE])
+            || self.parse_keywords(&[Keyword::ALL, Keyword::AS, Keyword::VALUE])
+        {
+            Some(ValueTableMode::AsValue)
+        } else if self.parse_keywords(&[Keyword::AS, Keyword::STRUCT])
+            || self.parse_keywords(&[Keyword::ALL, Keyword::AS, Keyword::STRUCT])
+        {
+            Some(ValueTableMode::AsStruct)
+        } else if self.parse_keyword(Keyword::AS) {
+            self.expected("VALUE or STRUCT", self.peek_token())?
+        } else {
+            None
+        };
+
+        Ok(mode)
+    }
+
     /// Invoke `f` after first setting the parser's `ParserState` to `state`.
     ///
     /// Upon return, restores the parser's state to what it started at.

tests/sqlparser_bigquery.rs

@@ -2313,16 +2313,46 @@ fn bigquery_select_expr_star() {
 
 #[test]
 fn test_select_as_struct() {
-    bigquery().verified_only_select("SELECT * FROM (SELECT AS VALUE STRUCT(123 AS a, false AS b))");
+    for (sql, parse_to) in [
+        (
+            "SELECT * FROM (SELECT AS STRUCT STRUCT(123 AS a, false AS b))",
+            "SELECT * FROM (SELECT AS STRUCT STRUCT(123 AS a, false AS b))",
+        ),
+        (
+            "SELECT * FROM (SELECT DISTINCT AS STRUCT STRUCT(123 AS a, false AS b))",
+            "SELECT * FROM (SELECT DISTINCT AS STRUCT STRUCT(123 AS a, false AS b))",
+        ),
+        (
+            "SELECT * FROM (SELECT ALL AS STRUCT STRUCT(123 AS a, false AS b))",
+            "SELECT * FROM (SELECT AS STRUCT STRUCT(123 AS a, false AS b))",
+        ),
+    ] {
+        bigquery().one_statement_parses_to(sql, parse_to);
+    }
+
     let select = bigquery().verified_only_select("SELECT AS STRUCT 1 AS a, 2 AS b");
     assert_eq!(Some(ValueTableMode::AsStruct), select.value_table_mode);
 }
 
 #[test]
 fn test_select_as_value() {
-    bigquery().verified_only_select(
-        "SELECT * FROM (SELECT AS VALUE STRUCT(5 AS star_rating, false AS up_down_rating))",
-    );
+    for (sql, parse_to) in [
+        (
+            "SELECT * FROM (SELECT AS VALUE STRUCT(5 AS star_rating, false AS up_down_rating))",
+            "SELECT * FROM (SELECT AS VALUE STRUCT(5 AS star_rating, false AS up_down_rating))",
+        ),
+        (
+            "SELECT * FROM (SELECT DISTINCT AS VALUE STRUCT(5 AS star_rating, false AS up_down_rating))",
+            "SELECT * FROM (SELECT DISTINCT AS VALUE STRUCT(5 AS star_rating, false AS up_down_rating))",
+        ),
+        (
+            "SELECT * FROM (SELECT ALL AS VALUE STRUCT(5 AS star_rating, false AS up_down_rating))",
+            "SELECT * FROM (SELECT AS VALUE STRUCT(5 AS star_rating, false AS up_down_rating))",
+        ),
+    ] {
+        bigquery().one_statement_parses_to(sql, parse_to);
+    }
+
     let select = bigquery().verified_only_select("SELECT AS VALUE STRUCT(1 AS a, 2 AS b) AS xyz");
     assert_eq!(Some(ValueTableMode::AsValue), select.value_table_mode);
 }