Add array::shuffle function (#3993)

surrealdb · May 15, 2024 · 97ce910 · 97ce910
1 parent 51aea4f
commit 97ce910
Show file tree

Hide file tree

Showing 7 changed files with 53 additions and 0 deletions.
diff --git a/core/src/fnc/array.rs b/core/src/fnc/array.rs
@@ -12,6 +12,8 @@ use crate::sql::array::Union;
 use crate::sql::array::Uniq;
 use crate::sql::value::Value;
 
+use rand::prelude::SliceRandom;
+
 pub fn add((mut array, value): (Array, Value)) -> Result<Value, Error> {
  match value {
  Value::Array(value) => {
@@ -331,6 +333,12 @@ pub fn reverse((mut array,): (Array,)) -> Result<Value, Error> {
  Ok(array.into())
 }
 
+pub fn shuffle((mut array,): (Array,)) -> Result<Value, Error> {
+ let mut rng = rand::thread_rng();
+ array.0.shuffle(&mut rng);
+ Ok(array.into())
+}
+
 pub fn slice((array, beg, lim): (Array, Option<isize>, Option<isize>)) -> Result<Value, Error> {
  let skip = match beg {
  Some(v) if v < 0 => array.len().saturating_sub(v.unsigned_abs()),

diff --git a/core/src/fnc/mod.rs b/core/src/fnc/mod.rs
@@ -125,6 +125,7 @@ pub fn synchronous(ctx: &Context<'_>, name: &str, args: Vec<Value>) -> Result<Va
  "array::push" => array::push,
  "array::remove" => array::remove,
  "array::reverse" => array::reverse,
+ "array::shuffle" => array::shuffle,
  "array::slice" => array::slice,
  "array::sort" => array::sort,
  "array::transpose" => array::transpose,

diff --git a/core/src/fnc/script/modules/surrealdb/functions/array.rs b/core/src/fnc/script/modules/surrealdb/functions/array.rs
@@ -45,6 +45,7 @@ impl_module_def!(
  "prepend" => run,
  "remove" => run,
  "reverse" => run,
+ "shuffle" => run,
  "slice" => run,
  "sort" => (sort::Package),
  "transpose" => run,

diff --git a/core/src/syn/parser/builtin.rs b/core/src/syn/parser/builtin.rs
@@ -115,6 +115,7 @@ pub(crate) static PATHS: phf::Map<UniCase<&'static str>, PathKind> = phf_map! {
  UniCase::ascii("array::push") => PathKind::Function,
  UniCase::ascii("array::remove") => PathKind::Function,
  UniCase::ascii("array::reverse") => PathKind::Function,
+ UniCase::ascii("array::shuffle") => PathKind::Function,
  UniCase::ascii("array::slice") => PathKind::Function,
  UniCase::ascii("array::sort") => PathKind::Function,
  UniCase::ascii("array::transpose") => PathKind::Function,

diff --git a/lib/fuzz/fuzz_targets/fuzz_executor.dict b/lib/fuzz/fuzz_targets/fuzz_executor.dict
@@ -184,6 +184,7 @@
 "array::push("
 "array::remove("
 "array::reverse("
+"array::shuffle("
 "array::slice("
 "array::sort("
 "array::sort::asc("

diff --git a/lib/fuzz/fuzz_targets/fuzz_sql_parser.dict b/lib/fuzz/fuzz_targets/fuzz_sql_parser.dict
@@ -184,6 +184,7 @@
 "array::push("
 "array::remove("
 "array::reverse("
+"array::shuffle("
 "array::slice("
 "array::sort("
 "array::sort::asc("

diff --git a/lib/tests/function.rs b/lib/tests/function.rs
@@ -1050,6 +1050,46 @@ async fn function_array_reverse() -> Result<(), Error> {
  Ok(())
 }
 
+#[tokio::test]
+async fn function_array_shuffle() -> Result<(), Error> {
+ let sql = r#"
+ RETURN array::shuffle([]);
+ RETURN array::shuffle(3);
+ RETURN array::shuffle([4]);
+ RETURN array::shuffle([1,1,1]);
+ RETURN array::shuffle([1,2,"text",3,3,4]); // find a way to check randomness
+ "#;
+ let dbs = new_ds().await?;
+ let ses = Session::owner().with_ns("test").with_db("test");
+ let res = &mut dbs.execute(sql, &ses, None).await?;
+ assert_eq!(res.len(), 5);
+ //
+ let tmp = res.remove(0).result?;
+ let val = Value::parse("[]");
+ assert_eq!(tmp, val);
+ //
+ let tmp = res.remove(0).result;
+ assert!(
+ matches!(
+ &tmp,
+ Err(e) if e.to_string() == "Incorrect arguments for function array::shuffle(). Argument 1 was the wrong type. Expected a array but found 3"
+ ),
+ "{tmp:?}"
+ );
+ //
+ let tmp = res.remove(0).result?;
+ let val = Value::parse("[4]");
+ assert_eq!(tmp, val);
+ //
+ let tmp = res.remove(0).result?;
+ let val = Value::parse("[1,1,1]");
+ assert_eq!(tmp, val);
+ //
+ let _ = res.remove(0).result?;
+ //
+ Ok(())
+}
+
 #[tokio::test]
 async fn function_array_slice() -> Result<(), Error> {
  let sql = r#"