Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: support extended array slicing with step #12983

Closed
wants to merge 1 commit into from
Closed

feat: support extended array slicing with step #12983

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 27 additions & 7 deletions src/expr/impl/src/scalar/array_range_access.rs
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading