Support optional stride in array access

Author: alamb

src/ast/mod.rs

@@ -820,13 +820,22 @@ 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.
     ///
     /// See <https://www.postgresql.org/docs/current/arrays.html#ARRAYS-ACCESSING>.
     Slice {
         lower_bound: Option<Expr>,
         upper_bound: Option<Expr>,
+        stride: Option<Expr>,
     },
 }
 
@@ -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(())
             }
         }

src/parser/mod.rs

@@ -2610,31 +2610,84 @@ impl<'a> Parser<'a> {
         }
     }
 
-    pub fn parse_subscript(&mut self, expr: Expr) -> Result<Expr, ParserError> {
-        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<Subscript, ParserError> {
+        // at either `<lower>:(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 `]`, `<upper>(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<Expr, ParserError> {
+        let subscript = self.parse_subscript_inner()?;
         Ok(Expr::Subscript {
             expr: Box::new(expr),
             subscript: Box::new(subscript),

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,27 +2016,31 @@ fn parse_array_subscript() {
                     op: BinaryOperator::Minus,
                     right: Box::new(Expr::Value(number("1"))),
                 }),
+                stride: None,
             },
         ),
         (
             "(ARRAY[1, 2, 3, 4, 5, 6])[:5]",
             Subscript::Slice {
                 lower_bound: None,
                 upper_bound: Some(Expr::Value(number("5"))),
+                stride: None,
             },
         ),
         (
             "(ARRAY[1, 2, 3, 4, 5, 6])[2:]",
             Subscript::Slice {
                 lower_bound: Some(Expr::Value(number("2"))),
                 upper_bound: None,
+                stride: None,
             },
         ),
         (
             "(ARRAY[1, 2, 3, 4, 5, 6])[:]",
             Subscript::Slice {
                 lower_bound: None,
                 upper_bound: None,
+                stride: None,
             },
         ),
     ];