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
This commit is contained in:
sivchari 2022-05-27 19:27:51 +09:00 committed by GitHub
parent d19c6c323c
commit aa46e930c5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 130 additions and 1 deletions

View file

@ -355,6 +355,19 @@ pub enum TableFactor {
expr: Expr,
alias: Option<TableAlias>,
},
/// SELECT * FROM UNNEST ([10,20,30]) as numbers WITH OFFSET;
/// +---------+--------+
/// | numbers | offset |
/// +---------+--------+
/// | 10 | 0 |
/// | 20 | 1 |
/// | 30 | 2 |
/// +---------+--------+
UNNEST {
alias: Option<TableAlias>,
array_expr: Box<Expr>,
with_offset: bool,
},
/// Represents a parenthesized table factor. The SQL spec only allows a
/// join expression (`(foo <JOIN> bar [ <JOIN> 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),
}
}

View file

@ -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<TableWithJoins, ParserError> {
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:

View file

@ -2786,6 +2786,84 @@ fn parse_table_function() {
);
}
#[test]
fn parse_unnest() {
fn chk(alias: bool, with_offset: bool, dialects: &TestedDialects, want: Vec<TableWithJoins>) {
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