Skip to main content

robinpath_modules/modules/
table_mod.rs

1use robinpath::{RobinPath, Value};
2
3pub fn register(rp: &mut RobinPath) {
4    rp.register_builtin("table.create", |args, _| {
5        let data = args.first().cloned().unwrap_or(Value::Null);
6        match &data {
7            Value::Array(_) => Ok(data),
8            Value::Object(obj) => {
9                // {columns: [...], rows: [[...], ...]} format
10                let columns = obj.get("columns").cloned().unwrap_or(Value::Array(vec![]));
11                let rows = obj.get("rows").cloned().unwrap_or(Value::Array(vec![]));
12                if let (Value::Array(cols), Value::Array(row_data)) = (&columns, &rows) {
13                    let col_names: Vec<String> = cols.iter().map(|c| c.to_display_string()).collect();
14                    let result: Vec<Value> = row_data
15                        .iter()
16                        .map(|row| {
17                            if let Value::Array(cells) = row {
18                                let mut obj = indexmap::IndexMap::new();
19                                for (i, name) in col_names.iter().enumerate() {
20                                    obj.insert(
21                                        name.clone(),
22                                        cells.get(i).cloned().unwrap_or(Value::Null),
23                                    );
24                                }
25                                Value::Object(obj)
26                            } else {
27                                row.clone()
28                            }
29                        })
30                        .collect();
31                    Ok(Value::Array(result))
32                } else {
33                    Ok(Value::Array(vec![]))
34                }
35            }
36            _ => Ok(Value::Array(vec![])),
37        }
38    });
39
40    rp.register_builtin("table.select", |args, _| {
41        let table = args.first().cloned().unwrap_or(Value::Null);
42        let columns = args.get(1).cloned().unwrap_or(Value::Null);
43        if let (Value::Array(rows), Value::Array(cols)) = (&table, &columns) {
44            let col_names: Vec<String> = cols.iter().map(|c| c.to_display_string()).collect();
45            let result: Vec<Value> = rows
46                .iter()
47                .map(|row| {
48                    if let Value::Object(obj) = row {
49                        let mut new_obj = indexmap::IndexMap::new();
50                        for name in &col_names {
51                            if let Some(v) = obj.get(name) {
52                                new_obj.insert(name.clone(), v.clone());
53                            }
54                        }
55                        Value::Object(new_obj)
56                    } else {
57                        row.clone()
58                    }
59                })
60                .collect();
61            Ok(Value::Array(result))
62        } else {
63            Ok(Value::Array(vec![]))
64        }
65    });
66
67    rp.register_builtin("table.where", |args, _| {
68        let table = args.first().cloned().unwrap_or(Value::Null);
69        let field = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
70        let value = args.get(2).cloned().unwrap_or(Value::Null);
71        if let Value::Array(rows) = &table {
72            let result: Vec<Value> = rows
73                .iter()
74                .filter(|row| {
75                    if let Value::Object(obj) = row {
76                        obj.get(&field).map_or(false, |v| v.deep_eq(&value))
77                    } else {
78                        false
79                    }
80                })
81                .cloned()
82                .collect();
83            Ok(Value::Array(result))
84        } else {
85            Ok(Value::Array(vec![]))
86        }
87    });
88
89    rp.register_builtin("table.orderBy", |args, _| {
90        let table = args.first().cloned().unwrap_or(Value::Null);
91        let field = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
92        let order = args.get(2).map(|v| v.to_display_string()).unwrap_or_else(|| "asc".to_string());
93        let desc = order == "desc";
94        if let Value::Array(mut rows) = table {
95            rows.sort_by(|a, b| {
96                let va = get_field(a, &field);
97                let vb = get_field(b, &field);
98                let cmp = va.to_number().partial_cmp(&vb.to_number()).unwrap_or(std::cmp::Ordering::Equal);
99                if desc { cmp.reverse() } else { cmp }
100            });
101            Ok(Value::Array(rows))
102        } else {
103            Ok(Value::Array(vec![]))
104        }
105    });
106
107    rp.register_builtin("table.groupBy", |args, _| {
108        let table = args.first().cloned().unwrap_or(Value::Null);
109        let field = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
110        if let Value::Array(rows) = &table {
111            let mut groups: indexmap::IndexMap<String, Vec<Value>> = indexmap::IndexMap::new();
112            for row in rows {
113                let key = get_field(row, &field).to_display_string();
114                groups.entry(key).or_default().push(row.clone());
115            }
116            let result: indexmap::IndexMap<String, Value> = groups
117                .into_iter()
118                .map(|(k, v)| (k, Value::Array(v)))
119                .collect();
120            Ok(Value::Object(result))
121        } else {
122            Ok(Value::Object(indexmap::IndexMap::new()))
123        }
124    });
125
126    rp.register_builtin("table.distinct", |args, _| {
127        let table = args.first().cloned().unwrap_or(Value::Null);
128        if let Value::Array(rows) = &table {
129            let mut seen = Vec::new();
130            let mut result = Vec::new();
131            for row in rows {
132                let key = row.to_json_string();
133                if !seen.contains(&key) {
134                    seen.push(key);
135                    result.push(row.clone());
136                }
137            }
138            Ok(Value::Array(result))
139        } else {
140            Ok(Value::Array(vec![]))
141        }
142    });
143
144    rp.register_builtin("table.limit", |args, _| {
145        let table = args.first().cloned().unwrap_or(Value::Null);
146        let n = args.get(1).map(|v| v.to_number() as usize).unwrap_or(10);
147        if let Value::Array(rows) = &table {
148            Ok(Value::Array(rows.iter().take(n).cloned().collect()))
149        } else {
150            Ok(Value::Array(vec![]))
151        }
152    });
153
154    rp.register_builtin("table.offset", |args, _| {
155        let table = args.first().cloned().unwrap_or(Value::Null);
156        let n = args.get(1).map(|v| v.to_number() as usize).unwrap_or(0);
157        if let Value::Array(rows) = &table {
158            Ok(Value::Array(rows.iter().skip(n).cloned().collect()))
159        } else {
160            Ok(Value::Array(vec![]))
161        }
162    });
163
164    rp.register_builtin("table.addColumn", |args, _| {
165        let table = args.first().cloned().unwrap_or(Value::Null);
166        let col_name = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
167        let default = args.get(2).cloned().unwrap_or(Value::Null);
168        if let Value::Array(rows) = &table {
169            let result: Vec<Value> = rows
170                .iter()
171                .map(|row| {
172                    if let Value::Object(mut obj) = row.clone() {
173                        obj.insert(col_name.clone(), default.clone());
174                        Value::Object(obj)
175                    } else {
176                        row.clone()
177                    }
178                })
179                .collect();
180            Ok(Value::Array(result))
181        } else {
182            Ok(Value::Array(vec![]))
183        }
184    });
185
186    rp.register_builtin("table.removeColumn", |args, _| {
187        let table = args.first().cloned().unwrap_or(Value::Null);
188        let col_name = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
189        if let Value::Array(rows) = &table {
190            let result: Vec<Value> = rows
191                .iter()
192                .map(|row| {
193                    if let Value::Object(mut obj) = row.clone() {
194                        obj.shift_remove(&col_name);
195                        Value::Object(obj)
196                    } else {
197                        row.clone()
198                    }
199                })
200                .collect();
201            Ok(Value::Array(result))
202        } else {
203            Ok(Value::Array(vec![]))
204        }
205    });
206
207    rp.register_builtin("table.renameColumn", |args, _| {
208        let table = args.first().cloned().unwrap_or(Value::Null);
209        let old_name = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
210        let new_name = args.get(2).map(|v| v.to_display_string()).unwrap_or_default();
211        if let Value::Array(rows) = &table {
212            let result: Vec<Value> = rows
213                .iter()
214                .map(|row| {
215                    if let Value::Object(obj) = row {
216                        let mut new_obj = indexmap::IndexMap::new();
217                        for (k, v) in obj {
218                            if k == &old_name {
219                                new_obj.insert(new_name.clone(), v.clone());
220                            } else {
221                                new_obj.insert(k.clone(), v.clone());
222                            }
223                        }
224                        Value::Object(new_obj)
225                    } else {
226                        row.clone()
227                    }
228                })
229                .collect();
230            Ok(Value::Array(result))
231        } else {
232            Ok(Value::Array(vec![]))
233        }
234    });
235
236    rp.register_builtin("table.count", |args, _| {
237        let table = args.first().cloned().unwrap_or(Value::Null);
238        if let Value::Array(rows) = &table {
239            Ok(Value::Number(rows.len() as f64))
240        } else {
241            Ok(Value::Number(0.0))
242        }
243    });
244
245    rp.register_builtin("table.sum", |args, _| {
246        let table = args.first().cloned().unwrap_or(Value::Null);
247        let field = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
248        if let Value::Array(rows) = &table {
249            let sum: f64 = rows.iter().map(|r| get_field(r, &field).to_number()).sum();
250            Ok(Value::Number(sum))
251        } else {
252            Ok(Value::Number(0.0))
253        }
254    });
255
256    rp.register_builtin("table.avg", |args, _| {
257        let table = args.first().cloned().unwrap_or(Value::Null);
258        let field = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
259        if let Value::Array(rows) = &table {
260            if rows.is_empty() {
261                return Ok(Value::Number(0.0));
262            }
263            let sum: f64 = rows.iter().map(|r| get_field(r, &field).to_number()).sum();
264            Ok(Value::Number(sum / rows.len() as f64))
265        } else {
266            Ok(Value::Number(0.0))
267        }
268    });
269
270    rp.register_builtin("table.min", |args, _| {
271        let table = args.first().cloned().unwrap_or(Value::Null);
272        let field = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
273        if let Value::Array(rows) = &table {
274            let min = rows.iter().map(|r| get_field(r, &field).to_number()).fold(f64::INFINITY, f64::min);
275            Ok(Value::Number(min))
276        } else {
277            Ok(Value::Null)
278        }
279    });
280
281    rp.register_builtin("table.max", |args, _| {
282        let table = args.first().cloned().unwrap_or(Value::Null);
283        let field = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
284        if let Value::Array(rows) = &table {
285            let max = rows.iter().map(|r| get_field(r, &field).to_number()).fold(f64::NEG_INFINITY, f64::max);
286            Ok(Value::Number(max))
287        } else {
288            Ok(Value::Null)
289        }
290    });
291
292    rp.register_builtin("table.head", |args, _| {
293        let table = args.first().cloned().unwrap_or(Value::Null);
294        let n = args.get(1).map(|v| v.to_number() as usize).unwrap_or(5);
295        if let Value::Array(rows) = &table {
296            Ok(Value::Array(rows.iter().take(n).cloned().collect()))
297        } else {
298            Ok(Value::Array(vec![]))
299        }
300    });
301
302    rp.register_builtin("table.tail", |args, _| {
303        let table = args.first().cloned().unwrap_or(Value::Null);
304        let n = args.get(1).map(|v| v.to_number() as usize).unwrap_or(5);
305        if let Value::Array(rows) = &table {
306            let skip = rows.len().saturating_sub(n);
307            Ok(Value::Array(rows.iter().skip(skip).cloned().collect()))
308        } else {
309            Ok(Value::Array(vec![]))
310        }
311    });
312
313    rp.register_builtin("table.columns", |args, _| {
314        let table = args.first().cloned().unwrap_or(Value::Null);
315        if let Value::Array(rows) = &table {
316            if let Some(Value::Object(first)) = rows.first() {
317                let cols: Vec<Value> = first.keys().map(|k| Value::String(k.clone())).collect();
318                Ok(Value::Array(cols))
319            } else {
320                Ok(Value::Array(vec![]))
321            }
322        } else {
323            Ok(Value::Array(vec![]))
324        }
325    });
326
327    rp.register_builtin("table.shape", |args, _| {
328        let table = args.first().cloned().unwrap_or(Value::Null);
329        if let Value::Array(rows) = &table {
330            let row_count = rows.len();
331            let col_count = if let Some(Value::Object(first)) = rows.first() {
332                first.len()
333            } else {
334                0
335            };
336            let mut obj = indexmap::IndexMap::new();
337            obj.insert("rows".to_string(), Value::Number(row_count as f64));
338            obj.insert("columns".to_string(), Value::Number(col_count as f64));
339            Ok(Value::Object(obj))
340        } else {
341            let mut obj = indexmap::IndexMap::new();
342            obj.insert("rows".to_string(), Value::Number(0.0));
343            obj.insert("columns".to_string(), Value::Number(0.0));
344            Ok(Value::Object(obj))
345        }
346    });
347
348    rp.register_builtin("table.join", |args, _| {
349        let left = args.first().cloned().unwrap_or(Value::Null);
350        let right = args.get(1).cloned().unwrap_or(Value::Null);
351        let on = args.get(2).map(|v| v.to_display_string()).unwrap_or_default();
352        if let (Value::Array(left_rows), Value::Array(right_rows)) = (&left, &right) {
353            let mut result = Vec::new();
354            for l_row in left_rows {
355                let l_key = get_field(l_row, &on);
356                for r_row in right_rows {
357                    let r_key = get_field(r_row, &on);
358                    if l_key.deep_eq(&r_key) {
359                        if let (Value::Object(l_obj), Value::Object(r_obj)) = (l_row, r_row) {
360                            let mut merged = l_obj.clone();
361                            for (k, v) in r_obj {
362                                if !merged.contains_key(k) {
363                                    merged.insert(k.clone(), v.clone());
364                                }
365                            }
366                            result.push(Value::Object(merged));
367                        }
368                    }
369                }
370            }
371            Ok(Value::Array(result))
372        } else {
373            Ok(Value::Array(vec![]))
374        }
375    });
376}
377
378fn get_field(val: &Value, field: &str) -> Value {
379    if let Value::Object(obj) = val {
380        obj.get(field).cloned().unwrap_or(Value::Null)
381    } else {
382        Value::Null
383    }
384}