diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 190f9bf8f..796a046a9 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -820,6 +820,14 @@ pub enum Subscript { /// {2,3,4,5} /// ``` /// + /// Stride notation is also supported + /// + /// ```plaintext + /// => select (array[1,2,3,4,5,6])[1:6:2]; + /// ----------- + /// {1,3,5} + /// ``` + /// /// The lower and/or upper bound can be omitted to slice from the start or /// end of the array respectively. /// @@ -827,6 +835,7 @@ pub enum Subscript { Slice { lower_bound: Option, upper_bound: Option, + stride: Option, }, } @@ -837,6 +846,7 @@ impl fmt::Display for Subscript { Subscript::Slice { lower_bound, upper_bound, + stride, } => { if let Some(lower) = lower_bound { write!(f, "{lower}")?; @@ -845,6 +855,10 @@ impl fmt::Display for Subscript { if let Some(upper) = upper_bound { write!(f, "{upper}")?; } + if let Some(stride) = stride { + write!(f, ":")?; + write!(f, "{stride}")?; + } Ok(()) } } diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 5e0d78e61..3986f81ed 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -2610,31 +2610,84 @@ impl<'a> Parser<'a> { } } - pub fn parse_subscript(&mut self, expr: Expr) -> Result { - let parse_upper_bound = |p: &mut Parser<'a>| { - if let Token::RBracket = p.peek_token().token { - Ok(None) - } else { - p.parse_expr().map(Some) - } + /// Parses an array subscript like + /// * `[:]` + /// * `[l]` + /// * `[l:]` + /// * `[:u]` + /// * `[l:u]` + /// * `[l:u:s]` + /// + /// Parser is right after `[` + fn parse_subscript_inner(&mut self) -> Result { + // at either `:(rest)` or `:(rest)]` + let lower_bound = if self.consume_token(&Token::Colon) { + None + } else { + Some(self.parse_expr()?) }; - let subscript = if self.consume_token(&Token::Colon) { - Subscript::Slice { - lower_bound: None, - upper_bound: parse_upper_bound(self)?, - } + + // check for end + if self.consume_token(&Token::RBracket) { + if let Some(lower_bound) = lower_bound { + return Ok(Subscript::Index { index: lower_bound }); + }; + return Ok(Subscript::Slice { + lower_bound, + upper_bound: None, + stride: None, + }); + } + + // consume the `:` + if lower_bound.is_some() { + self.expect_token(&Token::Colon)?; + } + + // we are now at either `]`, `(rest)]` + let upper_bound = if self.consume_token(&Token::RBracket) { + return Ok(Subscript::Slice { + lower_bound, + upper_bound: None, + stride: None, + }); } else { - let expr = self.parse_expr()?; - if self.consume_token(&Token::Colon) { - Subscript::Slice { - lower_bound: Some(expr), - upper_bound: parse_upper_bound(self)?, - } - } else { - Subscript::Index { index: expr } - } + Some(self.parse_expr()?) }; - self.expect_token(&Token::RBracket)?; + + // check for end + if self.consume_token(&Token::RBracket) { + return Ok(Subscript::Slice { + lower_bound, + upper_bound, + stride: None, + }); + } + + // we are now at `:]` or `:stride]` + self.expect_token(&Token::Colon)?; + let stride = if self.consume_token(&Token::RBracket) { + None + } else { + Some(self.parse_expr()?) + }; + + if stride.is_some() { + self.expect_token(&Token::RBracket)?; + } + + Ok(Subscript::Slice { + lower_bound, + upper_bound, + stride, + }) + } + + /// Parses an array subscript like `[1:3]` + /// + /// Parser is right after `[` + pub fn parse_subscript(&mut self, expr: Expr) -> Result { + let subscript = self.parse_subscript_inner()?; Ok(Expr::Subscript { expr: Box::new(expr), subscript: Box::new(subscript), diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index ad4165e79..43a5d2530 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -1992,6 +1992,15 @@ fn parse_array_subscript() { Subscript::Slice { lower_bound: Some(Expr::Value(number("2"))), upper_bound: Some(Expr::Value(number("5"))), + stride: None, + }, + ), + ( + "(ARRAY[1, 2, 3, 4, 5, 6])[2:5:3]", + Subscript::Slice { + lower_bound: Some(Expr::Value(number("2"))), + upper_bound: Some(Expr::Value(number("5"))), + stride: Some(Expr::Value(number("3"))), }, ), ( @@ -2007,6 +2016,7 @@ fn parse_array_subscript() { op: BinaryOperator::Minus, right: Box::new(Expr::Value(number("1"))), }), + stride: None, }, ), ( @@ -2014,6 +2024,7 @@ fn parse_array_subscript() { Subscript::Slice { lower_bound: None, upper_bound: Some(Expr::Value(number("5"))), + stride: None, }, ), ( @@ -2021,6 +2032,7 @@ fn parse_array_subscript() { Subscript::Slice { lower_bound: Some(Expr::Value(number("2"))), upper_bound: None, + stride: None, }, ), ( @@ -2028,6 +2040,7 @@ fn parse_array_subscript() { Subscript::Slice { lower_bound: None, upper_bound: None, + stride: None, }, ), ];