Add array::windows() function (#4445)

This commit is contained in:
Dave MacLeod 2024-08-02 21:50:40 +09:00 committed by GitHub
parent e8f55b60f2
commit 27e17d7cdd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 57 additions and 2 deletions

View file

@ -10,6 +10,7 @@ use crate::sql::array::Matches;
use crate::sql::array::Transpose; use crate::sql::array::Transpose;
use crate::sql::array::Union; use crate::sql::array::Union;
use crate::sql::array::Uniq; use crate::sql::array::Uniq;
use crate::sql::array::Windows;
use crate::sql::value::Value; use crate::sql::value::Value;
use rand::prelude::SliceRandom; use rand::prelude::SliceRandom;
@ -399,6 +400,11 @@ pub fn union((array, other): (Array, Array)) -> Result<Value, Error> {
Ok(array.union(other).into()) Ok(array.union(other).into())
} }
pub fn windows((array, window_size): (Array, i64)) -> Result<Value, Error> {
let window_size = window_size.max(0) as usize;
Ok(array.windows(window_size)?.into())
}
pub mod sort { pub mod sort {
use crate::err::Error; use crate::err::Error;

View file

@ -137,6 +137,7 @@ pub fn synchronous(
"array::union" => array::union, "array::union" => array::union,
"array::sort::asc" => array::sort::asc, "array::sort::asc" => array::sort::asc,
"array::sort::desc" => array::sort::desc, "array::sort::desc" => array::sort::desc,
"array::windows" => array::windows,
// //
"bytes::len" => bytes::len, "bytes::len" => bytes::len,
// //

View file

@ -49,5 +49,6 @@ impl_module_def!(
"slice" => run, "slice" => run,
"sort" => (sort::Package), "sort" => (sort::Package),
"transpose" => run, "transpose" => run,
"union" => run "union" => run,
"windows" => run
); );

View file

@ -498,3 +498,27 @@ impl Uniq<Array> for Array {
self self
} }
} }
// ------------------------------
pub(crate) trait Windows<T> {
fn windows(self, window_size: usize) -> Result<T, Error>;
}
impl Windows<Array> for Array {
fn windows(self, window_size: usize) -> Result<Array, Error> {
if window_size < 1 {
return Err(Error::InvalidArguments {
name: "array::windows".to_string(),
message: "The second argument must be an integer greater than 0".to_string(),
});
}
Ok(self
.0
.windows(window_size)
.map::<Value, _>(|chunk| chunk.to_vec().into())
.collect::<Vec<_>>()
.into())
}
}

View file

@ -125,6 +125,7 @@ pub(crate) static PATHS: phf::Map<UniCase<&'static str>, PathKind> = phf_map! {
UniCase::ascii("array::union") => PathKind::Function, UniCase::ascii("array::union") => PathKind::Function,
UniCase::ascii("array::sort::asc") => PathKind::Function, UniCase::ascii("array::sort::asc") => PathKind::Function,
UniCase::ascii("array::sort::desc") => PathKind::Function, UniCase::ascii("array::sort::desc") => PathKind::Function,
UniCase::ascii("array::windows") => PathKind::Function,
// //
UniCase::ascii("object::entries") => PathKind::Function, UniCase::ascii("object::entries") => PathKind::Function,
UniCase::ascii("object::from_entries") => PathKind::Function, UniCase::ascii("object::from_entries") => PathKind::Function,
@ -461,7 +462,7 @@ impl Parser<'_> {
.await .await
.map(|x| Value::Function(Box::new(x))), .map(|x| Value::Function(Box::new(x))),
None => { None => {
// Generate an suggestion. // Generate a suggestion.
// don't search further if the levenshtein distance is further then 10. // don't search further if the levenshtein distance is further then 10.
let mut cut_off = MAX_LEVENSTHEIN_CUT_OFF; let mut cut_off = MAX_LEVENSTHEIN_CUT_OFF;

View file

@ -1036,6 +1036,28 @@ async fn function_array_union() -> Result<(), Error> {
Ok(()) Ok(())
} }
#[tokio::test]
async fn function_array_windows() -> Result<(), Error> {
let sql = r#"
RETURN array::windows([0, 1, 2, 3], 2);
RETURN array::windows([0, 1, 2], 2);
RETURN array::windows([0, 1, 2], 3);
RETURN array::windows([0, 1, 2, 3, 4, 5], 3);
RETURN array::windows([0, 1, 2], 4);
RETURN array::windows([0, 1, 2], 0);
"#;
let error = "Incorrect arguments for function array::windows(). The second argument must be an integer greater than 0";
Test::new(sql)
.await?
.expect_val("[[0, 1], [1, 2], [2, 3]]")?
.expect_val("[[0, 1], [1, 2]]")?
.expect_val("[[0, 1, 2]]")?
.expect_val("[[0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5]]")?
.expect_val("[]")?
.expect_error(error)?;
Ok(())
}
// -------------------------------------------------- // --------------------------------------------------
// bytes // bytes
// -------------------------------------------------- // --------------------------------------------------