dsq_functions/builtin/
array_shift.rs

1use 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                // For testing, return original
24                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}