polars-df 0.1.4 → 0.1.5

data/ext/polars/src/lib.rs

@@ -1,3 +1,4 @@
+mod apply;
 mod batched_csv;
 mod conversion;
 mod dataframe;
@@ -25,7 +26,7 @@ use polars::datatypes::{DataType, TimeUnit};
 use polars::error::PolarsResult;
 use polars::frame::DataFrame;
 use polars::functions::{diag_concat_df, hor_concat_df};
-use polars::prelude::{ClosedWindow, Duration, IntoSeries, TimeZone};
+use polars::prelude::{ClosedWindow, Duration, DurationArgs, IntoSeries, TimeZone};
 use series::RbSeries;
 
 #[cfg(target_os = "linux")]
@@ -55,13 +56,18 @@ fn series() -> RClass {
 #[magnus::init]
 fn init() -> RbResult<()> {
     let module = module();
+    module.define_singleton_method("_rb_duration", function!(rb_duration, 8))?;
     module.define_singleton_method("_concat_df", function!(concat_df, 1))?;
+    module.define_singleton_method("_concat_lf", function!(concat_lf, 3))?;
     module.define_singleton_method("_diag_concat_df", function!(rb_diag_concat_df, 1))?;
     module.define_singleton_method("_hor_concat_df", function!(rb_hor_concat_df, 1))?;
     module.define_singleton_method("_concat_series", function!(concat_series, 1))?;
     module.define_singleton_method("_ipc_schema", function!(ipc_schema, 1))?;
     module.define_singleton_method("_parquet_schema", function!(parquet_schema, 1))?;
+    module.define_singleton_method("_collect_all", function!(collect_all, 1))?;
     module.define_singleton_method("_rb_date_range", function!(rb_date_range, 7))?;
+    module.define_singleton_method("_coalesce_exprs", function!(coalesce_exprs, 1))?;
+    module.define_singleton_method("_sum_exprs", function!(sum_exprs, 1))?;
     module.define_singleton_method("_as_struct", function!(as_struct, 1))?;
     module.define_singleton_method("_arg_where", function!(arg_where, 1))?;
 
@@ -74,10 +80,13 @@ fn init() -> RbResult<()> {
     class.define_singleton_method("read_csv", function!(RbDataFrame::read_csv, -1))?;
     class.define_singleton_method("read_parquet", function!(RbDataFrame::read_parquet, 7))?;
     class.define_singleton_method("read_ipc", function!(RbDataFrame::read_ipc, 6))?;
+    class.define_singleton_method("read_avro", function!(RbDataFrame::read_avro, 4))?;
+    class.define_singleton_method("read_hashes", function!(RbDataFrame::read_hashes, 3))?;
     class.define_singleton_method("read_hash", function!(RbDataFrame::read_hash, 1))?;
     class.define_singleton_method("read_json", function!(RbDataFrame::read_json, 1))?;
     class.define_singleton_method("read_ndjson", function!(RbDataFrame::read_ndjson, 1))?;
     class.define_method("estimated_size", method!(RbDataFrame::estimated_size, 0))?;
+    class.define_method("write_avro", method!(RbDataFrame::write_avro, 2))?;
     class.define_method("write_json", method!(RbDataFrame::write_json, 3))?;
     class.define_method("write_ndjson", method!(RbDataFrame::write_ndjson, 1))?;
     class.define_method("write_csv", method!(RbDataFrame::write_csv, 10))?;
@@ -143,6 +152,7 @@ fn init() -> RbResult<()> {
     class.define_method("with_row_count", method!(RbDataFrame::with_row_count, 2))?;
     class.define_method("_clone", method!(RbDataFrame::clone, 0))?;
     class.define_method("melt", method!(RbDataFrame::melt, 4))?;
+    class.define_method("pivot_expr", method!(RbDataFrame::pivot_expr, 6))?;
     class.define_method("partition_by", method!(RbDataFrame::partition_by, 2))?;
     class.define_method("shift", method!(RbDataFrame::shift, 1))?;
     class.define_method("unique", method!(RbDataFrame::unique, 3))?;
@@ -161,7 +171,9 @@ fn init() -> RbResult<()> {
     class.define_method("quantile", method!(RbDataFrame::quantile, 2))?;
     class.define_method("to_dummies", method!(RbDataFrame::to_dummies, 1))?;
     class.define_method("null_count", method!(RbDataFrame::null_count, 0))?;
+    class.define_method("apply", method!(RbDataFrame::apply, 3))?;
     class.define_method("shrink_to_fit", method!(RbDataFrame::shrink_to_fit, 0))?;
+    class.define_method("hash_rows", method!(RbDataFrame::hash_rows, 4))?;
     class.define_method("transpose", method!(RbDataFrame::transpose, 2))?;
     class.define_method("upsample", method!(RbDataFrame::upsample, 5))?;
     class.define_method("to_struct", method!(RbDataFrame::to_struct, 1))?;
@@ -361,12 +373,14 @@ fn init() -> RbResult<()> {
     class.define_method("dt_tz_localize", method!(RbExpr::dt_tz_localize, 1))?;
     class.define_method("dt_truncate", method!(RbExpr::dt_truncate, 2))?;
     class.define_method("dt_round", method!(RbExpr::dt_round, 2))?;
+    class.define_method("map", method!(RbExpr::map, 3))?;
     class.define_method("dot", method!(RbExpr::dot, 1))?;
     class.define_method("reinterpret", method!(RbExpr::reinterpret, 1))?;
     class.define_method("mode", method!(RbExpr::mode, 0))?;
     class.define_method("keep_name", method!(RbExpr::keep_name, 0))?;
     class.define_method("prefix", method!(RbExpr::prefix, 1))?;
     class.define_method("suffix", method!(RbExpr::suffix, 1))?;
+    class.define_method("map_alias", method!(RbExpr::map_alias, 1))?;
     class.define_method("exclude", method!(RbExpr::exclude, 1))?;
     class.define_method("interpolate", method!(RbExpr::interpolate, 0))?;
     class.define_method("rolling_sum", method!(RbExpr::rolling_sum, 6))?;
@@ -431,6 +445,7 @@ fn init() -> RbResult<()> {
     class.define_method("log", method!(RbExpr::log, 1))?;
     class.define_method("exp", method!(RbExpr::exp, 0))?;
     class.define_method("entropy", method!(RbExpr::entropy, 2))?;
+    class.define_method("_hash", method!(RbExpr::hash, 4))?;
 
     // meta
     class.define_method("meta_pop", method!(RbExpr::meta_pop, 0))?;
@@ -446,6 +461,7 @@ fn init() -> RbResult<()> {
     class.define_singleton_method("last", function!(crate::lazy::dsl::last, 0))?;
     class.define_singleton_method("cols", function!(crate::lazy::dsl::cols, 1))?;
     class.define_singleton_method("fold", function!(crate::lazy::dsl::fold, 3))?;
+    class.define_singleton_method("cumfold", function!(crate::lazy::dsl::cumfold, 4))?;
     class.define_singleton_method("lit", function!(crate::lazy::dsl::lit, 1))?;
     class.define_singleton_method("arange", function!(crate::lazy::dsl::arange, 3))?;
     class.define_singleton_method("repeat", function!(crate::lazy::dsl::repeat, 2))?;
@@ -455,11 +471,13 @@ fn init() -> RbResult<()> {
         function!(crate::lazy::dsl::spearman_rank_corr, 4),
     )?;
     class.define_singleton_method("cov", function!(crate::lazy::dsl::cov, 2))?;
+    class.define_singleton_method("argsort_by", function!(crate::lazy::dsl::argsort_by, 2))?;
     class.define_singleton_method("when", function!(crate::lazy::dsl::when, 1))?;
     class.define_singleton_method("concat_str", function!(crate::lazy::dsl::concat_str, 2))?;
     class.define_singleton_method("concat_lst", function!(crate::lazy::dsl::concat_lst, 1))?;
 
     let class = module.define_class("RbLazyFrame", Default::default())?;
+    class.define_singleton_method("read_json", function!(RbLazyFrame::read_json, 1))?;
     class.define_singleton_method(
         "new_from_ndjson",
         function!(RbLazyFrame::new_from_ndjson, 7),
@@ -490,6 +508,8 @@ fn init() -> RbResult<()> {
     class.define_method("groupby", method!(RbLazyFrame::groupby, 2))?;
     class.define_method("groupby_rolling", method!(RbLazyFrame::groupby_rolling, 5))?;
     class.define_method("groupby_dynamic", method!(RbLazyFrame::groupby_dynamic, 8))?;
+    class.define_method("with_context", method!(RbLazyFrame::with_context, 1))?;
+    class.define_method("join_asof", method!(RbLazyFrame::join_asof, 11))?;
     class.define_method("join", method!(RbLazyFrame::join, 7))?;
     class.define_method("with_columns", method!(RbLazyFrame::with_columns, 1))?;
     class.define_method("rename", method!(RbLazyFrame::rename, 2))?;
@@ -597,6 +617,7 @@ fn init() -> RbResult<()> {
     class.define_method("median", method!(RbSeries::median, 0))?;
     class.define_method("quantile", method!(RbSeries::quantile, 2))?;
     class.define_method("_clone", method!(RbSeries::clone, 0))?;
+    class.define_method("apply_lambda", method!(RbSeries::apply_lambda, 3))?;
     class.define_method("zip_with", method!(RbSeries::zip_with, 2))?;
     class.define_method("to_dummies", method!(RbSeries::to_dummies, 0))?;
     class.define_method("peak_max", method!(RbSeries::peak_max, 0))?;
@@ -611,6 +632,79 @@ fn init() -> RbResult<()> {
     class.define_method("time_unit", method!(RbSeries::time_unit, 0))?;
     class.define_method("set_at_idx", method!(RbSeries::set_at_idx, 2))?;
 
+    // set
+    // class.define_method("set_with_mask_str", method!(RbSeries::set_with_mask_str, 2))?;
+    class.define_method("set_with_mask_f64", method!(RbSeries::set_with_mask_f64, 2))?;
+    class.define_method("set_with_mask_f32", method!(RbSeries::set_with_mask_f32, 2))?;
+    class.define_method("set_with_mask_u8", method!(RbSeries::set_with_mask_u8, 2))?;
+    class.define_method("set_with_mask_u16", method!(RbSeries::set_with_mask_u16, 2))?;
+    class.define_method("set_with_mask_u32", method!(RbSeries::set_with_mask_u32, 2))?;
+    class.define_method("set_with_mask_u64", method!(RbSeries::set_with_mask_u64, 2))?;
+    class.define_method("set_with_mask_i8", method!(RbSeries::set_with_mask_i8, 2))?;
+    class.define_method("set_with_mask_i16", method!(RbSeries::set_with_mask_i16, 2))?;
+    class.define_method("set_with_mask_i32", method!(RbSeries::set_with_mask_i32, 2))?;
+    class.define_method("set_with_mask_i64", method!(RbSeries::set_with_mask_i64, 2))?;
+    class.define_method(
+        "set_with_mask_bool",
+        method!(RbSeries::set_with_mask_bool, 2),
+    )?;
+
+    // arithmetic
+    class.define_method("add_u8", method!(RbSeries::add_u8, 1))?;
+    class.define_method("add_u16", method!(RbSeries::add_u16, 1))?;
+    class.define_method("add_u32", method!(RbSeries::add_u32, 1))?;
+    class.define_method("add_u64", method!(RbSeries::add_u64, 1))?;
+    class.define_method("add_i8", method!(RbSeries::add_i8, 1))?;
+    class.define_method("add_i16", method!(RbSeries::add_i16, 1))?;
+    class.define_method("add_i32", method!(RbSeries::add_i32, 1))?;
+    class.define_method("add_i64", method!(RbSeries::add_i64, 1))?;
+    class.define_method("add_datetime", method!(RbSeries::add_datetime, 1))?;
+    class.define_method("add_duration", method!(RbSeries::add_duration, 1))?;
+    class.define_method("add_f32", method!(RbSeries::add_f32, 1))?;
+    class.define_method("add_f64", method!(RbSeries::add_f64, 1))?;
+    class.define_method("sub_u8", method!(RbSeries::sub_u8, 1))?;
+    class.define_method("sub_u16", method!(RbSeries::sub_u16, 1))?;
+    class.define_method("sub_u32", method!(RbSeries::sub_u32, 1))?;
+    class.define_method("sub_u64", method!(RbSeries::sub_u64, 1))?;
+    class.define_method("sub_i8", method!(RbSeries::sub_i8, 1))?;
+    class.define_method("sub_i16", method!(RbSeries::sub_i16, 1))?;
+    class.define_method("sub_i32", method!(RbSeries::sub_i32, 1))?;
+    class.define_method("sub_i64", method!(RbSeries::sub_i64, 1))?;
+    class.define_method("sub_datetime", method!(RbSeries::sub_datetime, 1))?;
+    class.define_method("sub_duration", method!(RbSeries::sub_duration, 1))?;
+    class.define_method("sub_f32", method!(RbSeries::sub_f32, 1))?;
+    class.define_method("sub_f64", method!(RbSeries::sub_f64, 1))?;
+    class.define_method("div_u8", method!(RbSeries::div_u8, 1))?;
+    class.define_method("div_u16", method!(RbSeries::div_u16, 1))?;
+    class.define_method("div_u32", method!(RbSeries::div_u32, 1))?;
+    class.define_method("div_u64", method!(RbSeries::div_u64, 1))?;
+    class.define_method("div_i8", method!(RbSeries::div_i8, 1))?;
+    class.define_method("div_i16", method!(RbSeries::div_i16, 1))?;
+    class.define_method("div_i32", method!(RbSeries::div_i32, 1))?;
+    class.define_method("div_i64", method!(RbSeries::div_i64, 1))?;
+    class.define_method("div_f32", method!(RbSeries::div_f32, 1))?;
+    class.define_method("div_f64", method!(RbSeries::div_f64, 1))?;
+    class.define_method("mul_u8", method!(RbSeries::mul_u8, 1))?;
+    class.define_method("mul_u16", method!(RbSeries::mul_u16, 1))?;
+    class.define_method("mul_u32", method!(RbSeries::mul_u32, 1))?;
+    class.define_method("mul_u64", method!(RbSeries::mul_u64, 1))?;
+    class.define_method("mul_i8", method!(RbSeries::mul_i8, 1))?;
+    class.define_method("mul_i16", method!(RbSeries::mul_i16, 1))?;
+    class.define_method("mul_i32", method!(RbSeries::mul_i32, 1))?;
+    class.define_method("mul_i64", method!(RbSeries::mul_i64, 1))?;
+    class.define_method("mul_f32", method!(RbSeries::mul_f32, 1))?;
+    class.define_method("mul_f64", method!(RbSeries::mul_f64, 1))?;
+    class.define_method("rem_u8", method!(RbSeries::rem_u8, 1))?;
+    class.define_method("rem_u16", method!(RbSeries::rem_u16, 1))?;
+    class.define_method("rem_u32", method!(RbSeries::rem_u32, 1))?;
+    class.define_method("rem_u64", method!(RbSeries::rem_u64, 1))?;
+    class.define_method("rem_i8", method!(RbSeries::rem_i8, 1))?;
+    class.define_method("rem_i16", method!(RbSeries::rem_i16, 1))?;
+    class.define_method("rem_i32", method!(RbSeries::rem_i32, 1))?;
+    class.define_method("rem_i64", method!(RbSeries::rem_i64, 1))?;
+    class.define_method("rem_f32", method!(RbSeries::rem_f32, 1))?;
+    class.define_method("rem_f64", method!(RbSeries::rem_f64, 1))?;
+
     // eq
     class.define_method("eq_u8", method!(RbSeries::eq_u8, 1))?;
     class.define_method("eq_u16", method!(RbSeries::eq_u16, 1))?;
@@ -698,6 +792,31 @@ fn init() -> RbResult<()> {
     Ok(())
 }
 
+#[allow(clippy::too_many_arguments)]
+fn rb_duration(
+    days: Option<&RbExpr>,
+    seconds: Option<&RbExpr>,
+    nanoseconds: Option<&RbExpr>,
+    microseconds: Option<&RbExpr>,
+    milliseconds: Option<&RbExpr>,
+    minutes: Option<&RbExpr>,
+    hours: Option<&RbExpr>,
+    weeks: Option<&RbExpr>,
+) -> RbExpr {
+    let args = DurationArgs {
+        days: days.map(|e| e.inner.clone()),
+        seconds: seconds.map(|e| e.inner.clone()),
+        nanoseconds: nanoseconds.map(|e| e.inner.clone()),
+        microseconds: microseconds.map(|e| e.inner.clone()),
+        milliseconds: milliseconds.map(|e| e.inner.clone()),
+        minutes: minutes.map(|e| e.inner.clone()),
+        hours: hours.map(|e| e.inner.clone()),
+        weeks: weeks.map(|e| e.inner.clone()),
+    };
+
+    polars::lazy::dsl::duration(args).into()
+}
+
 fn concat_df(seq: RArray) -> RbResult<RbDataFrame> {
     let mut iter = seq.each();
     let first = iter.next().unwrap()?;
@@ -726,6 +845,20 @@ fn concat_df(seq: RArray) -> RbResult<RbDataFrame> {
     Ok(df.into())
 }
 
+fn concat_lf(lfs: Value, rechunk: bool, parallel: bool) -> RbResult<RbLazyFrame> {
+    let (seq, len) = get_rbseq(lfs)?;
+    let mut lfs = Vec::with_capacity(len);
+
+    for res in seq.each() {
+        let item = res?;
+        let lf = get_lf(item)?;
+        lfs.push(lf);
+    }
+
+    let lf = polars::lazy::dsl::concat(lfs, rechunk, parallel).map_err(RbPolarsErr::from)?;
+    Ok(lf.into())
+}
+
 fn rb_diag_concat_df(seq: RArray) -> RbResult<RbDataFrame> {
     let mut dfs = Vec::new();
     for item in seq.each() {
@@ -786,6 +919,27 @@ fn parquet_schema(rb_f: Value) -> RbResult<Value> {
     Ok(dict.into())
 }
 
+fn collect_all(lfs: RArray) -> RbResult<Vec<RbDataFrame>> {
+    use polars_core::utils::rayon::prelude::*;
+
+    let lfs = lfs
+        .each()
+        .map(|v| v?.try_convert::<&RbLazyFrame>())
+        .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
+
+    let out = polars_core::POOL.install(|| {
+        lfs.par_iter()
+            .map(|lf| {
+                let df = lf.ldf.clone().collect()?;
+                Ok(RbDataFrame::new(df))
+            })
+            .collect::<polars_core::error::PolarsResult<Vec<_>>>()
+            .map_err(RbPolarsErr::from)
+    });
+
+    Ok(out?)
+}
+
 fn rb_date_range(
     start: i64,
     stop: i64,
@@ -808,6 +962,16 @@ fn rb_date_range(
     .into()
 }
 
+fn coalesce_exprs(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    Ok(polars::lazy::dsl::coalesce(&exprs).into())
+}
+
+fn sum_exprs(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    Ok(polars::lazy::dsl::sum_exprs(exprs).into())
+}
+
 fn as_struct(exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
     Ok(polars::lazy::dsl::as_struct(&exprs).into())

data/ext/polars/src/series.rs

@@ -4,6 +4,8 @@ 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;
@@ -529,6 +531,198 @@ impl RbSeries {
         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)?;
@@ -627,6 +821,108 @@ impl RbSeries {
     }
 }
 
+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 {

data/ext/polars/src/utils.rs

@@ -17,3 +17,28 @@ pub fn reinterpret(s: &Series, signed: bool) -> polars::prelude::PolarsResult<Se
         )),
     }
 }
+
+#[macro_export]
+macro_rules! apply_method_all_arrow_series2 {
+    ($self:expr, $method:ident, $($args:expr),*) => {
+        match $self.dtype() {
+            DataType::Boolean => $self.bool().unwrap().$method($($args),*),
+            DataType::Utf8 => $self.utf8().unwrap().$method($($args),*),
+            DataType::UInt8 => $self.u8().unwrap().$method($($args),*),
+            DataType::UInt16 => $self.u16().unwrap().$method($($args),*),
+            DataType::UInt32 => $self.u32().unwrap().$method($($args),*),
+            DataType::UInt64 => $self.u64().unwrap().$method($($args),*),
+            DataType::Int8 => $self.i8().unwrap().$method($($args),*),
+            DataType::Int16 => $self.i16().unwrap().$method($($args),*),
+            DataType::Int32 => $self.i32().unwrap().$method($($args),*),
+            DataType::Int64 => $self.i64().unwrap().$method($($args),*),
+            DataType::Float32 => $self.f32().unwrap().$method($($args),*),
+            DataType::Float64 => $self.f64().unwrap().$method($($args),*),
+            DataType::Date => $self.date().unwrap().$method($($args),*),
+            DataType::Datetime(_, _) => $self.datetime().unwrap().$method($($args),*),
+            // DataType::List(_) => $self.list().unwrap().$method($($args),*),
+            DataType::Struct(_) => $self.struct_().unwrap().$method($($args),*),
+            dt => panic!("dtype {:?} not supported", dt)
+        }
+    }
+}