Skip to content

Commit

Permalink
Support Map literal syntax for DuckDB and Generic (#1344)
Browse files Browse the repository at this point in the history
  • Loading branch information
goldmedal authored Jul 21, 2024
1 parent 71dc966 commit 48ea564
Show file tree
Hide file tree
Showing 7 changed files with 219 additions and 0 deletions.
42 changes: 42 additions & 0 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -329,6 +329,37 @@ impl fmt::Display for DictionaryField {
}
}

/// Represents a Map expression.
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct Map {
pub entries: Vec<MapEntry>,
}

impl Display for Map {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "MAP {{{}}}", display_comma_separated(&self.entries))
}
}

/// A map field within a map.
///
/// [duckdb]: https://duckdb.org/docs/sql/data_types/map.html#creating-maps
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct MapEntry {
pub key: Box<Expr>,
pub value: Box<Expr>,
}

impl fmt::Display for MapEntry {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}: {}", self.key, self.value)
}
}

/// Options for `CAST` / `TRY_CAST`
/// BigQuery: <https://cloud.google.com/bigquery/docs/reference/standard-sql/format-elements#formatting_syntax>
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
Expand Down Expand Up @@ -764,6 +795,14 @@ pub enum Expr {
/// ```
/// [1]: https://duckdb.org/docs/sql/data_types/struct#creating-structs
Dictionary(Vec<DictionaryField>),
/// `DuckDB` specific `Map` literal expression [1]
///
/// Syntax:
/// ```sql
/// syntax: Map {key1: value1[, ... ]}
/// ```
/// [1]: https://duckdb.org/docs/sql/data_types/map#creating-maps
Map(Map),
/// An access of nested data using subscript syntax, for example `array[2]`.
Subscript {
expr: Box<Expr>,
Expand Down Expand Up @@ -1331,6 +1370,9 @@ impl fmt::Display for Expr {
Expr::Dictionary(fields) => {
write!(f, "{{{}}}", display_comma_separated(fields))
}
Expr::Map(map) => {
write!(f, "{map}")
}
Expr::Subscript {
expr,
subscript: key,
Expand Down
7 changes: 7 additions & 0 deletions src/dialect/duckdb.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,4 +48,11 @@ impl Dialect for DuckDbDialect {
fn supports_dictionary_syntax(&self) -> bool {
true
}

// DuckDB uses this syntax for `MAP`s.
//
// https://duckdb.org/docs/sql/data_types/map.html#creating-maps
fn support_map_literal_syntax(&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 @@ -70,4 +70,8 @@ impl Dialect for GenericDialect {
fn supports_select_wildcard_except(&self) -> bool {
true
}

fn support_map_literal_syntax(&self) -> bool {
true
}
}
5 changes: 5 additions & 0 deletions src/dialect/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,11 @@ pub trait Dialect: Debug + Any {
fn supports_dictionary_syntax(&self) -> bool {
false
}
/// Returns true if the dialect supports defining object using the
/// syntax like `Map {1: 10, 2: 20}`.
fn support_map_literal_syntax(&self) -> bool {
false
}
/// Returns true if the dialect supports lambda functions, for example:
///
/// ```sql
Expand Down
44 changes: 44 additions & 0 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1078,6 +1078,9 @@ impl<'a> Parser<'a> {
let expr = self.parse_subexpr(Self::PLUS_MINUS_PREC)?;
Ok(Expr::Prior(Box::new(expr)))
}
Keyword::MAP if self.peek_token() == Token::LBrace && self.dialect.support_map_literal_syntax() => {
self.parse_duckdb_map_literal()
}
// Here `w` is a word, check if it's a part of a multipart
// identifier, a function call, or a simple identifier:
_ => match self.peek_token().token {
Expand Down Expand Up @@ -2322,6 +2325,47 @@ impl<'a> Parser<'a> {
})
}

/// DuckDB specific: Parse a duckdb [map]
///
/// Syntax:
///
/// ```sql
/// Map {key1: value1[, ... ]}
/// ```
///
/// [map]: https://duckdb.org/docs/sql/data_types/map.html#creating-maps
fn parse_duckdb_map_literal(&mut self) -> Result<Expr, ParserError> {
self.expect_token(&Token::LBrace)?;

let fields = self.parse_comma_separated(Self::parse_duckdb_map_field)?;

self.expect_token(&Token::RBrace)?;

Ok(Expr::Map(Map { entries: fields }))
}

/// Parse a field for a duckdb [map]
///
/// Syntax
///
/// ```sql
/// key: value
/// ```
///
/// [map]: https://duckdb.org/docs/sql/data_types/map.html#creating-maps
fn parse_duckdb_map_field(&mut self) -> Result<MapEntry, ParserError> {
let key = self.parse_expr()?;

self.expect_token(&Token::Colon)?;

let value = self.parse_expr()?;

Ok(MapEntry {
key: Box::new(key),
value: Box::new(value),
})
}

/// Parse clickhouse [map]
///
/// Syntax
Expand Down
95 changes: 95 additions & 0 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10077,6 +10077,101 @@ fn test_dictionary_syntax() {
)
}

#[test]
fn test_map_syntax() {
fn check(sql: &str, expect: Expr) {
assert_eq!(
all_dialects_where(|d| d.support_map_literal_syntax()).verified_expr(sql),
expect
);
}

check(
"MAP {'Alberta': 'Edmonton', 'Manitoba': 'Winnipeg'}",
Expr::Map(Map {
entries: vec![
MapEntry {
key: Box::new(Expr::Value(Value::SingleQuotedString("Alberta".to_owned()))),
value: Box::new(Expr::Value(Value::SingleQuotedString(
"Edmonton".to_owned(),
))),
},
MapEntry {
key: Box::new(Expr::Value(Value::SingleQuotedString(
"Manitoba".to_owned(),
))),
value: Box::new(Expr::Value(Value::SingleQuotedString(
"Winnipeg".to_owned(),
))),
},
],
}),
);

fn number_expr(s: &str) -> Expr {
Expr::Value(number(s))
}

check(
"MAP {1: 10.0, 2: 20.0}",
Expr::Map(Map {
entries: vec![
MapEntry {
key: Box::new(number_expr("1")),
value: Box::new(number_expr("10.0")),
},
MapEntry {
key: Box::new(number_expr("2")),
value: Box::new(number_expr("20.0")),
},
],
}),
);

check(
"MAP {[1, 2, 3]: 10.0, [4, 5, 6]: 20.0}",
Expr::Map(Map {
entries: vec![
MapEntry {
key: Box::new(Expr::Array(Array {
elem: vec![number_expr("1"), number_expr("2"), number_expr("3")],
named: false,
})),
value: Box::new(Expr::Value(number("10.0"))),
},
MapEntry {
key: Box::new(Expr::Array(Array {
elem: vec![number_expr("4"), number_expr("5"), number_expr("6")],
named: false,
})),
value: Box::new(Expr::Value(number("20.0"))),
},
],
}),
);

check(
"MAP {'a': 10, 'b': 20}['a']",
Expr::Subscript {
expr: Box::new(Expr::Map(Map {
entries: vec![
MapEntry {
key: Box::new(Expr::Value(Value::SingleQuotedString("a".to_owned()))),
value: Box::new(number_expr("10")),
},
MapEntry {
key: Box::new(Expr::Value(Value::SingleQuotedString("b".to_owned()))),
value: Box::new(number_expr("20")),
},
],
})),
subscript: Box::new(Subscript::Index {
index: Expr::Value(Value::SingleQuotedString("a".to_owned())),
}),
},
);
}

#[test]
fn parse_within_group() {
verified_expr("PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY sales_amount)");
Expand Down
22 changes: 22 additions & 0 deletions tests/sqlparser_custom_dialect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,28 @@ fn custom_statement_parser() -> Result<(), ParserError> {
Ok(())
}

#[test]
fn test_map_syntax_not_support_default() -> Result<(), ParserError> {
#[derive(Debug)]
struct MyDialect {}

impl Dialect for MyDialect {
fn is_identifier_start(&self, ch: char) -> bool {
is_identifier_start(ch)
}

fn is_identifier_part(&self, ch: char) -> bool {
is_identifier_part(ch)
}
}

let dialect = MyDialect {};
let sql = "SELECT MAP {1: 2}";
let ast = Parser::parse_sql(&dialect, sql);
assert!(ast.is_err());
Ok(())
}

fn is_identifier_start(ch: char) -> bool {
ch.is_ascii_lowercase() || ch.is_ascii_uppercase() || ch == '_'
}
Expand Down

0 comments on commit 48ea564

Please sign in to comment.