ext/polars/src/series.rs in polars-df-0.1.4 vs ext/polars/src/series.rs in polars-df-0.1.5
- old
+ new
@@ -2,10 +2,12 @@
use magnus::{Error, RArray, Value};
use polars::prelude::*;
use polars::series::IsSorted;
use std::cell::RefCell;
+use crate::apply::series::{call_lambda_and_extract, ApplyLambda};
+use crate::apply_method_all_arrow_series2;
use crate::conversion::*;
use crate::list_construction::rb_seq_to_list;
use crate::set::set_at_idx;
use crate::{RbDataFrame, RbPolarsErr, RbResult, RbValueError};
@@ -527,10 +529,202 @@
pub fn clone(&self) -> Self {
RbSeries::new(self.series.borrow().clone())
}
+ pub fn apply_lambda(
+ &self,
+ lambda: Value,
+ output_type: Option<Wrap<DataType>>,
+ skip_nulls: bool,
+ ) -> RbResult<Self> {
+ let series = &self.series.borrow();
+
+ let output_type = output_type.map(|dt| dt.0);
+
+ macro_rules! dispatch_apply {
+ ($self:expr, $method:ident, $($args:expr),*) => {
+ if matches!($self.dtype(), DataType::Object(_)) {
+ // let ca = $self.0.unpack::<ObjectType<ObjectValue>>().unwrap();
+ // ca.$method($($args),*)
+ todo!()
+ } else {
+ apply_method_all_arrow_series2!(
+ $self,
+ $method,
+ $($args),*
+ )
+ }
+
+ }
+
+ }
+
+ if matches!(
+ series.dtype(),
+ DataType::Datetime(_, _)
+ | DataType::Date
+ | DataType::Duration(_)
+ | DataType::Categorical(_)
+ | DataType::Time
+ ) || !skip_nulls
+ {
+ let mut avs = Vec::with_capacity(series.len());
+ let iter = series.iter().map(|av| {
+ let input = Wrap(av);
+ call_lambda_and_extract::<_, Wrap<AnyValue>>(lambda, input)
+ .unwrap()
+ .0
+ });
+ avs.extend(iter);
+ return Ok(Series::new(&self.name(), &avs).into());
+ }
+
+ let out = match output_type {
+ Some(DataType::Int8) => {
+ let ca: Int8Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::Int16) => {
+ let ca: Int16Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::Int32) => {
+ let ca: Int32Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::Int64) => {
+ let ca: Int64Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::UInt8) => {
+ let ca: UInt8Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::UInt16) => {
+ let ca: UInt16Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::UInt32) => {
+ let ca: UInt32Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::UInt64) => {
+ let ca: UInt64Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::Float32) => {
+ let ca: Float32Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::Float64) => {
+ let ca: Float64Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_series()
+ }
+ Some(DataType::Boolean) => {
+ let ca: BooleanChunked =
+ dispatch_apply!(series, apply_lambda_with_bool_out_type, lambda, 0, None)?;
+ ca.into_series()
+ }
+ Some(DataType::Date) => {
+ let ca: Int32Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_date().into_series()
+ }
+ Some(DataType::Datetime(tu, tz)) => {
+ let ca: Int64Chunked = dispatch_apply!(
+ series,
+ apply_lambda_with_primitive_out_type,
+ lambda,
+ 0,
+ None
+ )?;
+ ca.into_datetime(tu, tz).into_series()
+ }
+ Some(DataType::Utf8) => {
+ let ca = dispatch_apply!(series, apply_lambda_with_utf8_out_type, lambda, 0, None)?;
+
+ ca.into_series()
+ }
+ Some(DataType::Object(_)) => {
+ let ca =
+ dispatch_apply!(series, apply_lambda_with_object_out_type, lambda, 0, None)?;
+ ca.into_series()
+ }
+ None => return dispatch_apply!(series, apply_lambda_unknown, lambda),
+
+ _ => return dispatch_apply!(series, apply_lambda_unknown, lambda),
+ };
+
+ Ok(RbSeries::new(out))
+ }
+
pub fn zip_with(&self, mask: &RbSeries, other: &RbSeries) -> RbResult<Self> {
let binding = mask.series.borrow();
let mask = binding.bool().map_err(RbPolarsErr::from)?;
let s = self
.series
@@ -624,9 +818,111 @@
}
Err(e) => Err(RbPolarsErr::from(e)),
}
}
}
+
+macro_rules! impl_set_with_mask {
+ ($name:ident, $native:ty, $cast:ident, $variant:ident) => {
+ fn $name(
+ series: &Series,
+ filter: &RbSeries,
+ value: Option<$native>,
+ ) -> PolarsResult<Series> {
+ let binding = filter.series.borrow();
+ let mask = binding.bool()?;
+ let ca = series.$cast()?;
+ let new = ca.set(mask, value)?;
+ Ok(new.into_series())
+ }
+
+ impl RbSeries {
+ pub fn $name(&self, filter: &RbSeries, value: Option<$native>) -> RbResult<Self> {
+ let series =
+ $name(&self.series.borrow(), filter, value).map_err(RbPolarsErr::from)?;
+ Ok(Self::new(series))
+ }
+ }
+ };
+}
+
+// impl_set_with_mask!(set_with_mask_str, &str, utf8, Utf8);
+impl_set_with_mask!(set_with_mask_f64, f64, f64, Float64);
+impl_set_with_mask!(set_with_mask_f32, f32, f32, Float32);
+impl_set_with_mask!(set_with_mask_u8, u8, u8, UInt8);
+impl_set_with_mask!(set_with_mask_u16, u16, u16, UInt16);
+impl_set_with_mask!(set_with_mask_u32, u32, u32, UInt32);
+impl_set_with_mask!(set_with_mask_u64, u64, u64, UInt64);
+impl_set_with_mask!(set_with_mask_i8, i8, i8, Int8);
+impl_set_with_mask!(set_with_mask_i16, i16, i16, Int16);
+impl_set_with_mask!(set_with_mask_i32, i32, i32, Int32);
+impl_set_with_mask!(set_with_mask_i64, i64, i64, Int64);
+impl_set_with_mask!(set_with_mask_bool, bool, bool, Boolean);
+
+macro_rules! impl_arithmetic {
+ ($name:ident, $type:ty, $operand:tt) => {
+ impl RbSeries {
+ pub fn $name(&self, other: $type) -> RbResult<Self> {
+ Ok(RbSeries::new(&*self.series.borrow() $operand other))
+ }
+ }
+ };
+}
+
+impl_arithmetic!(add_u8, u8, +);
+impl_arithmetic!(add_u16, u16, +);
+impl_arithmetic!(add_u32, u32, +);
+impl_arithmetic!(add_u64, u64, +);
+impl_arithmetic!(add_i8, i8, +);
+impl_arithmetic!(add_i16, i16, +);
+impl_arithmetic!(add_i32, i32, +);
+impl_arithmetic!(add_i64, i64, +);
+impl_arithmetic!(add_datetime, i64, +);
+impl_arithmetic!(add_duration, i64, +);
+impl_arithmetic!(add_f32, f32, +);
+impl_arithmetic!(add_f64, f64, +);
+impl_arithmetic!(sub_u8, u8, -);
+impl_arithmetic!(sub_u16, u16, -);
+impl_arithmetic!(sub_u32, u32, -);
+impl_arithmetic!(sub_u64, u64, -);
+impl_arithmetic!(sub_i8, i8, -);
+impl_arithmetic!(sub_i16, i16, -);
+impl_arithmetic!(sub_i32, i32, -);
+impl_arithmetic!(sub_i64, i64, -);
+impl_arithmetic!(sub_datetime, i64, -);
+impl_arithmetic!(sub_duration, i64, -);
+impl_arithmetic!(sub_f32, f32, -);
+impl_arithmetic!(sub_f64, f64, -);
+impl_arithmetic!(div_u8, u8, /);
+impl_arithmetic!(div_u16, u16, /);
+impl_arithmetic!(div_u32, u32, /);
+impl_arithmetic!(div_u64, u64, /);
+impl_arithmetic!(div_i8, i8, /);
+impl_arithmetic!(div_i16, i16, /);
+impl_arithmetic!(div_i32, i32, /);
+impl_arithmetic!(div_i64, i64, /);
+impl_arithmetic!(div_f32, f32, /);
+impl_arithmetic!(div_f64, f64, /);
+impl_arithmetic!(mul_u8, u8, *);
+impl_arithmetic!(mul_u16, u16, *);
+impl_arithmetic!(mul_u32, u32, *);
+impl_arithmetic!(mul_u64, u64, *);
+impl_arithmetic!(mul_i8, i8, *);
+impl_arithmetic!(mul_i16, i16, *);
+impl_arithmetic!(mul_i32, i32, *);
+impl_arithmetic!(mul_i64, i64, *);
+impl_arithmetic!(mul_f32, f32, *);
+impl_arithmetic!(mul_f64, f64, *);
+impl_arithmetic!(rem_u8, u8, %);
+impl_arithmetic!(rem_u16, u16, %);
+impl_arithmetic!(rem_u32, u32, %);
+impl_arithmetic!(rem_u64, u64, %);
+impl_arithmetic!(rem_i8, i8, %);
+impl_arithmetic!(rem_i16, i16, %);
+impl_arithmetic!(rem_i32, i32, %);
+impl_arithmetic!(rem_i64, i64, %);
+impl_arithmetic!(rem_f32, f32, %);
+impl_arithmetic!(rem_f64, f64, %);
macro_rules! impl_eq_num {
($name:ident, $type:ty) => {
impl RbSeries {
pub fn $name(&self, rhs: $type) -> RbResult<Self> {