ext/polars/src/series.rs in polars-df-0.1.2 vs ext/polars/src/series.rs in polars-df-0.1.3

- old
+ new

@@ -3,10 +3,11 @@ use polars::prelude::*; use polars::series::IsSorted; use std::cell::RefCell; use crate::conversion::*; +use crate::set::set_at_idx; use crate::{RbDataFrame, RbPolarsErr, RbResult, RbValueError}; #[magnus::wrap(class = "Polars::RbSeries")] pub struct RbSeries { pub series: RefCell<Series>, @@ -472,10 +473,16 @@ s.into_iter().collect() } else if let Ok(s) = series.bool() { s.into_iter().collect() } else if let Ok(s) = series.utf8() { s.into_iter().collect() + } else if let Ok(_s) = series.date() { + let a = RArray::with_capacity(series.len()); + for v in series.iter() { + a.push::<Value>(Wrap(v).into()).unwrap(); + } + a } else { unimplemented!(); } } @@ -592,46 +599,192 @@ } else { None } } - // dispatch dynamically in future? - - pub fn cumsum(&self, reverse: bool) -> Self { - self.series.borrow().cumsum(reverse).into() + pub fn set_at_idx(&self, idx: &RbSeries, values: &RbSeries) -> RbResult<()> { + let mut s = self.series.borrow_mut(); + match set_at_idx(s.clone(), &idx.series.borrow(), &values.series.borrow()) { + Ok(out) => { + *s = out; + Ok(()) + } + Err(e) => Err(RbPolarsErr::from(e)), + } } +} - pub fn cummax(&self, reverse: bool) -> Self { - self.series.borrow().cummax(reverse).into() - } +macro_rules! impl_eq_num { + ($name:ident, $type:ty) => { + impl RbSeries { + pub fn $name(&self, rhs: $type) -> RbResult<Self> { + let s = self.series.borrow().equal(rhs).map_err(RbPolarsErr::from)?; + Ok(RbSeries::new(s.into_series())) + } + } + }; +} - pub fn cummin(&self, reverse: bool) -> Self { - self.series.borrow().cummin(reverse).into() - } +impl_eq_num!(eq_u8, u8); +impl_eq_num!(eq_u16, u16); +impl_eq_num!(eq_u32, u32); +impl_eq_num!(eq_u64, u64); +impl_eq_num!(eq_i8, i8); +impl_eq_num!(eq_i16, i16); +impl_eq_num!(eq_i32, i32); +impl_eq_num!(eq_i64, i64); +impl_eq_num!(eq_f32, f32); +impl_eq_num!(eq_f64, f64); +// impl_eq_num!(eq_str, &str); - pub fn cumprod(&self, reverse: bool) -> Self { - self.series.borrow().cumprod(reverse).into() - } +macro_rules! impl_neq_num { + ($name:ident, $type:ty) => { + impl RbSeries { + pub fn $name(&self, rhs: $type) -> RbResult<Self> { + let s = self + .series + .borrow() + .not_equal(rhs) + .map_err(RbPolarsErr::from)?; + Ok(RbSeries::new(s.into_series())) + } + } + }; +} - pub fn slice(&self, offset: i64, length: usize) -> Self { - let series = self.series.borrow().slice(offset, length); - series.into() - } +impl_neq_num!(neq_u8, u8); +impl_neq_num!(neq_u16, u16); +impl_neq_num!(neq_u32, u32); +impl_neq_num!(neq_u64, u64); +impl_neq_num!(neq_i8, i8); +impl_neq_num!(neq_i16, i16); +impl_neq_num!(neq_i32, i32); +impl_neq_num!(neq_i64, i64); +impl_neq_num!(neq_f32, f32); +impl_neq_num!(neq_f64, f64); +// impl_neq_num!(neq_str, &str); - pub fn ceil(&self) -> RbResult<Self> { - let s = self.series.borrow().ceil().map_err(RbPolarsErr::from)?; - Ok(s.into()) - } +macro_rules! impl_gt_num { + ($name:ident, $type:ty) => { + impl RbSeries { + pub fn $name(&self, rhs: $type) -> RbResult<Self> { + let s = self.series.borrow().gt(rhs).map_err(RbPolarsErr::from)?; + Ok(RbSeries::new(s.into_series())) + } + } + }; +} - pub fn round(&self, decimals: u32) -> RbResult<Self> { - let s = self - .series - .borrow() - .round(decimals) - .map_err(RbPolarsErr::from)?; - Ok(s.into()) +impl_gt_num!(gt_u8, u8); +impl_gt_num!(gt_u16, u16); +impl_gt_num!(gt_u32, u32); +impl_gt_num!(gt_u64, u64); +impl_gt_num!(gt_i8, i8); +impl_gt_num!(gt_i16, i16); +impl_gt_num!(gt_i32, i32); +impl_gt_num!(gt_i64, i64); +impl_gt_num!(gt_f32, f32); +impl_gt_num!(gt_f64, f64); +// impl_gt_num!(gt_str, &str); + +macro_rules! impl_gt_eq_num { + ($name:ident, $type:ty) => { + impl RbSeries { + pub fn $name(&self, rhs: $type) -> RbResult<Self> { + let s = self.series.borrow().gt_eq(rhs).map_err(RbPolarsErr::from)?; + Ok(RbSeries::new(s.into_series())) + } + } + }; +} + +impl_gt_eq_num!(gt_eq_u8, u8); +impl_gt_eq_num!(gt_eq_u16, u16); +impl_gt_eq_num!(gt_eq_u32, u32); +impl_gt_eq_num!(gt_eq_u64, u64); +impl_gt_eq_num!(gt_eq_i8, i8); +impl_gt_eq_num!(gt_eq_i16, i16); +impl_gt_eq_num!(gt_eq_i32, i32); +impl_gt_eq_num!(gt_eq_i64, i64); +impl_gt_eq_num!(gt_eq_f32, f32); +impl_gt_eq_num!(gt_eq_f64, f64); +// impl_gt_eq_num!(gt_eq_str, &str); + +macro_rules! impl_lt_num { + ($name:ident, $type:ty) => { + impl RbSeries { + pub fn $name(&self, rhs: $type) -> RbResult<RbSeries> { + let s = self.series.borrow().lt(rhs).map_err(RbPolarsErr::from)?; + Ok(RbSeries::new(s.into_series())) + } + } + }; +} + +impl_lt_num!(lt_u8, u8); +impl_lt_num!(lt_u16, u16); +impl_lt_num!(lt_u32, u32); +impl_lt_num!(lt_u64, u64); +impl_lt_num!(lt_i8, i8); +impl_lt_num!(lt_i16, i16); +impl_lt_num!(lt_i32, i32); +impl_lt_num!(lt_i64, i64); +impl_lt_num!(lt_f32, f32); +impl_lt_num!(lt_f64, f64); +// impl_lt_num!(lt_str, &str); + +macro_rules! impl_lt_eq_num { + ($name:ident, $type:ty) => { + impl RbSeries { + pub fn $name(&self, rhs: $type) -> RbResult<Self> { + let s = self.series.borrow().lt_eq(rhs).map_err(RbPolarsErr::from)?; + Ok(RbSeries::new(s.into_series())) + } + } + }; +} + +impl_lt_eq_num!(lt_eq_u8, u8); +impl_lt_eq_num!(lt_eq_u16, u16); +impl_lt_eq_num!(lt_eq_u32, u32); +impl_lt_eq_num!(lt_eq_u64, u64); +impl_lt_eq_num!(lt_eq_i8, i8); +impl_lt_eq_num!(lt_eq_i16, i16); +impl_lt_eq_num!(lt_eq_i32, i32); +impl_lt_eq_num!(lt_eq_i64, i64); +impl_lt_eq_num!(lt_eq_f32, f32); +impl_lt_eq_num!(lt_eq_f64, f64); +// impl_lt_eq_num!(lt_eq_str, &str); + +pub fn to_series_collection(rs: RArray) -> RbResult<Vec<Series>> { + let mut series = Vec::new(); + for item in rs.each() { + series.push(item?.try_convert::<&RbSeries>()?.series.borrow().clone()); } + Ok(series) } pub fn to_rbseries_collection(s: Vec<Series>) -> Vec<RbSeries> { s.into_iter().map(RbSeries::new).collect() +} + +impl RbSeries { + pub fn new_opt_date(name: String, values: RArray, _strict: Option<bool>) -> RbResult<Self> { + let len = values.len(); + let mut builder = PrimitiveChunkedBuilder::<Int32Type>::new(&name, len); + for item in values.each() { + let v = item?; + if v.is_nil() { + builder.append_null(); + } else { + // convert to DateTime for UTC + let v: Value = v.funcall("to_datetime", ())?; + let v: Value = v.funcall("to_time", ())?; + let v: Value = v.funcall("to_i", ())?; + // TODO use strict + builder.append_value(v.try_convert::<i32>()? / 86400); + } + } + let ca: ChunkedArray<Int32Type> = builder.finish(); + Ok(ca.into_date().into_series().into()) + } }