Support array expressions such as ARRAY[1,2] , foo[1] and INT[][] (#419)

* feat: add array expression

* test: add back the the existing test

Co-authored-by: gamife <gamife9886@gmail.com>

src/ast/mod.rs

@@ -286,6 +286,13 @@ pub enum Expr {
     Rollup(Vec<Vec<Expr>>),
     /// ROW / TUPLE a single value, such as `SELECT (1, 2)`
     Tuple(Vec<Expr>),
+    /// An array index expression e.g. `(ARRAY[1, 2])[1]` or `(current_schemas(FALSE))[1]`
+    ArrayIndex {
+        obj: Box<Expr>,
+        indexs: Vec<Expr>,
+    },
+    /// An array expression e.g. `ARRAY[1, 2]`
+    Array(Vec<Expr>),
 }
 
 impl fmt::Display for Expr {
@@ -450,6 +457,16 @@ impl fmt::Display for Expr {
             Expr::Tuple(exprs) => {
                 write!(f, "({})", display_comma_separated(exprs))
             }
+            Expr::ArrayIndex { obj, indexs } => {
+                write!(f, "{}", obj)?;
+                for i in indexs {
+                    write!(f, "[{}]", i)?;
+                }
+                Ok(())
+            }
+            Expr::Array(set) => {
+                write!(f, "ARRAY[{}]", display_comma_separated(set))
+            }
         }
     }
 }

src/parser.rs

@@ -420,6 +420,7 @@ impl<'a> Parser<'a> {
                 Keyword::TRIM => self.parse_trim_expr(),
                 Keyword::INTERVAL => self.parse_literal_interval(),
                 Keyword::LISTAGG => self.parse_listagg_expr(),
+                Keyword::ARRAY => self.parse_array_expr(),
                 Keyword::NOT => Ok(Expr::UnaryOp {
                     op: UnaryOperator::Not,
                     expr: Box::new(self.parse_subexpr(Self::UNARY_NOT_PREC)?),
@@ -820,6 +821,13 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_array_expr(&mut self) -> Result<Expr, ParserError> {
+        self.expect_token(&Token::LBracket)?;
+        let exprs = self.parse_comma_separated(Parser::parse_expr)?;
+        self.expect_token(&Token::RBracket)?;
+        Ok(Expr::Array(exprs))
+    }
+
     /// Parse a SQL LISTAGG expression, e.g. `LISTAGG(...) WITHIN GROUP (ORDER BY ...)`.
     pub fn parse_listagg_expr(&mut self) -> Result<Expr, ParserError> {
         self.expect_token(&Token::LParen)?;
@@ -1083,6 +1091,10 @@ impl<'a> Parser<'a> {
                 expr: Box::new(expr),
             })
         } else if Token::LBracket == tok {
+            if dialect_of!(self is PostgreSqlDialect) {
+                // parse index
+                return self.parse_array_index(expr);
+            }
             self.parse_map_access(expr)
         } else {
             // Can only happen if `get_next_precedence` got out of sync with this function
@@ -1090,6 +1102,21 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_array_index(&mut self, expr: Expr) -> Result<Expr, ParserError> {
+        let index = self.parse_expr()?;
+        self.expect_token(&Token::RBracket)?;
+        let mut indexs: Vec<Expr> = vec![index];
+        while self.consume_token(&Token::LBracket) {
+            let index = self.parse_expr()?;
+            self.expect_token(&Token::RBracket)?;
+            indexs.push(index);
+        }
+        Ok(Expr::ArrayIndex {
+            obj: Box::new(expr),
+            indexs,
+        })
+    }
+
     pub fn parse_map_access(&mut self, expr: Expr) -> Result<Expr, ParserError> {
         let key = self.parse_map_key()?;
         let tok = self.consume_token(&Token::RBracket);
@@ -1202,7 +1229,7 @@ impl<'a> Parser<'a> {
             Token::Mul | Token::Div | Token::Mod | Token::StringConcat => Ok(40),
             Token::DoubleColon => Ok(50),
             Token::ExclamationMark => Ok(50),
-            Token::LBracket | Token::RBracket => Ok(10),
+            Token::LBracket => Ok(10),
             _ => Ok(0),
         }
     }
@@ -2287,7 +2314,7 @@ impl<'a> Parser<'a> {
 
     /// Parse a SQL datatype (in the context of a CREATE TABLE statement for example)
     pub fn parse_data_type(&mut self) -> Result<DataType, ParserError> {
-        match self.next_token() {
+        let mut data = match self.next_token() {
             Token::Word(w) => match w.keyword {
                 Keyword::BOOLEAN => Ok(DataType::Boolean),
                 Keyword::FLOAT => Ok(DataType::Float(self.parse_optional_precision()?)),
@@ -2332,15 +2359,7 @@ impl<'a> Parser<'a> {
                 Keyword::INTERVAL => Ok(DataType::Interval),
                 Keyword::REGCLASS => Ok(DataType::Regclass),
                 Keyword::STRING => Ok(DataType::String),
-                Keyword::TEXT => {
-                    if self.consume_token(&Token::LBracket) {
-                        // Note: this is postgresql-specific
-                        self.expect_token(&Token::RBracket)?;
-                        Ok(DataType::Array(Box::new(DataType::Text)))
-                    } else {
-                        Ok(DataType::Text)
-                    }
-                }
+                Keyword::TEXT => Ok(DataType::Text),
                 Keyword::BYTEA => Ok(DataType::Bytea),
                 Keyword::NUMERIC | Keyword::DECIMAL | Keyword::DEC => {
                     let (precision, scale) = self.parse_optional_precision_scale()?;
@@ -2355,7 +2374,14 @@ impl<'a> Parser<'a> {
                 }
             },
             unexpected => self.expected("a data type name", unexpected),
+        }?;
+
+        // Parse array data types. Note: this is postgresql-specific
+        while self.consume_token(&Token::LBracket) {
+            self.expect_token(&Token::RBracket)?;
+            data = DataType::Array(Box::new(data))
         }
+        Ok(data)
     }
 
     pub fn parse_string_values(&mut self) -> Result<Vec<String>, ParserError> {

tests/sqlparser_postgres.rs

@@ -18,9 +18,6 @@
 mod test_utils;
 use test_utils::*;
 
-#[cfg(feature = "bigdecimal")]
-use bigdecimal::BigDecimal;
-use sqlparser::ast::Expr::{Identifier, MapAccess};
 use sqlparser::ast::*;
 use sqlparser::dialect::{GenericDialect, PostgreSqlDialect};
 use sqlparser::parser::ParserError;
@@ -774,51 +771,72 @@ fn parse_pg_regex_match_ops() {
 }
 
 #[test]
-fn parse_map_access_expr() {
-    #[cfg(not(feature = "bigdecimal"))]
-    let zero = "0".to_string();
+fn parse_array_index_expr() {
     #[cfg(feature = "bigdecimal")]
-    let zero = BigDecimal::parse_bytes(b"0", 10).unwrap();
+    let num: Vec<Expr> = (0..=10)
+        .into_iter()
+        .map(|s| Expr::Value(Value::Number(bigdecimal::BigDecimal::from(s), false)))
+        .collect();
+    #[cfg(not(feature = "bigdecimal"))]
+    let num: Vec<Expr> = (0..=10)
+        .into_iter()
+        .map(|s| Expr::Value(Value::Number(s.to_string(), false)))
+        .collect();
+
     let sql = "SELECT foo[0] FROM foos";
-    let select = pg_and_generic().verified_only_select(sql);
+    let select = pg().verified_only_select(sql);
     assert_eq!(
-        &MapAccess {
-            column: Box::new(Identifier(Ident {
-                value: "foo".to_string(),
-                quote_style: None
-            })),
-            keys: vec![Expr::Value(Value::Number(zero.clone(), false))]
+        &Expr::ArrayIndex {
+            obj: Box::new(Expr::Identifier(Ident::new("foo"))),
+            indexs: vec![num[0].clone()],
         },
         expr_from_projection(only(&select.projection)),
     );
+
     let sql = "SELECT foo[0][0] FROM foos";
-    let select = pg_and_generic().verified_only_select(sql);
+    let select = pg().verified_only_select(sql);
     assert_eq!(
-        &MapAccess {
-            column: Box::new(Identifier(Ident {
-                value: "foo".to_string(),
-                quote_style: None
-            })),
-            keys: vec![
-                Expr::Value(Value::Number(zero.clone(), false)),
-                Expr::Value(Value::Number(zero.clone(), false))
-            ]
+        &Expr::ArrayIndex {
+            obj: Box::new(Expr::Identifier(Ident::new("foo"))),
+            indexs: vec![num[0].clone(), num[0].clone()],
         },
         expr_from_projection(only(&select.projection)),
     );
+
     let sql = r#"SELECT bar[0]["baz"]["fooz"] FROM foos"#;
-    let select = pg_and_generic().verified_only_select(sql);
+    let select = pg().verified_only_select(sql);
     assert_eq!(
-        &MapAccess {
-            column: Box::new(Identifier(Ident {
-                value: "bar".to_string(),
-                quote_style: None
-            })),
-            keys: vec![
-                Expr::Value(Value::Number(zero, false)),
-                Expr::Value(Value::SingleQuotedString("baz".to_string())),
-                Expr::Value(Value::SingleQuotedString("fooz".to_string()))
-            ]
+        &Expr::ArrayIndex {
+            obj: Box::new(Expr::Identifier(Ident::new("bar"))),
+            indexs: vec![
+                num[0].clone(),
+                Expr::Identifier(Ident {
+                    value: "baz".to_string(),
+                    quote_style: Some('"')
+                }),
+                Expr::Identifier(Ident {
+                    value: "fooz".to_string(),
+                    quote_style: Some('"')
+                })
+            ],
+        },
+        expr_from_projection(only(&select.projection)),
+    );
+
+    let sql = "SELECT (CAST(ARRAY[ARRAY[2, 3]] AS INT[][]))[1][2]";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        &Expr::ArrayIndex {
+            obj: Box::new(Expr::Nested(Box::new(Expr::Cast {
+                expr: Box::new(Expr::Array(vec![Expr::Array(vec![
+                    num[2].clone(),
+                    num[3].clone(),
+                ])])),
+                data_type: DataType::Array(Box::new(DataType::Array(Box::new(DataType::Int(
+                    None
+                )))))
+            }))),
+            indexs: vec![num[1].clone(), num[2].clone()],
         },
         expr_from_projection(only(&select.projection)),
     );