Skip to content

Commit

Permalink
feat: support extended array slicing with step
Browse files Browse the repository at this point in the history
  • Loading branch information
@jetjinser
jetjinser committed Oct 20, 2023
1 parent fa66cbd commit f25a5fc
Show file tree
Hide file tree
Showing 5 changed files with 201 additions and 49 deletions.
34 changes: 27 additions & 7 deletions src/expr/impl/src/scalar/array_range_access.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,38 @@ use risingwave_expr::function;

/// If the case is `array[1,2,3][:2]`, then start will be 0 set by the frontend
/// If the case is `array[1,2,3][1:]`, then end will be `i32::MAX` set by the frontend
#[function("array_range_access(anyarray, int4, int4) -> anyarray")]
pub fn array_range_access(list: ListRef<'_>, start: i32, end: i32) -> Option<ListValue> {
#[function("array_range_access(anyarray, int4, int4, int4) -> anyarray")]
pub fn array_range_access(list: ListRef<'_>, start: i32, end: i32, step: i32) -> Option<ListValue> {
let mut data = vec![];
let list_all_values = list.iter();
let start = std::cmp::max(start, 1) as usize;
let end = std::cmp::min(std::cmp::max(0, end), list_all_values.len() as i32) as usize;
if start > end {

let mut new_start = std::cmp::max(start, 1) as usize;
let mut new_end = std::cmp::min(std::cmp::max(0, end), list_all_values.len() as i32) as usize;

let (step, need_rev) = if step < 0 {
if start != 0 && end != i32::MAX {
(new_start, new_end) = (new_end, new_start);
}
(-step, true)
} else {
(step, false)
};

if new_start > new_end {
return Some(ListValue::new(data));
}

for datumref in list_all_values.take(end).skip(start - 1) {
for datumref in list_all_values
.take(new_end)
.skip(new_start - 1)
.step_by(step as _)
{
data.push(datumref.to_owned_datum());
}
Some(ListValue::new(data))

Some(ListValue::new(if need_rev {
data.into_iter().rev().collect()
} else {
data
}))
}
9 changes: 6 additions & 3 deletions src/frontend/src/binder/expr/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,9 +122,12 @@ impl Binder {
Expr::Nested(expr) => self.bind_expr_inner(*expr),
Expr::Array(Array { elem: exprs, .. }) => self.bind_array(exprs),
Expr::ArrayIndex { obj, index } => self.bind_array_index(*obj, *index),
Expr::ArrayRangeIndex { obj, start, end } => {
self.bind_array_range_index(*obj, start, end)
}
Expr::ArrayRangeIndex {
obj,
start,
end,
step,
} => self.bind_array_range_index(*obj, start, end, step),
Expr::Function(f) => self.bind_function(f),
Expr::Subquery(q) => self.bind_subquery_expr(*q, SubqueryKind::Scalar),
Expr::Exists(q) => self.bind_subquery_expr(*q, SubqueryKind::Existential),
Expand Down
9 changes: 8 additions & 1 deletion src/frontend/src/binder/expr/value.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,7 @@ impl Binder {
obj: Expr,
start: Option<Box<Expr>>,
end: Option<Box<Expr>>,
step: Option<Box<Expr>>,
) -> Result<ExprImpl> {
let obj = self.bind_expr_inner(obj)?;
let start = match start {
Expand All @@ -169,10 +170,16 @@ impl Binder {
.bind_expr_inner(*expr)?
.cast_implicit(DataType::Int32)?,
};
let step = match step {
None => ExprImpl::literal_int(1),
Some(expr) => self
.bind_expr_inner(*expr)?
.cast_implicit(DataType::Int32)?,
};
match obj.return_type() {
DataType::List(return_type) => Ok(FunctionCall::new_unchecked(
ExprType::ArrayRangeAccess,
vec![obj, start, end],
vec![obj, start, end, step],
DataType::List(return_type),
)
.into()),
Expand Down
16 changes: 13 additions & 3 deletions src/sqlparser/src/ast/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -442,11 +442,12 @@ pub enum Expr {
obj: Box<Expr>,
index: Box<Expr>,
},
/// A slice expression `arr[1:3]`
/// A slice expression `arr[1:3:2]`
ArrayRangeIndex {
obj: Box<Expr>,
start: Option<Box<Expr>>,
end: Option<Box<Expr>>,
step: Option<Box<Expr>>,
},
LambdaFunction {
args: Vec<Ident>,
Expand Down Expand Up @@ -670,7 +671,12 @@ impl fmt::Display for Expr {
write!(f, "{}[{}]", obj, index)?;
Ok(())
}
Expr::ArrayRangeIndex { obj, start, end } => {
Expr::ArrayRangeIndex {
obj,
start,
end,
step,
} => {
let start_str = match start {
None => "".to_string(),
Some(start) => format!("{}", start),
Expand All @@ -679,7 +685,11 @@ impl fmt::Display for Expr {
None => "".to_string(),
Some(end) => format!("{}", end),
};
write!(f, "{}[{}:{}]", obj, start_str, end_str)?;
let step_str = match step {
None => "".to_string(),
Some(step) => format!(":{}", step),
};
write!(f, "{}[{}:{}{}]", obj, start_str, end_str, step_str)?;
Ok(())
}
Expr::Array(exprs) => write!(f, "{}", exprs),
Expand Down
182 changes: 147 additions & 35 deletions src/sqlparser/src/parser.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1540,63 +1540,175 @@ impl Parser {
}
}

/// We parse both `array[1,9][1]`, `array[1,9][1:2]`, `array[1,9][:2]`, `array[1,9][1:]` and
/// `array[1,9][:]` in this function.
pub fn parse_array_index(&mut self, expr: Expr) -> Result<Expr, ParserError> {
let new_expr = match self.peek_token().token {
fn parse_array_index_inner(&mut self, expr: Expr) -> Result<Expr, ParserError> {
// [N]
// [:], [N:], [:N], [N:M]
// [::], [N::], [:N:], [::N]
// [N:M:], [N::M], [:N:P], [N:M:P]
match self.peek_token().token {
// [:]
// [:N]
// [:N:M], [:N:]
Token::Colon => {
// [:] or [:N]
assert!(self.consume_token(&Token::Colon));
let end = match self.peek_token().token {
Token::RBracket => None,
match self.peek_token().token {
// [:]
Token::RBracket => Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: None,
end: None,
step: None,
}),
// [:N], [:N:]
// [:N:M]
_ => {
let end_index = Box::new(self.parse_expr()?);
Some(end_index)
let end = Some(Box::new(self.parse_expr()?));
match self.peek_token().token {
// [:N], [:N:]
Token::RBracket | Token::Colon => Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: None,
end,
step: None,
}),
// [:N:M]
_ => {
let step = Some(Box::new(self.parse_expr()?));
Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: None,
end,
step,
})
}
}
}
};
Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: None,
end,
}
}
// [::], [::N]
Token::DoubleColon => {
assert!(self.consume_token(&Token::DoubleColon));
match self.peek_token().token {
// [::]
Token::RBracket => Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: None,
end: None,
step: None,
}),
// [::N]
_ => {
let step = Some(Box::new(self.parse_expr()?));
Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: None,
end: None,
step,
})
}
}
}
// [N]
// [N:], [N:M]
// [N:M:P], [N:M:], [N::M]
// [N::]
_ => {
// [N], [N:], [N:M]
let index = Box::new(self.parse_expr()?);
match self.peek_token().token {
// [N:], [N:M]
// [N:M:P], [N:M:]
Token::Colon => {
// [N:], [N:M]
assert!(self.consume_token(&Token::Colon));

match self.peek_token().token {
Token::RBracket => {
// [N:]
Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: Some(index),
end: None,
}
}
// [N:]
Token::RBracket => Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: Some(index),
end: None,
step: None,
}),

// [N:M]
// [N:M:P], [N:M:]
_ => {
// [N:M]
let end = Some(Box::new(self.parse_expr()?));
Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: Some(index),
end,
match self.peek_token().token {
// [N:M:P], [N:M:]
Token::Colon => {
assert!(self.consume_token(&Token::Colon));
match self.peek_token().token {
// [N:M:]
Token::RBracket => Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: Some(index),
end,
step: None,
}),
// [N:M:P]
_ => {
let step = Some(Box::new(self.parse_expr()?));
Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: Some(index),
end,
step,
})
}
}
}
// [N:M]
_ => Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: Some(index),
end,
step: None,
}),
}
}
}
}
_ => {
// [N]
Expr::ArrayIndex {
obj: Box::new(expr),
index,
// [N::], [N::M]
Token::DoubleColon => {
assert!(self.consume_token(&Token::DoubleColon));
match self.peek_token().token {
// [N::]
Token::RBracket => Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: Some(index),
end: None,
step: None,
}),
// [N::M]
_ => {
let step = Some(Box::new(self.parse_expr()?));
Ok(Expr::ArrayRangeIndex {
obj: Box::new(expr),
start: Some(index),
end: None,
step,
})
}
}
}
// [N]
_ => Ok(Expr::ArrayIndex {
obj: Box::new(expr),
index,
}),
}
}
};
}
}

/// We parse both array index and slice in this function.
/// [N]
/// [:], [N:], [:N], [N:M]
/// [::], [N::], [:N:], [::N]
/// [N:M:], [N::M], [:N:P], [N:M:P]
pub fn parse_array_index(&mut self, expr: Expr) -> Result<Expr, ParserError> {
let new_expr = self.parse_array_index_inner(expr)?;

self.expect_token(&Token::RBracket)?;
// recursively checking for more indices
if self.consume_token(&Token::LBracket) {
Expand Down

0 comments on commit f25a5fc

Please sign in to comment.