Skip to content

Commit

Permalink
Support Databricks struct literal (#1542)
Browse files Browse the repository at this point in the history
  • Loading branch information
ayman-sigma authored Dec 2, 2024
1 parent 4ab3ab9 commit bd750df
Show file tree
Hide file tree
Showing 7 changed files with 71 additions and 9 deletions.
6 changes: 4 additions & 2 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -931,12 +931,14 @@ pub enum Expr {
Rollup(Vec<Vec<Expr>>),
/// ROW / TUPLE a single value, such as `SELECT (1, 2)`
Tuple(Vec<Expr>),
/// `BigQuery` specific `Struct` literal expression [1]
/// `Struct` literal expression
/// Syntax:
/// ```sql
/// STRUCT<[field_name] field_type, ...>( expr1 [, ... ])
///
/// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#struct_type)
/// [Databricks](https://docs.databricks.com/en/sql/language-manual/functions/struct.html)
/// ```
/// [1]: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#struct_type
Struct {
/// Struct values.
values: Vec<Expr>,
Expand Down
5 changes: 5 additions & 0 deletions src/dialect/bigquery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,4 +72,9 @@ impl Dialect for BigQueryDialect {
fn require_interval_qualifier(&self) -> bool {
true
}

// See https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#constructing_a_struct
fn supports_struct_literal(&self) -> bool {
true
}
}
5 changes: 5 additions & 0 deletions src/dialect/databricks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -59,4 +59,9 @@ impl Dialect for DatabricksDialect {
fn require_interval_qualifier(&self) -> bool {
true
}

// See https://docs.databricks.com/en/sql/language-manual/functions/struct.html
fn supports_struct_literal(&self) -> bool {
true
}
}
4 changes: 4 additions & 0 deletions src/dialect/generic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,4 +123,8 @@ impl Dialect for GenericDialect {
fn supports_named_fn_args_with_assignment_operator(&self) -> bool {
true
}

fn supports_struct_literal(&self) -> bool {
true
}
}
10 changes: 10 additions & 0 deletions src/dialect/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -375,6 +375,16 @@ pub trait Dialect: Debug + Any {
false
}

/// Return true if the dialect supports the STRUCT literal
///
/// Example
/// ```sql
/// SELECT STRUCT(1 as one, 'foo' as foo, false)
/// ```
fn supports_struct_literal(&self) -> bool {
false
}

/// Dialect-specific infix parser override
///
/// This method is called to parse the next infix expression.
Expand Down
15 changes: 8 additions & 7 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1123,9 +1123,8 @@ impl<'a> Parser<'a> {
Keyword::MATCH if dialect_of!(self is MySqlDialect | GenericDialect) => {
Ok(Some(self.parse_match_against()?))
}
Keyword::STRUCT if dialect_of!(self is BigQueryDialect | GenericDialect) => {
self.prev_token();
Ok(Some(self.parse_bigquery_struct_literal()?))
Keyword::STRUCT if self.dialect.supports_struct_literal() => {
Ok(Some(self.parse_struct_literal()?))
}
Keyword::PRIOR if matches!(self.state, ParserState::ConnectBy) => {
let expr = self.parse_subexpr(self.dialect.prec_value(Precedence::PlusMinus))?;
Expand Down Expand Up @@ -2383,15 +2382,16 @@ impl<'a> Parser<'a> {
}
}

/// Bigquery specific: Parse a struct literal
/// Syntax
/// ```sql
/// -- typed
/// STRUCT<[field_name] field_type, ...>( expr1 [, ... ])
/// -- typeless
/// STRUCT( expr1 [AS field_name] [, ... ])
/// ```
fn parse_bigquery_struct_literal(&mut self) -> Result<Expr, ParserError> {
fn parse_struct_literal(&mut self) -> Result<Expr, ParserError> {
// Parse the fields definition if exist `<[field_name] field_type, ...>`
self.prev_token();
let (fields, trailing_bracket) =
self.parse_struct_type_def(Self::parse_struct_field_def)?;
if trailing_bracket.0 {
Expand All @@ -2401,6 +2401,7 @@ impl<'a> Parser<'a> {
);
}

// Parse the struct values `(expr1 [, ... ])`
self.expect_token(&Token::LParen)?;
let values = self
.parse_comma_separated(|parser| parser.parse_struct_field_expr(!fields.is_empty()))?;
Expand All @@ -2409,13 +2410,13 @@ impl<'a> Parser<'a> {
Ok(Expr::Struct { values, fields })
}

/// Parse an expression value for a bigquery struct [1]
/// Parse an expression value for a struct literal
/// Syntax
/// ```sql
/// expr [AS name]
/// ```
///
/// Parameter typed_syntax is set to true if the expression
/// For biquery [1], Parameter typed_syntax is set to true if the expression
/// is to be parsed as a field expression declared using typed
/// struct syntax [2], and false if using typeless struct syntax [3].
///
Expand Down
35 changes: 35 additions & 0 deletions tests/sqlparser_databricks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -278,3 +278,38 @@ fn parse_use() {
);
}
}

#[test]
fn parse_databricks_struct_function() {
assert_eq!(
databricks_and_generic()
.verified_only_select("SELECT STRUCT(1, 'foo')")
.projection[0],
SelectItem::UnnamedExpr(Expr::Struct {
values: vec![
Expr::Value(number("1")),
Expr::Value(Value::SingleQuotedString("foo".to_string()))
],
fields: vec![]
})
);
assert_eq!(
databricks_and_generic()
.verified_only_select("SELECT STRUCT(1 AS one, 'foo' AS foo, false)")
.projection[0],
SelectItem::UnnamedExpr(Expr::Struct {
values: vec![
Expr::Named {
expr: Expr::Value(number("1")).into(),
name: Ident::new("one")
},
Expr::Named {
expr: Expr::Value(Value::SingleQuotedString("foo".to_string())).into(),
name: Ident::new("foo")
},
Expr::Value(Value::Boolean(false))
],
fields: vec![]
})
);
}

0 comments on commit bd750df

Please sign in to comment.