From 27e17d7cdd37411b3646d3bfed339ff5f18e4651 Mon Sep 17 00:00:00 2001
From: Dave MacLeod <56599343+Dhghomon@users.noreply.github.com>
Date: Fri, 2 Aug 2024 21:50:40 +0900
Subject: [PATCH] Add array::windows() function (#4445)

---
 core/src/fnc/array.rs                         |  6 +++++
 core/src/fnc/mod.rs                           |  1 +
 .../modules/surrealdb/functions/array.rs      |  3 ++-
 core/src/sql/array.rs                         | 24 +++++++++++++++++++
 core/src/syn/parser/builtin.rs                |  3 ++-
 lib/tests/function.rs                         | 22 +++++++++++++++++
 6 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/core/src/fnc/array.rs b/core/src/fnc/array.rs
index 11837e08..197e981a 100644
--- a/core/src/fnc/array.rs
+++ b/core/src/fnc/array.rs
@@ -10,6 +10,7 @@ use crate::sql::array::Matches;
 use crate::sql::array::Transpose;
 use crate::sql::array::Union;
 use crate::sql::array::Uniq;
+use crate::sql::array::Windows;
 use crate::sql::value::Value;
 
 use rand::prelude::SliceRandom;
@@ -399,6 +400,11 @@ pub fn union((array, other): (Array, Array)) -> Result<Value, Error> {
 	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 {
 
 	use crate::err::Error;
diff --git a/core/src/fnc/mod.rs b/core/src/fnc/mod.rs
index 24f8859c..2c77381c 100644
--- a/core/src/fnc/mod.rs
+++ b/core/src/fnc/mod.rs
@@ -137,6 +137,7 @@ pub fn synchronous(
 		"array::union" => array::union,
 		"array::sort::asc" => array::sort::asc,
 		"array::sort::desc" => array::sort::desc,
+		"array::windows" => array::windows,
 		//
 		"bytes::len" => bytes::len,
 		//
diff --git a/core/src/fnc/script/modules/surrealdb/functions/array.rs b/core/src/fnc/script/modules/surrealdb/functions/array.rs
index 7dbddb37..9bf1600c 100644
--- a/core/src/fnc/script/modules/surrealdb/functions/array.rs
+++ b/core/src/fnc/script/modules/surrealdb/functions/array.rs
@@ -49,5 +49,6 @@ impl_module_def!(
 	"slice" => run,
 	"sort" => (sort::Package),
 	"transpose" => run,
-	"union" => run
+	"union" => run,
+	"windows" => run
 );
diff --git a/core/src/sql/array.rs b/core/src/sql/array.rs
index 144dd762..e271f498 100644
--- a/core/src/sql/array.rs
+++ b/core/src/sql/array.rs
@@ -498,3 +498,27 @@ impl Uniq<Array> for Array {
 		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())
+	}
+}
diff --git a/core/src/syn/parser/builtin.rs b/core/src/syn/parser/builtin.rs
index 3b6d4663..987d85d6 100644
--- a/core/src/syn/parser/builtin.rs
+++ b/core/src/syn/parser/builtin.rs
@@ -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::sort::asc") => PathKind::Function,
 		UniCase::ascii("array::sort::desc") => PathKind::Function,
+		UniCase::ascii("array::windows") => PathKind::Function,
 		//
 		UniCase::ascii("object::entries") => PathKind::Function,
 		UniCase::ascii("object::from_entries") => PathKind::Function,
@@ -461,7 +462,7 @@ impl Parser<'_> {
 				.await
 				.map(|x| Value::Function(Box::new(x))),
 			None => {
-				// Generate an suggestion.
+				// Generate a suggestion.
 				// don't search further if the levenshtein distance is further then 10.
 				let mut cut_off = MAX_LEVENSTHEIN_CUT_OFF;
 
diff --git a/lib/tests/function.rs b/lib/tests/function.rs
index 70f91ab3..5c24fe36 100644
--- a/lib/tests/function.rs
+++ b/lib/tests/function.rs
@@ -1036,6 +1036,28 @@ async fn function_array_union() -> Result<(), Error> {
 	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
 // --------------------------------------------------