From aa46e930c5ed1125820538d3e75492de4b35ee3d Mon Sep 17 00:00:00 2001 From: sivchari Date: Fri, 27 May 2022 19:27:51 +0900 Subject: [PATCH] Support `UNNEST` as a table factor (#493) * support unnest * add test * fix ast * Update condition for BigQueryDialect or GenericDialect I updated condition. This changes conditionalize parsing for only BigQueryDialect or GenericDialect. * Add some tests * fix test --- src/ast/query.rs | 27 ++++++++++++++ src/parser.rs | 26 ++++++++++++- tests/sqlparser_common.rs | 78 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 130 insertions(+), 1 deletion(-) diff --git a/src/ast/query.rs b/src/ast/query.rs index 07295f44f..6212469f4 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -355,6 +355,19 @@ pub enum TableFactor { expr: Expr, alias: Option, }, + /// SELECT * FROM UNNEST ([10,20,30]) as numbers WITH OFFSET; + /// +---------+--------+ + /// | numbers | offset | + /// +---------+--------+ + /// | 10 | 0 | + /// | 20 | 1 | + /// | 30 | 2 | + /// +---------+--------+ + UNNEST { + alias: Option, + array_expr: Box, + with_offset: bool, + }, /// Represents a parenthesized table factor. The SQL spec only allows a /// join expression (`(foo bar [ baz ... ])`) to be nested, /// possibly several times. @@ -406,6 +419,20 @@ impl fmt::Display for TableFactor { } Ok(()) } + TableFactor::UNNEST { + alias, + array_expr, + with_offset, + } => { + write!(f, "UNNEST({})", array_expr)?; + if let Some(alias) = alias { + write!(f, " AS {}", alias)?; + } + if *with_offset { + write!(f, " WITH OFFSET")?; + } + Ok(()) + } TableFactor::NestedJoin(table_reference) => write!(f, "({})", table_reference), } } diff --git a/src/parser.rs b/src/parser.rs index 23fd79abb..36c2075c0 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -3242,6 +3242,7 @@ impl<'a> Parser<'a> { } else { vec![] }; + let mut lateral_views = vec![]; loop { if self.parse_keywords(&[Keyword::LATERAL, Keyword::VIEW]) { @@ -3490,7 +3491,6 @@ impl<'a> Parser<'a> { pub fn parse_table_and_joins(&mut self) -> Result { let relation = self.parse_table_factor()?; - // Note that for keywords to be properly handled here, they need to be // added to `RESERVED_FOR_TABLE_ALIAS`, otherwise they may be parsed as // a table alias. @@ -3635,6 +3635,7 @@ impl<'a> Parser<'a> { match &mut table_and_joins.relation { TableFactor::Derived { alias, .. } | TableFactor::Table { alias, .. } + | TableFactor::UNNEST { alias, .. } | TableFactor::TableFunction { alias, .. } => { // but not `FROM (mytable AS alias1) AS alias2`. if let Some(inner_alias) = alias { @@ -3658,6 +3659,29 @@ impl<'a> Parser<'a> { // appearing alone in parentheses (e.g. `FROM (mytable)`) self.expected("joined table", self.peek_token()) } + } else if dialect_of!(self is BigQueryDialect | GenericDialect) + && self.parse_keyword(Keyword::UNNEST) + { + self.expect_token(&Token::LParen)?; + let expr = self.parse_expr()?; + self.expect_token(&Token::RParen)?; + + let alias = match self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS) { + Ok(Some(alias)) => Some(alias), + Ok(None) => None, + Err(e) => return Err(e), + }; + + let with_offset = match self.expect_keywords(&[Keyword::WITH, Keyword::OFFSET]) { + Ok(()) => true, + Err(_) => false, + }; + + Ok(TableFactor::UNNEST { + alias, + array_expr: Box::new(expr), + with_offset, + }) } else { let name = self.parse_object_name()?; // Postgres, MSSQL: table-valued functions: diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index e931cd96f..83dacb046 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -2786,6 +2786,84 @@ fn parse_table_function() { ); } +#[test] +fn parse_unnest() { + fn chk(alias: bool, with_offset: bool, dialects: &TestedDialects, want: Vec) { + let sql = &format!( + "SELECT * FROM UNNEST(expr){}{}", + if alias { " AS numbers" } else { "" }, + if with_offset { " WITH OFFSET" } else { "" }, + ); + let select = dialects.verified_only_select(sql); + assert_eq!(select.from, want); + } + let dialects = TestedDialects { + dialects: vec![Box::new(BigQueryDialect {}), Box::new(GenericDialect {})], + }; + // 1. both Alias and WITH OFFSET clauses. + chk( + true, + true, + &dialects, + vec![TableWithJoins { + relation: TableFactor::UNNEST { + alias: Some(TableAlias { + name: Ident::new("numbers"), + columns: vec![], + }), + array_expr: Box::new(Expr::Identifier(Ident::new("expr"))), + with_offset: true, + }, + joins: vec![], + }], + ); + // 2. neither Alias nor WITH OFFSET clause. + chk( + false, + false, + &dialects, + vec![TableWithJoins { + relation: TableFactor::UNNEST { + alias: None, + array_expr: Box::new(Expr::Identifier(Ident::new("expr"))), + with_offset: false, + }, + joins: vec![], + }], + ); + // 3. Alias but no WITH OFFSET clause. + chk( + false, + true, + &dialects, + vec![TableWithJoins { + relation: TableFactor::UNNEST { + alias: None, + array_expr: Box::new(Expr::Identifier(Ident::new("expr"))), + with_offset: true, + }, + joins: vec![], + }], + ); + // 4. WITH OFFSET but no Alias. + chk( + true, + false, + &dialects, + vec![TableWithJoins { + relation: TableFactor::UNNEST { + alias: Some(TableAlias { + name: Ident::new("numbers"), + columns: vec![], + }), + array_expr: Box::new(Expr::Identifier(Ident::new("expr"))), + with_offset: false, + }, + joins: vec![], + }], + ); +} + #[test] fn parse_delimited_identifiers() { // check that quoted identifiers in any position remain quoted after serialization