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

This commit is contained in:
Denys Tsomenko 2025-05-15 17:40:21 +03:00 committed by GitHub
parent c6e897dc12
commit 3c59950060
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 111 additions and 44 deletions

View file

@ -9633,6 +9633,26 @@ impl fmt::Display for OpenStatement {
} }
} }
/// 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)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;

View file

@ -1336,7 +1336,7 @@ pub enum TableFactor {
/// ///
/// Syntax: /// Syntax:
/// ```sql /// ```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>. /// See <https://docs.snowflake.com/en/sql-reference/constructs/unpivot>.
@ -1345,6 +1345,7 @@ pub enum TableFactor {
value: Ident, value: Ident,
name: Ident, name: Ident,
columns: Vec<Ident>, columns: Vec<Ident>,
null_inclusion: Option<NullInclusion>,
alias: Option<TableAlias>, alias: Option<TableAlias>,
}, },
/// A `MATCH_RECOGNIZE` operation on a table. /// A `MATCH_RECOGNIZE` operation on a table.
@ -2015,15 +2016,19 @@ impl fmt::Display for TableFactor {
} }
TableFactor::Unpivot { TableFactor::Unpivot {
table, table,
null_inclusion,
value, value,
name, name,
columns, columns,
alias, alias,
} => { } => {
write!(f, "{table} UNPIVOT")?;
if let Some(null_inclusion) = null_inclusion {
write!(f, " {null_inclusion} ")?;
}
write!( write!(
f, f,
"{} UNPIVOT({} FOR {} IN ({}))", "({} FOR {} IN ({}))",
table,
value, value,
name, name,
display_comma_separated(columns) display_comma_separated(columns)

View file

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

View file

@ -13368,6 +13368,15 @@ impl<'a> Parser<'a> {
&mut self, &mut self,
table: TableFactor, table: TableFactor,
) -> Result<TableFactor, ParserError> { ) -> 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)?; self.expect_token(&Token::LParen)?;
let value = self.parse_identifier()?; let value = self.parse_identifier()?;
self.expect_keyword_is(Keyword::FOR)?; self.expect_keyword_is(Keyword::FOR)?;
@ -13379,6 +13388,7 @@ impl<'a> Parser<'a> {
Ok(TableFactor::Unpivot { Ok(TableFactor::Unpivot {
table: Box::new(table), table: Box::new(table),
value, value,
null_inclusion,
name, name,
columns, columns,
alias, alias,

View file

@ -10746,49 +10746,47 @@ fn parse_unpivot_table() {
"SELECT * FROM sales AS s ", "SELECT * FROM sales AS s ",
"UNPIVOT(quantity FOR quarter IN (Q1, Q2, Q3, Q4)) AS u (product, quarter, quantity)" "UNPIVOT(quantity FOR quarter IN (Q1, Q2, Q3, Q4)) AS u (product, quarter, quantity)"
); );
let base_unpivot = Unpivot {
pretty_assertions::assert_eq!( table: Box::new(TableFactor::Table {
verified_only_select(sql).from[0].relation, name: ObjectName::from(vec![Ident::new("sales")]),
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![],
}),
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"]
.into_iter()
.map(Ident::new)
.collect(),
alias: Some(TableAlias { alias: Some(TableAlias {
name: Ident::new("u"), name: Ident::new("s"),
columns: ["product", "quarter", "quantity"] columns: vec![],
.into_iter()
.map(TableAliasColumnDef::from_name)
.collect(),
}), }),
} 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"]
.into_iter()
.map(Ident::new)
.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); assert_eq!(verified_stmt(sql).to_string(), sql);
let sql_without_aliases = concat!( let sql_without_aliases = concat!(
@ -10808,6 +10806,38 @@ fn parse_unpivot_table() {
verified_stmt(sql_without_aliases).to_string(), verified_stmt(sql_without_aliases).to_string(),
sql_without_aliases 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] #[test]
@ -10904,6 +10934,7 @@ fn parse_pivot_unpivot_table() {
sample: None, sample: None,
index_hints: vec![], index_hints: vec![],
}), }),
null_inclusion: None,
value: Ident { value: Ident {
value: "population".to_string(), value: "population".to_string(),
quote_style: None, quote_style: None,