map access on a column

Author: hntd187

src/ast/ddl.rs

@@ -295,7 +295,7 @@ impl fmt::Display for ColumnOption {
     }
 }
 
-fn display_constraint_name<'a>(name: &'a Option<Ident>) -> impl fmt::Display + 'a {
+fn display_constraint_name(name: &Option<Ident>) -> impl fmt::Display + '_ {
     struct ConstraintName<'a>(&'a Option<Ident>);
     impl<'a> fmt::Display for ConstraintName<'a> {
         fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

src/ast/mod.rs

@@ -192,7 +192,10 @@ pub enum Expr {
         right: Box<Expr>,
     },
     /// Unary operation e.g. `NOT foo`
-    UnaryOp { op: UnaryOperator, expr: Box<Expr> },
+    UnaryOp {
+        op: UnaryOperator,
+        expr: Box<Expr>,
+    },
     /// CAST an expression to a different data type e.g. `CAST(foo AS VARCHAR(123))`
     Cast {
         expr: Box<Expr>,
@@ -214,7 +217,14 @@ pub enum Expr {
     /// A constant of form `<data_type> 'value'`.
     /// This can represent ANSI SQL `DATE`, `TIME`, and `TIMESTAMP` literals (such as `DATE '2020-01-01'`),
     /// as well as constants of other types (a non-standard PostgreSQL extension).
-    TypedString { data_type: DataType, value: String },
+    TypedString {
+        data_type: DataType,
+        value: String,
+    },
+    MapAccess {
+        column: Box<Expr>,
+        key: String,
+    },
     /// Scalar function call e.g. `LEFT(foo, 5)`
     Function(Function),
     /// `CASE [<operand>] WHEN <condition> THEN <result> ... [ELSE <result>] END`
@@ -242,6 +252,7 @@ impl fmt::Display for Expr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Expr::Identifier(s) => write!(f, "{}", s),
+            Expr::MapAccess { column, key } => write!(f, "{}[\"{}\"]", column, key),
             Expr::Wildcard => f.write_str("*"),
             Expr::QualifiedWildcard(q) => write!(f, "{}.*", display_separated(q, ".")),
             Expr::CompoundIdentifier(s) => write!(f, "{}", display_separated(s, ".")),

src/ast/query.rs

@@ -413,7 +413,7 @@ impl fmt::Display for Join {
                 _ => "",
             }
         }
-        fn suffix<'a>(constraint: &'a JoinConstraint) -> impl fmt::Display + 'a {
+        fn suffix(constraint: &JoinConstraint) -> impl fmt::Display + '_ {
             struct Suffix<'a>(&'a JoinConstraint);
             impl<'a> fmt::Display for Suffix<'a> {
                 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

src/dialect/ansi.rs

@@ -17,13 +17,13 @@ pub struct AnsiDialect {}
 
 impl Dialect for AnsiDialect {
     fn is_identifier_start(&self, ch: char) -> bool {
-        (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z')
+        ('a'..='z').contains(&ch) || ('A'..='Z').contains(&ch)
     }
 
     fn is_identifier_part(&self, ch: char) -> bool {
-        (ch >= 'a' && ch <= 'z')
-            || (ch >= 'A' && ch <= 'Z')
-            || (ch >= '0' && ch <= '9')
+        ('a'..='z').contains(&ch)
+            || ('A'..='Z').contains(&ch)
+            || ('0'..='9').contains(&ch)
             || ch == '_'
     }
 }

src/dialect/generic.rs

@@ -17,13 +17,17 @@ pub struct GenericDialect;
 
 impl Dialect for GenericDialect {
     fn is_identifier_start(&self, ch: char) -> bool {
-        (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || ch == '_' || ch == '#' || ch == '@'
+        ('a'..='z').contains(&ch)
+            || ('A'..='Z').contains(&ch)
+            || ch == '_'
+            || ch == '#'
+            || ch == '@'
     }
 
     fn is_identifier_part(&self, ch: char) -> bool {
-        (ch >= 'a' && ch <= 'z')
-            || (ch >= 'A' && ch <= 'Z')
-            || (ch >= '0' && ch <= '9')
+        ('a'..='z').contains(&ch)
+            || ('A'..='Z').contains(&ch)
+            || ('0'..='9').contains(&ch)
             || ch == '@'
             || ch == '$'
             || ch == '#'

src/parser.rs

@@ -48,12 +48,14 @@ pub enum IsOptional {
     Optional,
     Mandatory,
 }
+
 use IsOptional::*;
 
 pub enum IsLateral {
     Lateral,
     NotLateral,
 }
+
 use crate::ast::Statement::CreateVirtualTable;
 use IsLateral::*;
 
@@ -281,6 +283,7 @@ impl<'a> Parser<'a> {
         loop {
             let next_precedence = self.get_next_precedence()?;
             debug!("next precedence: {:?}", next_precedence);
+
             if precedence >= next_precedence {
                 break;
             }
@@ -422,6 +425,7 @@ impl<'a> Parser<'a> {
                 self.prev_token();
                 Ok(Expr::Value(self.parse_value()?))
             }
+
             Token::LParen => {
                 let expr =
                     if self.parse_keyword(Keyword::SELECT) || self.parse_keyword(Keyword::WITH) {
@@ -845,12 +849,27 @@ impl<'a> Parser<'a> {
                 op: UnaryOperator::PGPostfixFactorial,
                 expr: Box::new(expr),
             })
+        } else if Token::LBracket == tok {
+            self.parse_map_access(expr)
         } else {
             // Can only happen if `get_next_precedence` got out of sync with this function
             panic!("No infix parser for token {:?}", tok)
         }
     }
 
+    pub fn parse_map_access(&mut self, expr: Expr) -> Result<Expr, ParserError> {
+        let key = self.parse_literal_string()?;
+        let tok = self.consume_token(&Token::RBracket);
+        debug!("Tok: {}", tok);
+        match expr {
+            e @ Expr::Identifier(_) | e @ Expr::CompoundIdentifier(_) => Ok(Expr::MapAccess {
+                column: Box::new(e),
+                key,
+            }),
+            _ => Ok(expr),
+        }
+    }
+
     /// Parses the parens following the `[ NOT ] IN` operator
     pub fn parse_in(&mut self, expr: Expr, negated: bool) -> Result<Expr, ParserError> {
         self.expect_token(&Token::LParen)?;
@@ -936,6 +955,7 @@ impl<'a> Parser<'a> {
             Token::Mult | Token::Div | Token::Mod | Token::StringConcat => Ok(40),
             Token::DoubleColon => Ok(50),
             Token::ExclamationMark => Ok(50),
+            Token::LBracket | Token::RBracket => Ok(10),
             _ => Ok(0),
         }
     }
@@ -2460,7 +2480,7 @@ impl<'a> Parser<'a> {
                         }
                     }
                     Keyword::OUTER => {
-                        return self.expected("LEFT, RIGHT, or FULL", self.peek_token())
+                        return self.expected("LEFT, RIGHT, or FULL", self.peek_token());
                     }
                     _ if natural => {
                         return self.expected("a join type after NATURAL", self.peek_token());

tests/sqlparser_hive.rs

@@ -192,6 +192,12 @@ fn rename_table() {
     hive().verified_stmt(rename);
 }
 
+#[test]
+fn map_access() {
+    let rename = "SELECT a.b[\"asdf\"] FROM db.table WHERE a = 2";
+    hive().verified_stmt(rename);
+}
+
 fn hive() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(HiveDialect {})],