From 3e4b69ee4944cdb7813125ae31f36ea40c67e885 Mon Sep 17 00:00:00 2001
From: Joe Bellus <joe@5sigma.io>
Date: Wed, 19 Oct 2022 20:24:17 -0400
Subject: [PATCH] Added Simple Series Aggregation Functions

Added Series.max, Series.min, Series.head, Series.tail, Series.mean

Added Range collection

(start..end).collect() now produces an array
---
 abacus-core/src/dataframe.rs | 74 +++++++++++++++++++++++++++++++++++-
 1 file changed, 73 insertions(+), 1 deletion(-)

diff --git a/abacus-core/src/dataframe.rs b/abacus-core/src/dataframe.rs
index d8edf37..c4d1675 100644
--- a/abacus-core/src/dataframe.rs
+++ b/abacus-core/src/dataframe.rs
@@ -25,8 +25,14 @@ pub fn setup_engine(engine: &mut rhai::Engine) {
     engine.register_fn("to_series", script_functions::to_series_unnamed);
     engine.register_fn("series", script_functions::series_unnamed);
     engine.register_fn("head", Series::s_head);
+    engine.register_fn("head", Series::s_head_len);
+    engine.register_fn("tail", Series::s_tail);
+    engine.register_fn("tail", Series::s_tail_len);
     engine.register_fn("sort", Series::s_sort);
     engine.register_fn("sum", Series::s_sum);
+    engine.register_fn("mean", Series::s_mean);
+    engine.register_fn("min", Series::s_min);
+    engine.register_fn("max", Series::s_max);
     engine.register_fn("add", Series::s_op_add);
 
     engine.register_fn("+", Series::add);
@@ -67,6 +73,7 @@ pub fn setup_engine(engine: &mut rhai::Engine) {
     engine.register_fn("min", script_functions::min);
     engine.register_fn("max", script_functions::max);
     engine.register_fn("first", script_functions::first);
+    engine.register_fn("collect", script_functions::range_to_array);
     let _ = engine.register_custom_operator("gt", 200);
     let _ = engine.register_custom_operator("gte", 200);
     let _ = engine.register_custom_operator("<<", 200);
@@ -217,9 +224,22 @@ impl Add for Series {
 }
 
 impl Series {
-    pub fn s_head(&mut self, n: i64) -> Series {
+    pub fn s_head(&mut self) -> Series {
+        Series(self.0.head(Some(1)))
+    }
+
+    pub fn s_head_len(&mut self, n: i64) -> Series {
         Series(self.0.head(Some(n as usize)))
     }
+
+    pub fn s_tail(&mut self) -> Series {
+        Series(self.0.tail(Some(self.0.len() - 1)))
+    }
+
+    pub fn s_tail_len(&mut self, size: i64) -> Series {
+        Series(self.0.tail(Some(size as usize)))
+    }
+
     pub fn s_sort(&mut self, reverse: bool) -> Series {
         Series(self.0.sort(reverse))
     }
@@ -228,6 +248,18 @@ impl Series {
         self.0.sum().unwrap_or_default()
     }
 
+    pub fn s_mean(&mut self) -> f64 {
+        self.0.mean().unwrap_or_default()
+    }
+
+    pub fn s_min(&mut self) -> f64 {
+        self.0.min().unwrap_or_default()
+    }
+
+    pub fn s_max(&mut self) -> f64 {
+        self.0.max().unwrap_or_default()
+    }
+
     pub fn s_get(&mut self, n: i64) -> ScriptResult<rhai::Dynamic> {
         let value = self.get(n as usize);
         match value {
@@ -436,6 +468,10 @@ mod script_functions {
     pub fn gte_op(a: &str, b: rhai::Dynamic) -> DataFrameExpression {
         DataFrameExpression(polars::prelude::col(a).gt_eq(DataFrameExpression::from(b)))
     }
+
+    pub fn range_to_array(range: std::ops::Range<i64>) -> rhai::Array {
+        range.map(Dynamic::from).collect()
+    }
 }
 
 fn implementation_df_select(
@@ -690,4 +726,40 @@ from data ["age"] : ["age" gt 18];
             .collect::<Vec<_>>();
         assert_eq!(s, vec![Some(22), Some(32)]);
     }
+
+    #[test]
+    pub fn test_range_to_array() {
+        let res = process(r#"series((1..4).collect())"#);
+        let s = res
+            .into_series()
+            .i64()
+            .unwrap()
+            .into_iter()
+            .collect::<Vec<_>>();
+        assert_eq!(s, vec![Some(1), Some(2), Some(3)]);
+    }
+
+    #[test]
+    pub fn test_series_min() {
+        let res = process(r#"series((2..4).collect()).min()"#);
+        dbg!(&res);
+        let s = res.into_scalar().as_float().unwrap();
+        assert_eq!(s, 2.0);
+    }
+
+    #[test]
+    pub fn test_series_max() {
+        let res = process(r#"series((2..4).collect()).max()"#);
+        dbg!(&res);
+        let s = res.into_scalar().as_float().unwrap();
+        assert_eq!(s, 3.0);
+    }
+
+    #[test]
+    pub fn test_series_mean() {
+        let res = process(r#"series((1..51).collect()).mean()"#);
+        dbg!(&res);
+        let s = res.into_scalar().as_float().unwrap();
+        assert_eq!(s, 25.5);
+    }
 }