polars-df 0.6.0 → 0.8.0

data/ext/polars/src/expr/general.rs

@@ -1,12 +1,11 @@
-use magnus::{block::Proc, IntoValue, RArray, Value};
+use magnus::{prelude::*, value::Opaque, IntoValue, RArray, Ruby, Value};
 use polars::lazy::dsl;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 use polars_core::series::IsSorted;
-use std::any::Any;
 
-use crate::apply::lazy::map_single;
 use crate::conversion::{parse_fill_null_strategy, Wrap};
+use crate::map::lazy::map_single;
 use crate::rb_exprs_to_exprs;
 use crate::utils::reinterpret;
 use crate::{RbExpr, RbResult};
@@ -163,12 +162,71 @@ impl RbExpr {
             .into()
     }
 
+    pub fn cut(
+        &self,
+        breaks: Vec<f64>,
+        labels: Option<Vec<String>>,
+        left_closed: bool,
+        include_breaks: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .cut(breaks, labels, left_closed, include_breaks)
+            .into()
+    }
+
+    pub fn qcut(
+        &self,
+        probs: Vec<f64>,
+        labels: Option<Vec<String>>,
+        left_closed: bool,
+        allow_duplicates: bool,
+        include_breaks: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .qcut(probs, labels, left_closed, allow_duplicates, include_breaks)
+            .into()
+    }
+
+    pub fn qcut_uniform(
+        &self,
+        n_bins: usize,
+        labels: Option<Vec<String>>,
+        left_closed: bool,
+        allow_duplicates: bool,
+        include_breaks: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .qcut_uniform(
+                n_bins,
+                labels,
+                left_closed,
+                allow_duplicates,
+                include_breaks,
+            )
+            .into()
+    }
+
+    pub fn rle(&self) -> Self {
+        self.inner.clone().rle().into()
+    }
+
+    pub fn rle_id(&self) -> Self {
+        self.inner.clone().rle_id().into()
+    }
+
     pub fn agg_groups(&self) -> Self {
-        self.clone().inner.agg_groups().into()
+        self.inner.clone().agg_groups().into()
     }
 
     pub fn count(&self) -> Self {
-        self.clone().inner.count().into()
+        self.inner.clone().count().into()
+    }
+
+    pub fn len(&self) -> Self {
+        self.inner.clone().len().into()
     }
 
     pub fn value_counts(&self, multithreaded: bool, sorted: bool) -> Self {
@@ -220,12 +278,20 @@ impl RbExpr {
             .into()
     }
 
-    pub fn top_k(&self, k: usize) -> Self {
-        self.inner.clone().top_k(k).into()
+    pub fn top_k(&self, k: &Self) -> Self {
+        self.inner.clone().top_k(k.inner.clone()).into()
+    }
+
+    pub fn bottom_k(&self, k: &Self) -> Self {
+        self.inner.clone().bottom_k(k.inner.clone()).into()
+    }
+
+    pub fn peak_min(&self) -> Self {
+        self.inner.clone().peak_min().into()
     }
 
-    pub fn bottom_k(&self, k: usize) -> Self {
-        self.inner.clone().bottom_k(k).into()
+    pub fn peak_max(&self) -> Self {
+        self.inner.clone().peak_max().into()
     }
 
     pub fn arg_max(&self) -> Self {
@@ -243,8 +309,8 @@ impl RbExpr {
             .into()
     }
 
-    pub fn take(&self, idx: &RbExpr) -> Self {
-        self.clone().inner.take(idx.inner.clone()).into()
+    pub fn gather(&self, idx: &RbExpr) -> Self {
+        self.clone().inner.gather(idx.inner.clone()).into()
     }
 
     pub fn sort_by(&self, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
@@ -260,14 +326,13 @@ impl RbExpr {
         self.clone().inner.forward_fill(limit).into()
     }
 
-    pub fn shift(&self, periods: i64) -> Self {
-        self.clone().inner.shift(periods).into()
-    }
-    pub fn shift_and_fill(&self, periods: i64, fill_value: &RbExpr) -> Self {
-        self.clone()
-            .inner
-            .shift_and_fill(periods, fill_value.inner.clone())
-            .into()
+    pub fn shift(&self, n: &Self, fill_value: Option<&Self>) -> Self {
+        let expr = self.inner.clone();
+        let out = match fill_value {
+            Some(v) => expr.shift_and_fill(n.inner.clone(), v.inner.clone()),
+            None => expr.shift(n.inner.clone()),
+        };
+        out.into()
     }
 
     pub fn fill_null(&self, expr: &RbExpr) -> Self {
@@ -323,26 +388,30 @@ impl RbExpr {
         self.clone().inner.is_unique().into()
     }
 
-    pub fn approx_unique(&self) -> Self {
-        self.clone().inner.approx_unique().into()
+    pub fn approx_n_unique(&self) -> Self {
+        self.clone().inner.approx_n_unique().into()
+    }
+
+    pub fn is_first_distinct(&self) -> Self {
+        self.clone().inner.is_first_distinct().into()
     }
 
-    pub fn is_first(&self) -> Self {
-        self.clone().inner.is_first().into()
+    pub fn is_last_distinct(&self) -> Self {
+        self.clone().inner.is_last_distinct().into()
     }
 
     pub fn explode(&self) -> Self {
         self.clone().inner.explode().into()
     }
 
-    pub fn take_every(&self, n: usize) -> Self {
+    pub fn gather_every(&self, n: usize, offset: usize) -> Self {
         self.clone()
             .inner
             .map(
-                move |s: Series| Ok(Some(s.take_every(n))),
+                move |s: Series| Ok(Some(s.gather_every(n, offset))),
                 GetOutput::same_type(),
             )
-            .with_fmt("take_every")
+            .with_fmt("gather_every")
             .into()
     }
 
@@ -387,20 +456,15 @@ impl RbExpr {
         self.clone().inner.ceil().into()
     }
 
-    pub fn clip(&self, min: Value, max: Value) -> Self {
-        let min = min.try_convert::<Wrap<AnyValue>>().unwrap().0;
-        let max = max.try_convert::<Wrap<AnyValue>>().unwrap().0;
-        self.clone().inner.clip(min, max).into()
-    }
-
-    pub fn clip_min(&self, min: Value) -> Self {
-        let min = min.try_convert::<Wrap<AnyValue>>().unwrap().0;
-        self.clone().inner.clip_min(min).into()
-    }
-
-    pub fn clip_max(&self, max: Value) -> Self {
-        let max = max.try_convert::<Wrap<AnyValue>>().unwrap().0;
-        self.clone().inner.clip_max(max).into()
+    pub fn clip(&self, min: Option<&Self>, max: Option<&Self>) -> Self {
+        let expr = self.inner.clone();
+        let out = match (min, max) {
+            (Some(min), Some(max)) => expr.clip(min.inner.clone(), max.inner.clone()),
+            (Some(min), None) => expr.clip_min(min.inner.clone()),
+            (None, Some(max)) => expr.clip_max(max.inner.clone()),
+            (None, None) => expr,
+        };
+        out.into()
     }
 
     pub fn abs(&self) -> Self {
@@ -492,20 +556,20 @@ impl RbExpr {
         self.clone().inner.pow(exponent.inner.clone()).into()
     }
 
-    pub fn cumsum(&self, reverse: bool) -> Self {
-        self.clone().inner.cumsum(reverse).into()
+    pub fn cum_sum(&self, reverse: bool) -> Self {
+        self.clone().inner.cum_sum(reverse).into()
     }
 
-    pub fn cummax(&self, reverse: bool) -> Self {
-        self.clone().inner.cummax(reverse).into()
+    pub fn cum_max(&self, reverse: bool) -> Self {
+        self.clone().inner.cum_max(reverse).into()
     }
 
-    pub fn cummin(&self, reverse: bool) -> Self {
-        self.clone().inner.cummin(reverse).into()
+    pub fn cum_min(&self, reverse: bool) -> Self {
+        self.clone().inner.cum_min(reverse).into()
     }
 
-    pub fn cumprod(&self, reverse: bool) -> Self {
-        self.clone().inner.cumprod(reverse).into()
+    pub fn cum_prod(&self, reverse: bool) -> Self {
+        self.clone().inner.cum_prod(reverse).into()
     }
 
     pub fn product(&self) -> Self {
@@ -541,33 +605,6 @@ impl RbExpr {
         self.inner.clone().mode().into()
     }
 
-    pub fn keep_name(&self) -> Self {
-        self.inner.clone().keep_name().into()
-    }
-
-    pub fn prefix(&self, prefix: String) -> Self {
-        self.inner.clone().prefix(&prefix).into()
-    }
-
-    pub fn suffix(&self, suffix: String) -> Self {
-        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()
     }
@@ -576,200 +613,6 @@ impl RbExpr {
         self.inner.clone().interpolate(method.0).into()
     }
 
-    pub fn rolling_sum(
-        &self,
-        window_size: String,
-        weights: Option<Vec<f64>>,
-        min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-    ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
-            min_periods,
-            center,
-            by,
-            closed_window: closed.map(|c| c.0),
-            ..Default::default()
-        };
-        self.inner.clone().rolling_sum(options).into()
-    }
-
-    pub fn rolling_min(
-        &self,
-        window_size: String,
-        weights: Option<Vec<f64>>,
-        min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-    ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
-            min_periods,
-            center,
-            by,
-            closed_window: closed.map(|c| c.0),
-            ..Default::default()
-        };
-        self.inner.clone().rolling_min(options).into()
-    }
-
-    pub fn rolling_max(
-        &self,
-        window_size: String,
-        weights: Option<Vec<f64>>,
-        min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-    ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
-            min_periods,
-            center,
-            by,
-            closed_window: closed.map(|c| c.0),
-            ..Default::default()
-        };
-        self.inner.clone().rolling_max(options).into()
-    }
-
-    pub fn rolling_mean(
-        &self,
-        window_size: String,
-        weights: Option<Vec<f64>>,
-        min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-    ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
-            min_periods,
-            center,
-            by,
-            closed_window: closed.map(|c| c.0),
-            ..Default::default()
-        };
-
-        self.inner.clone().rolling_mean(options).into()
-    }
-
-    #[allow(clippy::too_many_arguments)]
-    pub fn rolling_std(
-        &self,
-        window_size: String,
-        weights: Option<Vec<f64>>,
-        min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-        ddof: u8,
-    ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
-            min_periods,
-            center,
-            by,
-            closed_window: closed.map(|c| c.0),
-            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
-        };
-
-        self.inner.clone().rolling_std(options).into()
-    }
-
-    #[allow(clippy::too_many_arguments)]
-    pub fn rolling_var(
-        &self,
-        window_size: String,
-        weights: Option<Vec<f64>>,
-        min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-        ddof: u8,
-    ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
-            min_periods,
-            center,
-            by,
-            closed_window: closed.map(|c| c.0),
-            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
-        };
-
-        self.inner.clone().rolling_var(options).into()
-    }
-
-    pub fn rolling_median(
-        &self,
-        window_size: String,
-        weights: Option<Vec<f64>>,
-        min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-    ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
-            min_periods,
-            center,
-            by,
-            closed_window: closed.map(|c| c.0),
-            fn_params: Some(Arc::new(RollingQuantileParams {
-                prob: 0.5,
-                interpol: QuantileInterpolOptions::Linear,
-            }) as Arc<dyn Any + Send + Sync>),
-        };
-        self.inner.clone().rolling_quantile(options).into()
-    }
-
-    #[allow(clippy::too_many_arguments)]
-    pub fn rolling_quantile(
-        &self,
-        quantile: f64,
-        interpolation: Wrap<QuantileInterpolOptions>,
-        window_size: String,
-        weights: Option<Vec<f64>>,
-        min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-    ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
-            min_periods,
-            center,
-            by,
-            closed_window: closed.map(|c| c.0),
-            fn_params: Some(Arc::new(RollingQuantileParams {
-                prob: quantile,
-                interpol: interpolation.0,
-            }) as Arc<dyn Any + Send + Sync>),
-        };
-
-        self.inner.clone().rolling_quantile(options).into()
-    }
-
-    pub fn rolling_skew(&self, window_size: usize, bias: bool) -> Self {
-        self.inner
-            .clone()
-            .rolling_apply_float(window_size, move |ca| {
-                ca.clone().into_series().skew(bias).unwrap()
-            })
-            .into()
-    }
-
     pub fn lower_bound(&self) -> Self {
         self.inner.clone().lower_bound().into()
     }
@@ -797,8 +640,8 @@ impl RbExpr {
         self.inner.clone().diff(n, null_behavior.0).into()
     }
 
-    pub fn pct_change(&self, n: i64) -> Self {
-        self.inner.clone().pct_change(n).into()
+    pub fn pct_change(&self, n: &Self) -> Self {
+        self.inner.clone().pct_change(n.inner.clone()).into()
     }
 
     pub fn skew(&self, bias: bool) -> Self {
@@ -813,8 +656,8 @@ impl RbExpr {
         self.inner.clone().reshape(&dims).into()
     }
 
-    pub fn cumcount(&self, reverse: bool) -> Self {
-        self.inner.clone().cumcount(reverse).into()
+    pub fn cum_count(&self, reverse: bool) -> Self {
+        self.inner.clone().cum_count(reverse).into()
     }
 
     pub fn to_physical(&self) -> Self {
@@ -828,35 +671,33 @@ impl RbExpr {
             .into()
     }
 
-    pub fn shuffle(&self, seed: Option<u64>, fixed_seed: bool) -> Self {
-        self.inner.clone().shuffle(seed, fixed_seed).into()
+    pub fn shuffle(&self, seed: Option<u64>) -> Self {
+        self.inner.clone().shuffle(seed).into()
     }
 
     pub fn sample_n(
         &self,
-        n: usize,
+        n: &Self,
         with_replacement: bool,
         shuffle: bool,
         seed: Option<u64>,
-        fixed_seed: bool,
     ) -> Self {
         self.inner
             .clone()
-            .sample_n(n, with_replacement, shuffle, seed, fixed_seed)
+            .sample_n(n.inner.clone(), with_replacement, shuffle, seed)
             .into()
     }
 
     pub fn sample_frac(
         &self,
-        frac: f64,
+        frac: &Self,
         with_replacement: bool,
         shuffle: bool,
         seed: Option<u64>,
-        fixed_seed: bool,
     ) -> Self {
         self.inner
             .clone()
-            .sample_frac(frac, with_replacement, shuffle, seed, fixed_seed)
+            .sample_frac(frac.inner.clone(), with_replacement, shuffle, seed)
             .into()
     }
 
@@ -915,11 +756,13 @@ impl RbExpr {
 
     pub fn extend_constant(&self, value: Wrap<AnyValue>, n: usize) -> Self {
         let value = value.into_value();
+        let value = Opaque::from(value);
         self.inner
             .clone()
             .apply(
                 move |s| {
-                    let value = value.try_convert::<Wrap<AnyValue>>().unwrap().0;
+                    let value = Ruby::get().unwrap().get_inner(value);
+                    let value = Wrap::<AnyValue>::try_convert(value).unwrap().0;
                     s.extend_constant(value, n).map(Some)
                 },
                 GetOutput::same_type(),
@@ -928,12 +771,12 @@ impl RbExpr {
             .into()
     }
 
-    pub fn any(&self) -> Self {
-        self.inner.clone().any().into()
+    pub fn any(&self, drop_nulls: bool) -> Self {
+        self.inner.clone().any(drop_nulls).into()
     }
 
-    pub fn all(&self) -> Self {
-        self.inner.clone().all().into()
+    pub fn all(&self, drop_nulls: bool) -> Self {
+        self.inner.clone().all(drop_nulls).into()
     }
 
     pub fn log(&self, base: f64) -> Self {

data/ext/polars/src/expr/list.rs

@@ -23,11 +23,11 @@ impl RbExpr {
             .into()
     }
 
-    pub fn list_count_match(&self, expr: &RbExpr) -> Self {
+    pub fn list_count_matches(&self, expr: &RbExpr) -> Self {
         self.inner
             .clone()
             .list()
-            .count_match(expr.inner.clone())
+            .count_matches(expr.inner.clone())
             .into()
     }
 
@@ -47,12 +47,16 @@ impl RbExpr {
         self.inner.clone().list().get(index.inner.clone()).into()
     }
 
-    pub fn list_join(&self, separator: String) -> Self {
-        self.inner.clone().list().join(&separator).into()
+    pub fn list_join(&self, separator: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .join(separator.inner.clone())
+            .into()
     }
 
-    pub fn list_lengths(&self) -> Self {
-        self.inner.clone().list().lengths().into()
+    pub fn list_len(&self) -> Self {
+        self.inner.clone().list().len().into()
     }
 
     pub fn list_max(&self) -> Self {
@@ -76,8 +80,12 @@ impl RbExpr {
         self.inner.clone().list().reverse().into()
     }
 
-    pub fn list_shift(&self, periods: i64) -> Self {
-        self.inner.clone().list().shift(periods).into()
+    pub fn list_shift(&self, periods: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .shift(periods.inner.clone())
+            .into()
     }
 
     pub fn list_slice(&self, offset: &RbExpr, length: Option<&RbExpr>) -> Self {
@@ -127,7 +135,7 @@ impl RbExpr {
         // let name_gen = name_gen.map(|lambda| {
         //     Arc::new(move |idx: usize| {
         //         let out: Value = lambda.funcall("call", (idx,)).unwrap();
-        //         out.try_convert::<String>().unwrap()
+        //         String::try_convert(out).unwrap()
         //     }) as NameGenerator
         // });
 

data/ext/polars/src/{expr.rs → expr/mod.rs}

@@ -5,10 +5,12 @@ mod datetime;
 mod general;
 mod list;
 mod meta;
+mod name;
+mod rolling;
 mod string;
 mod r#struct;
 
-use magnus::RArray;
+use magnus::{prelude::*, RArray};
 use polars::lazy::dsl::Expr;
 
 use crate::RbResult;
@@ -28,7 +30,7 @@ impl From<Expr> for RbExpr {
 pub fn rb_exprs_to_exprs(rb_exprs: RArray) -> RbResult<Vec<Expr>> {
     let mut exprs = Vec::new();
     for item in rb_exprs.each() {
-        exprs.push(item?.try_convert::<&RbExpr>()?.inner.clone());
+        exprs.push(<&RbExpr>::try_convert(item?)?.inner.clone());
     }
     Ok(exprs)
 }

data/ext/polars/src/expr/name.rs

@@ -0,0 +1,44 @@
+use magnus::{block::Proc, value::Opaque, Ruby};
+use polars::prelude::*;
+
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn name_keep(&self) -> Self {
+        self.inner.clone().name().keep().into()
+    }
+
+    pub fn name_map(&self, lambda: Proc) -> Self {
+        let lambda = Opaque::from(lambda);
+        self.inner
+            .clone()
+            .name()
+            .map(move |name| {
+                let lambda = Ruby::get().unwrap().get_inner(lambda);
+                let out = lambda.call::<_, String>((name,));
+                match out {
+                    Ok(out) => Ok(out),
+                    Err(e) => Err(PolarsError::ComputeError(
+                        format!("Ruby function in 'name.map' produced an error: {}.", e).into(),
+                    )),
+                }
+            })
+            .into()
+    }
+
+    pub fn name_prefix(&self, prefix: String) -> Self {
+        self.inner.clone().name().prefix(&prefix).into()
+    }
+
+    pub fn name_suffix(&self, suffix: String) -> Self {
+        self.inner.clone().name().suffix(&suffix).into()
+    }
+
+    pub fn name_to_lowercase(&self) -> Self {
+        self.inner.clone().name().to_lowercase().into()
+    }
+
+    pub fn name_to_uppercase(&self) -> Self {
+        self.inner.clone().name().to_uppercase().into()
+    }
+}