Add support for INCLUDE/EXCLUDE NULLS for UNPIVOT (#10)

Author: Vedin

src/ast/query.rs

@@ -1248,7 +1248,7 @@ pub enum TableFactor {
     ///
     /// Syntax:
     /// ```sql
-    /// table UNPIVOT(value FOR name IN (column1, [ column2, ... ])) [ alias ]
+    /// table UNPIVOT [ { INCLUDE | EXCLUDE } NULLS ] (value FOR name IN (column1, [ column2, ... ])) [ alias ]
     /// ```
     ///
     /// See <https://docs.snowflake.com/en/sql-reference/constructs/unpivot>.
@@ -1257,6 +1257,7 @@ pub enum TableFactor {
         value: Ident,
         name: Ident,
         columns: Vec<Ident>,
+        include_nulls: Option<bool>,
         alias: Option<TableAlias>,
     },
     /// A `MATCH_RECOGNIZE` operation on a table.
@@ -1893,15 +1894,23 @@ impl fmt::Display for TableFactor {
             }
             TableFactor::Unpivot {
                 table,
+                include_nulls,
                 value,
                 name,
                 columns,
                 alias,
             } => {
+                write!(f, "{table} UNPIVOT")?;
+                if let Some(include_nulls) = include_nulls {
+                    if *include_nulls {
+                        write!(f, " INCLUDE NULLS ")?;
+                    } else {
+                        write!(f, " EXCLUDE NULLS ")?;
+                    }
+                }
                 write!(
                     f,
-                    "{} UNPIVOT({} FOR {} IN ({}))",
-                    table,
+                    "({} FOR {} IN ({}))",
                     value,
                     name,
                     display_comma_separated(columns)

src/ast/spans.rs

@@ -1830,6 +1830,7 @@ impl Spanned for TableFactor {
             TableFactor::Unpivot {
                 table,
                 value,
+                include_nulls: _,
                 name,
                 columns,
                 alias,

src/parser/mod.rs

@@ -12396,6 +12396,15 @@ impl<'a> Parser<'a> {
         &mut self,
         table: TableFactor,
     ) -> Result<TableFactor, ParserError> {
+        let include_nulls = if self.parse_keyword(Keyword::INCLUDE) {
+            self.expect_keyword_is(Keyword::NULLS)?;
+            Some(true)
+        } else if self.parse_keyword(Keyword::EXCLUDE) {
+            self.expect_keyword_is(Keyword::NULLS)?;
+            Some(false)
+        } else {
+            None
+        };
         self.expect_token(&Token::LParen)?;
         let value = self.parse_identifier()?;
         self.expect_keyword_is(Keyword::FOR)?;
@@ -12407,6 +12416,7 @@ impl<'a> Parser<'a> {
         Ok(TableFactor::Unpivot {
             table: Box::new(table),
             value,
+            include_nulls,
             name,
             columns,
             alias,

tests/sqlparser_common.rs

@@ -10568,49 +10568,47 @@ fn parse_unpivot_table() {
         "SELECT * FROM sales AS s ",
         "UNPIVOT(quantity FOR quarter IN (Q1, Q2, Q3, Q4)) AS u (product, quarter, quantity)"
     );
-
-    pretty_assertions::assert_eq!(
-        verified_only_select(sql).from[0].relation,
-        Unpivot {
-            table: Box::new(TableFactor::Table {
-                name: ObjectName::from(vec![Ident::new("sales")]),
-                alias: Some(TableAlias {
-                    name: Ident::new("s"),
-                    columns: vec![]
-                }),
-                args: None,
-                with_hints: vec![],
-                version: None,
-                partitions: vec![],
-                with_ordinality: false,
-                json_path: None,
-                sample: None,
-                index_hints: vec![],
+    let base_unpivot = Unpivot {
+        table: Box::new(TableFactor::Table {
+            name: ObjectName::from(vec![Ident::new("sales")]),
+            alias: Some(TableAlias {
+                name: Ident::new("s"),
+                columns: vec![],
             }),
-            value: Ident {
-                value: "quantity".to_string(),
-                quote_style: None,
-                span: Span::empty()
-            },
+            args: None,
+            with_hints: vec![],
+            version: None,
+            partitions: vec![],
+            with_ordinality: false,
+            json_path: None,
+            sample: None,
+            index_hints: vec![],
+        }),
+        include_nulls: None,
+        value: Ident {
+            value: "quantity".to_string(),
+            quote_style: None,
+            span: Span::empty(),
+        },
 
-            name: Ident {
-                value: "quarter".to_string(),
-                quote_style: None,
-                span: Span::empty()
-            },
-            columns: ["Q1", "Q2", "Q3", "Q4"]
+        name: Ident {
+            value: "quarter".to_string(),
+            quote_style: None,
+            span: Span::empty(),
+        },
+        columns: ["Q1", "Q2", "Q3", "Q4"]
+            .into_iter()
+            .map(Ident::new)
+            .collect(),
+        alias: Some(TableAlias {
+            name: Ident::new("u"),
+            columns: ["product", "quarter", "quantity"]
                 .into_iter()
-                .map(Ident::new)
+                .map(TableAliasColumnDef::from_name)
                 .collect(),
-            alias: Some(TableAlias {
-                name: Ident::new("u"),
-                columns: ["product", "quarter", "quantity"]
-                    .into_iter()
-                    .map(TableAliasColumnDef::from_name)
-                    .collect(),
-            }),
-        }
-    );
+        }),
+    };
+    pretty_assertions::assert_eq!(verified_only_select(sql).from[0].relation, base_unpivot);
     assert_eq!(verified_stmt(sql).to_string(), sql);
 
     let sql_without_aliases = concat!(
@@ -10630,6 +10628,38 @@ fn parse_unpivot_table() {
         verified_stmt(sql_without_aliases).to_string(),
         sql_without_aliases
     );
+
+    let sql_unpivot_exclude_nulls = concat!(
+    "SELECT * FROM sales AS s ",
+    "UNPIVOT EXCLUDE NULLS (quantity FOR quarter IN (Q1, Q2, Q3, Q4)) AS u (product, quarter, quantity)"
+    );
+
+    if let Unpivot { include_nulls, .. } =
+        &verified_only_select(sql_unpivot_exclude_nulls).from[0].relation
+    {
+        assert_eq!(*include_nulls, Some(false));
+    }
+
+    assert_eq!(
+        verified_stmt(sql_unpivot_exclude_nulls).to_string(),
+        sql_unpivot_exclude_nulls
+    );
+
+    let sql_unpivot_include_nulls = concat!(
+        "SELECT * FROM sales AS s ",
+        "UNPIVOT INCLUDE NULLS (quantity FOR quarter IN (Q1, Q2, Q3, Q4)) AS u (product, quarter, quantity)"
+    );
+
+    if let Unpivot { include_nulls, .. } =
+        &verified_only_select(sql_unpivot_include_nulls).from[0].relation
+    {
+        assert_eq!(*include_nulls, Some(true));
+    }
+
+    assert_eq!(
+        verified_stmt(sql_unpivot_include_nulls).to_string(),
+        sql_unpivot_include_nulls
+    );
 }
 
 #[test]
@@ -10726,6 +10756,7 @@ fn parse_pivot_unpivot_table() {
                     sample: None,
                     index_hints: vec![],
                 }),
+                include_nulls: None,
                 value: Ident {
                     value: "population".to_string(),
                     quote_style: None,