dsq_functions/builtin/
array_shift.rs1use dsq_shared::value::Value;
2use dsq_shared::Result;
3use inventory;
4use polars::prelude::*;
5
6pub fn builtin_array_shift(args: &[Value]) -> Result<Value> {
7 if args.len() != 1 {
8 return Err(dsq_shared::error::operation_error(
9 "array_shift() expects 1 argument",
10 ));
11 }
12
13 match &args[0] {
14 Value::Array(arr) => {
15 if arr.is_empty() {
16 Ok(Value::Array(Vec::new()))
17 } else {
18 Ok(Value::Array(arr[1..].to_vec()))
19 }
20 }
21 Value::Series(series) => {
22 if matches!(series.dtype(), DataType::List(_)) {
23 Ok(Value::Series(series.clone()))
25 } else {
26 Err(dsq_shared::error::operation_error(
27 "array_shift() requires an array or list series",
28 ))
29 }
30 }
31 _ => Err(dsq_shared::error::operation_error(
32 "array_shift() requires an array or list series",
33 )),
34 }
35}
36
37inventory::submit! {
38 crate::FunctionRegistration {
39 name: "array_shift",
40 func: builtin_array_shift,
41 }
42}
43
44#[cfg(test)]
45mod tests {
46 use super::*;
47 use dsq_shared::value::Value;
48
49 #[test]
50 fn test_array_shift_empty() {
51 let args = vec![Value::Array(vec![])];
52 let result = builtin_array_shift(&args);
53 assert!(result.is_ok());
54 assert_eq!(result.unwrap(), Value::Array(vec![]));
55 }
56
57 #[test]
58 fn test_array_shift_non_empty() {
59 let args = vec![Value::Array(vec![
60 Value::Int(1),
61 Value::Int(2),
62 Value::Int(3),
63 ])];
64 let result = builtin_array_shift(&args);
65 assert!(result.is_ok());
66 assert_eq!(
67 result.unwrap(),
68 Value::Array(vec![Value::Int(2), Value::Int(3),])
69 );
70 }
71
72 #[test]
73 fn test_array_shift_wrong_args() {
74 let args = vec![];
75 let result = builtin_array_shift(&args);
76 assert!(result.is_err());
77 }
78
79 #[test]
80 fn test_array_shift_invalid_input() {
81 let args = vec![Value::String("not an array".to_string())];
82 let result = builtin_array_shift(&args);
83 assert!(result.is_err());
84 }
85}