dsq_functions/builtin/
array_pop.rs

1use dsq_shared::value::{df_row_to_value, value_from_any_value, Value};
2use dsq_shared::Result;
3use polars::prelude::*;
4
5use crate::inventory;
6use crate::FunctionRegistration;
7
8pub fn builtin_array_pop(args: &[Value]) -> Result<Value> {
9    if args.len() != 1 {
10        return Err(dsq_shared::error::operation_error(
11            "array_pop() expects 1 argument",
12        ));
13    }
14
15    match &args[0] {
16        Value::Array(arr) => {
17            if arr.is_empty() {
18                Ok(Value::Null)
19            } else {
20                Ok(arr[arr.len() - 1].clone())
21            }
22        }
23        Value::DataFrame(df) => {
24            if df.height() == 0 {
25                Ok(Value::Null)
26            } else {
27                // Return last row as object using the robust conversion
28                let last_idx = df.height() - 1;
29                df_row_to_value(df, last_idx)
30            }
31        }
32        Value::Series(series) => {
33            if matches!(series.dtype(), DataType::List(_)) {
34                if series.len() == 1 {
35                    match series.list().unwrap().get_as_series(0) {
36                        Some(list_series) => {
37                            if !list_series.is_empty() {
38                                let last_idx = list_series.len() - 1;
39                                match list_series.get(last_idx) {
40                                    Ok(val) => Ok(value_from_any_value(val).unwrap_or(Value::Null)),
41                                    _ => Ok(Value::Null),
42                                }
43                            } else {
44                                Ok(Value::Null)
45                            }
46                        }
47                        _ => Ok(Value::Null),
48                    }
49                } else {
50                    Err(dsq_shared::error::operation_error(format!(
51                        "array_pop() on series with {} elements not supported",
52                        series.len()
53                    )))
54                }
55            } else {
56                Err(dsq_shared::error::operation_error(
57                    "array_pop() requires an array, DataFrame, or list series",
58                ))
59            }
60        }
61        _ => Err(dsq_shared::error::operation_error(
62            "array_pop() requires an array, DataFrame, or list series",
63        )),
64    }
65}
66
67inventory::submit! {
68    FunctionRegistration {
69        name: "array_pop",
70        func: builtin_array_pop,
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77    use dsq_shared::value::Value;
78
79    #[test]
80    fn test_builtin_array_pop_array() {
81        let arr = vec![Value::Int(1), Value::Int(2), Value::Int(3)];
82        let result = builtin_array_pop(&[Value::Array(arr)]).unwrap();
83        assert_eq!(result, Value::Int(3));
84    }
85
86    #[test]
87    fn test_builtin_array_pop_empty_array() {
88        let arr = vec![];
89        let result = builtin_array_pop(&[Value::Array(arr)]).unwrap();
90        assert_eq!(result, Value::Null);
91    }
92
93    #[test]
94    fn test_builtin_array_pop_dataframe() {
95        let s1 = Series::new(PlSmallStr::from("col1"), &[1i64, 2i64]);
96        let s2 = Series::new(PlSmallStr::from("col2"), &["a", "b"]);
97        let df = DataFrame::new(vec![s1.into(), s2.into()]).unwrap();
98        let result = builtin_array_pop(&[Value::DataFrame(df)]).unwrap();
99        match result {
100            Value::Object(obj) => {
101                assert_eq!(obj.get("col1"), Some(&Value::Int(2)));
102                assert_eq!(obj.get("col2"), Some(&Value::String("b".to_string())));
103            }
104            _ => panic!("Expected Object"),
105        }
106    }
107
108    #[test]
109    fn test_builtin_array_pop_empty_dataframe() {
110        let df = DataFrame::new(vec![Series::new(
111            PlSmallStr::from("empty"),
112            Vec::<String>::new(),
113        )
114        .into()])
115        .unwrap();
116        let result = builtin_array_pop(&[Value::DataFrame(df)]).unwrap();
117        assert_eq!(result, Value::Null);
118    }
119
120    #[test]
121    fn test_builtin_array_pop_series() {
122        let s1 = Series::new(PlSmallStr::from(""), &[1i64, 2i64, 3i64]);
123        let list_series = Series::new(PlSmallStr::from("list_col"), &[s1]);
124        let result = builtin_array_pop(&[Value::Series(list_series)]).unwrap();
125        assert_eq!(result, Value::Int(3));
126    }
127
128    #[test]
129    fn test_builtin_array_pop_empty_series() {
130        let s1 = Series::new(PlSmallStr::from(""), &[] as &[i64]);
131        let list_series = Series::new(PlSmallStr::from("list_col"), &[s1]);
132        let result = builtin_array_pop(&[Value::Series(list_series)]).unwrap();
133        assert_eq!(result, Value::Null);
134    }
135
136    #[test]
137    fn test_builtin_array_pop_invalid_args() {
138        let result = builtin_array_pop(&[]);
139        assert!(result.is_err());
140        assert!(result
141            .unwrap_err()
142            .to_string()
143            .contains("expects 1 argument"));
144
145        let result = builtin_array_pop(&[Value::Array(vec![Value::Int(1)]), Value::Int(2)]);
146        assert!(result.is_err());
147        assert!(result
148            .unwrap_err()
149            .to_string()
150            .contains("expects 1 argument"));
151
152        let result = builtin_array_pop(&[Value::String("not array".to_string())]);
153        assert!(result.is_err());
154        assert!(result
155            .unwrap_err()
156            .to_string()
157            .contains("requires an array"));
158    }
159
160    #[test]
161    fn test_array_pop_registered_via_inventory() {
162        use crate::BuiltinRegistry;
163        let registry = BuiltinRegistry::new();
164        assert!(registry.functions.contains_key("array_pop"));
165    }
166}
dsq_functions/builtin/array_pop.rs

dsq_functions/builtin/
array_pop.rs
