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
This commit is contained in:
Joe Bellus 2022-10-19 20:24:17 -04:00
parent 5943876569
commit 3e4b69ee49
1 changed files with 73 additions and 1 deletions

View File

@ -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);
}
}