Skip to content

Commit

Permalink
Support UNNEST as a table factor (#493)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
sivchari committed May 27, 2022
1 parent d19c6c3 commit aa46e93
Show file tree
Hide file tree
Showing 3 changed files with 130 additions and 1 deletion.
27 changes: 27 additions & 0 deletions src/ast/query.rs
Expand Up @@ -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.
Expand Down Expand Up @@ -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),
}
}
Expand Down
26 changes: 25 additions & 1 deletion src/parser.rs
Expand Up @@ -3242,6 +3242,7 @@ impl<'a> Parser<'a> {
} else {
vec![]
};

let mut lateral_views = vec![];
loop {
if self.parse_keywords(&[Keyword::LATERAL, Keyword::VIEW]) {
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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 {
Expand All @@ -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:
Expand Down
78 changes: 78 additions & 0 deletions tests/sqlparser_common.rs
Expand Up @@ -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
Expand Down

0 comments on commit aa46e93

Please sign in to comment.