Skip to content

Commit

Permalink
Support optional stride in array access
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed May 30, 2024
1 parent bd573b2 commit 4df6dc8
Show file tree
Hide file tree
Showing 3 changed files with 102 additions and 22 deletions.
14 changes: 14 additions & 0 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
},
}

Expand All @@ -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}")?;
Expand All @@ -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(())
}
}
Expand Down
97 changes: 75 additions & 22 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
13 changes: 13 additions & 0 deletions tests/sqlparser_postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"))),
},
),
(
Expand All @@ -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,
},
),
];
Expand Down

0 comments on commit 4df6dc8

Please sign in to comment.