Add support for INCLUDE/EXCLUDE NULLS for UNPIVOT (apache#1849)

Author: Vedin

src/ast/mod.rs

@@ -8701,6 +8701,26 @@ pub enum CopyIntoSnowflakeKind {
     Location,
 }
 
+/// Specifies Include / Exclude NULL within UNPIVOT command.
+/// For example
+/// `UNPIVOT (column1 FOR new_column IN (col3, col4, col5, col6))`
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum NullInclusion {
+    IncludeNulls,
+    ExcludeNulls,
+}
+
+impl fmt::Display for NullInclusion {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            NullInclusion::IncludeNulls => write!(f, "INCLUDE NULLS"),
+            NullInclusion::ExcludeNulls => write!(f, "EXCLUDE NULLS"),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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>,
+        null_inclusion: Option<NullInclusion>,
         alias: Option<TableAlias>,
     },
     /// A `MATCH_RECOGNIZE` operation on a table.
@@ -1893,15 +1894,19 @@ impl fmt::Display for TableFactor {
             }
             TableFactor::Unpivot {
                 table,
+                null_inclusion,
                 value,
                 name,
                 columns,
                 alias,
             } => {
+                write!(f, "{table} UNPIVOT")?;
+                if let Some(null_inclusion) = null_inclusion {
+                    write!(f, " {null_inclusion} ")?;
+                }
                 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,
+                null_inclusion: _,
                 name,
                 columns,
                 alias,

src/parser/mod.rs

@@ -12396,6 +12396,15 @@ impl<'a> Parser<'a> {
         &mut self,
         table: TableFactor,
     ) -> Result<TableFactor, ParserError> {
+        let null_inclusion = if self.parse_keyword(Keyword::INCLUDE) {
+            self.expect_keyword_is(Keyword::NULLS)?;
+            Some(NullInclusion::IncludeNulls)
+        } else if self.parse_keyword(Keyword::EXCLUDE) {
+            self.expect_keyword_is(Keyword::NULLS)?;
+            Some(NullInclusion::ExcludeNulls)
+        } 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,
+            null_inclusion,
             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![],
+        }),
+        null_inclusion: 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 { null_inclusion, .. } =
+        &verified_only_select(sql_unpivot_exclude_nulls).from[0].relation
+    {
+        assert_eq!(*null_inclusion, Some(NullInclusion::ExcludeNulls));
+    }
+
+    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 { null_inclusion, .. } =
+        &verified_only_select(sql_unpivot_include_nulls).from[0].relation
+    {
+        assert_eq!(*null_inclusion, Some(NullInclusion::IncludeNulls));
+    }
+
+    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![],
                 }),
+                null_inclusion: None,
                 value: Ident {
                     value: "population".to_string(),
                     quote_style: None,