polars-df 0.1.4 → 0.2.0

data/ext/polars/src/lazy/dataframe.rs

@@ -3,7 +3,7 @@ use polars::io::RowCount;
 use polars::lazy::frame::{LazyFrame, LazyGroupBy};
 use polars::prelude::*;
 use std::cell::RefCell;
-use std::io::BufWriter;
+use std::io::{BufWriter, Read};
 
 use crate::conversion::*;
 use crate::file::get_file_like;
@@ -53,6 +53,27 @@ impl From<LazyFrame> for RbLazyFrame {
 }
 
 impl RbLazyFrame {
+    pub fn read_json(rb_f: Value) -> RbResult<Self> {
+        // it is faster to first read to memory and then parse: https://github.com/serde-rs/json/issues/160
+        // so don't bother with files.
+        let mut json = String::new();
+        let _ = get_file_like(rb_f, false)?
+            .read_to_string(&mut json)
+            .unwrap();
+
+        // Safety
+        // we skipped the serializing/deserializing of the static in lifetime in `DataType`
+        // so we actually don't have a lifetime at all when serializing.
+
+        // &str still has a lifetime. Bit its ok, because we drop it immediately
+        // in this scope
+        let json = unsafe { std::mem::transmute::<&'_ str, &'static str>(json.as_str()) };
+
+        let lp = serde_json::from_str::<LogicalPlan>(json)
+            .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
+        Ok(LazyFrame::from(lp).into())
+    }
+
     pub fn new_from_ndjson(
         path: String,
         infer_schema_length: Option<usize>,
@@ -327,6 +348,7 @@ impl RbLazyFrame {
         include_boundaries: bool,
         closed: Wrap<ClosedWindow>,
         by: RArray,
+        start_by: Wrap<StartBy>,
     ) -> RbResult<RbLazyGroupBy> {
         let closed_window = closed.0;
         let by = rb_exprs_to_exprs(by)?;
@@ -341,6 +363,7 @@ impl RbLazyFrame {
                 truncate,
                 include_boundaries,
                 closed_window,
+                start_by: start_by.0,
             },
         );
 
@@ -349,6 +372,56 @@ impl RbLazyFrame {
         })
     }
 
+    pub fn with_context(&self, contexts: RArray) -> RbResult<Self> {
+        let contexts = contexts
+            .each()
+            .map(|v| v.unwrap().try_convert())
+            .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
+        let contexts = contexts
+            .into_iter()
+            .map(|ldf| ldf.ldf.clone())
+            .collect::<Vec<_>>();
+        Ok(self.ldf.clone().with_context(contexts).into())
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn join_asof(
+        &self,
+        other: &RbLazyFrame,
+        left_on: &RbExpr,
+        right_on: &RbExpr,
+        left_by: Option<Vec<String>>,
+        right_by: Option<Vec<String>>,
+        allow_parallel: bool,
+        force_parallel: bool,
+        suffix: String,
+        strategy: Wrap<AsofStrategy>,
+        tolerance: Option<Wrap<AnyValue<'_>>>,
+        tolerance_str: Option<String>,
+    ) -> RbResult<Self> {
+        let ldf = self.ldf.clone();
+        let other = other.ldf.clone();
+        let left_on = left_on.inner.clone();
+        let right_on = right_on.inner.clone();
+        Ok(ldf
+            .join_builder()
+            .with(other)
+            .left_on([left_on])
+            .right_on([right_on])
+            .allow_parallel(allow_parallel)
+            .force_parallel(force_parallel)
+            .how(JoinType::AsOf(AsOfOptions {
+                strategy: strategy.0,
+                left_by,
+                right_by,
+                tolerance: tolerance.map(|t| t.0.into_static().unwrap()),
+                tolerance_str,
+            }))
+            .suffix(suffix)
+            .finish()
+            .into())
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn join(
         &self,
@@ -443,9 +516,13 @@ impl RbLazyFrame {
         ldf.median().into()
     }
 
-    pub fn quantile(&self, quantile: f64, interpolation: Wrap<QuantileInterpolOptions>) -> Self {
+    pub fn quantile(
+        &self,
+        quantile: &RbExpr,
+        interpolation: Wrap<QuantileInterpolOptions>,
+    ) -> Self {
         let ldf = self.ldf.clone();
-        ldf.quantile(quantile, interpolation.0).into()
+        ldf.quantile(quantile.inner.clone(), interpolation.0).into()
     }
 
     pub fn explode(&self, column: RArray) -> RbResult<Self> {

data/ext/polars/src/lazy/dsl.rs

@@ -1,3 +1,4 @@
+use magnus::block::Proc;
 use magnus::{class, RArray, RString, Value};
 use polars::chunked_array::ops::SortOptions;
 use polars::lazy::dsl;
@@ -163,10 +164,14 @@ impl RbExpr {
         self.clone().inner.list().into()
     }
 
-    pub fn quantile(&self, quantile: f64, interpolation: Wrap<QuantileInterpolOptions>) -> Self {
+    pub fn quantile(
+        &self,
+        quantile: &RbExpr,
+        interpolation: Wrap<QuantileInterpolOptions>,
+    ) -> Self {
         self.clone()
             .inner
-            .quantile(quantile, interpolation.0)
+            .quantile(quantile.inner.clone(), interpolation.0)
             .into()
     }
 
@@ -505,7 +510,13 @@ impl RbExpr {
         self.inner.clone().shrink_dtype().into()
     }
 
-    pub fn str_parse_date(&self, fmt: Option<String>, strict: bool, exact: bool) -> Self {
+    pub fn str_parse_date(
+        &self,
+        fmt: Option<String>,
+        strict: bool,
+        exact: bool,
+        cache: bool,
+    ) -> Self {
         self.inner
             .clone()
             .str()
@@ -514,11 +525,20 @@ impl RbExpr {
                 fmt,
                 strict,
                 exact,
+                cache,
+                tz_aware: false,
             })
             .into()
     }
 
-    pub fn str_parse_datetime(&self, fmt: Option<String>, strict: bool, exact: bool) -> Self {
+    pub fn str_parse_datetime(
+        &self,
+        fmt: Option<String>,
+        strict: bool,
+        exact: bool,
+        cache: bool,
+        tz_aware: bool,
+    ) -> Self {
         let tu = match fmt {
             Some(ref fmt) => {
                 if fmt.contains("%.9f")
@@ -543,11 +563,19 @@ impl RbExpr {
                 fmt,
                 strict,
                 exact,
+                cache,
+                tz_aware,
             })
             .into()
     }
 
-    pub fn str_parse_time(&self, fmt: Option<String>, strict: bool, exact: bool) -> Self {
+    pub fn str_parse_time(
+        &self,
+        fmt: Option<String>,
+        strict: bool,
+        exact: bool,
+        cache: bool,
+    ) -> Self {
         self.inner
             .clone()
             .str()
@@ -556,6 +584,8 @@ impl RbExpr {
                 fmt,
                 strict,
                 exact,
+                cache,
+                tz_aware: false,
             })
             .into()
     }
@@ -722,8 +752,12 @@ impl RbExpr {
         self.inner.clone().str().extract(&pat, group_index).into()
     }
 
-    pub fn str_extract_all(&self, pat: String) -> Self {
-        self.inner.clone().str().extract_all(&pat).into()
+    pub fn str_extract_all(&self, pat: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .extract_all(pat.inner.clone())
+            .into()
     }
 
     pub fn count_match(&self, pat: String) -> Self {
@@ -946,6 +980,10 @@ impl RbExpr {
         self.inner.clone().dt().round(&every, &offset).into()
     }
 
+    pub fn map(&self, lambda: Value, output_type: Option<Wrap<DataType>>, agg_list: bool) -> Self {
+        map_single(self, lambda, output_type, agg_list)
+    }
+
     pub fn dot(&self, other: &RbExpr) -> Self {
         self.inner.clone().dot(other.inner.clone()).into()
     }
@@ -979,12 +1017,27 @@ impl RbExpr {
         self.inner.clone().suffix(&suffix).into()
     }
 
+    pub fn map_alias(&self, lambda: Proc) -> Self {
+        self.inner
+            .clone()
+            .map_alias(move |name| {
+                let out = lambda.call::<_, String>((name,));
+                match out {
+                    Ok(out) => Ok(out),
+                    Err(e) => Err(PolarsError::ComputeError(
+                        format!("Ruby function in 'map_alias' produced an error: {}.", e).into(),
+                    )),
+                }
+            })
+            .into()
+    }
+
     pub fn exclude(&self, columns: Vec<String>) -> Self {
         self.inner.clone().exclude(columns).into()
     }
 
-    pub fn interpolate(&self) -> Self {
-        self.inner.clone().interpolate().into()
+    pub fn interpolate(&self, method: Wrap<InterpolationMethod>) -> Self {
+        self.inner.clone().interpolate(method.0).into()
     }
 
     pub fn rolling_sum(
@@ -1275,6 +1328,7 @@ impl RbExpr {
         &self,
         width_strat: Wrap<ListToStructWidthStrategy>,
         _name_gen: Option<Value>,
+        upper_bound: usize,
     ) -> RbResult<Self> {
         // TODO fix
         let name_gen = None;
@@ -1289,7 +1343,7 @@ impl RbExpr {
             .inner
             .clone()
             .arr()
-            .to_struct(width_strat.0, name_gen)
+            .to_struct(width_strat.0, name_gen, upper_bound)
             .into())
     }
 
@@ -1450,6 +1504,10 @@ impl RbExpr {
     pub fn entropy(&self, base: f64, normalize: bool) -> Self {
         self.inner.clone().entropy(base, normalize).into()
     }
+
+    pub fn hash(&self, seed: u64, seed_1: u64, seed_2: u64, seed_3: u64) -> Self {
+        self.inner.clone().hash(seed, seed_1, seed_2, seed_3).into()
+    }
 }
 
 pub fn col(name: String) -> RbExpr {
@@ -1472,6 +1530,10 @@ pub fn cols(names: Vec<String>) -> RbExpr {
     dsl::cols(names).into()
 }
 
+pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
+    dsl::dtype_cols(dtypes).into()
+}
+
 pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
 
@@ -1479,6 +1541,13 @@ pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
     Ok(polars::lazy::dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
 }
 
+pub fn cumfold(acc: &RbExpr, lambda: Value, exprs: RArray, include_init: bool) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+
+    let func = move |a: Series, b: Series| binary_lambda(lambda, a, b);
+    Ok(polars::lazy::dsl::cumfold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+}
+
 // TODO improve
 pub fn lit(value: Value) -> RbResult<RbExpr> {
     if value.is_nil() {
@@ -1531,6 +1600,11 @@ pub fn cov(a: &RbExpr, b: &RbExpr) -> RbExpr {
     polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone()).into()
 }
 
+pub fn argsort_by(by: RArray, reverse: Vec<bool>) -> RbResult<RbExpr> {
+    let by = rb_exprs_to_exprs(by)?;
+    Ok(polars::lazy::dsl::argsort_by(by, &reverse).into())
+}
+
 #[magnus::wrap(class = "Polars::RbWhen")]
 #[derive(Clone)]
 pub struct RbWhen {

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,19 @@ fn series() -> RClass {
 #[magnus::init]
 fn init() -> RbResult<()> {
     let module = module();
+    module.define_singleton_method("_dtype_cols", function!(dtype_cols, 1))?;
+    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 +81,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 +153,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 +172,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))?;
@@ -277,9 +290,9 @@ fn init() -> RbResult<()> {
     class.define_method("cumprod", method!(RbExpr::cumprod, 1))?;
     class.define_method("product", method!(RbExpr::product, 0))?;
     class.define_method("shrink_dtype", method!(RbExpr::shrink_dtype, 0))?;
-    class.define_method("str_parse_date", method!(RbExpr::str_parse_date, 3))?;
-    class.define_method("str_parse_datetime", method!(RbExpr::str_parse_datetime, 3))?;
-    class.define_method("str_parse_time", method!(RbExpr::str_parse_time, 3))?;
+    class.define_method("str_parse_date", method!(RbExpr::str_parse_date, 4))?;
+    class.define_method("str_parse_datetime", method!(RbExpr::str_parse_datetime, 5))?;
+    class.define_method("str_parse_time", method!(RbExpr::str_parse_time, 4))?;
     class.define_method("str_strip", method!(RbExpr::str_strip, 1))?;
     class.define_method("str_rstrip", method!(RbExpr::str_rstrip, 1))?;
     class.define_method("str_lstrip", method!(RbExpr::str_lstrip, 1))?;
@@ -361,14 +374,16 @@ 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("interpolate", method!(RbExpr::interpolate, 1))?;
     class.define_method("rolling_sum", method!(RbExpr::rolling_sum, 6))?;
     class.define_method("rolling_min", method!(RbExpr::rolling_min, 6))?;
     class.define_method("rolling_max", method!(RbExpr::rolling_max, 6))?;
@@ -396,7 +411,7 @@ fn init() -> RbResult<()> {
     class.define_method("lst_slice", method!(RbExpr::lst_slice, 2))?;
     class.define_method("lst_eval", method!(RbExpr::lst_eval, 2))?;
     class.define_method("cumulative_eval", method!(RbExpr::cumulative_eval, 3))?;
-    class.define_method("lst_to_struct", method!(RbExpr::lst_to_struct, 2))?;
+    class.define_method("lst_to_struct", method!(RbExpr::lst_to_struct, 3))?;
     class.define_method("rank", method!(RbExpr::rank, 2))?;
     class.define_method("diff", method!(RbExpr::diff, 2))?;
     class.define_method("pct_change", method!(RbExpr::pct_change, 1))?;
@@ -431,6 +446,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 +462,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 +472,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),
@@ -489,7 +508,9 @@ fn init() -> RbResult<()> {
     class.define_method("select", method!(RbLazyFrame::select, 1))?;
     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("groupby_dynamic", method!(RbLazyFrame::groupby_dynamic, 9))?;
+    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 +618,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 +633,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 +793,40 @@ fn init() -> RbResult<()> {
     Ok(())
 }
 
+fn dtype_cols(dtypes: RArray) -> RbResult<RbExpr> {
+    let dtypes = dtypes
+        .each()
+        .map(|v| v?.try_convert::<Wrap<DataType>>())
+        .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
+    let dtypes = vec_extract_wrapped(dtypes);
+    Ok(crate::lazy::dsl::dtype_cols(dtypes))
+}
+
+#[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 +855,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 +929,25 @@ 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>>>()?;
+
+    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)
+    })
+}
+
 fn rb_date_range(
     start: i64,
     stop: i64,
@@ -802,12 +964,22 @@ fn rb_date_range(
         Duration::parse(&every),
         closed.0,
         tu.0,
-        tz,
+        tz.as_ref(),
     )
     .into_series()
     .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())