polars-df 0.5.0 → 0.7.0

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

@@ -1,10 +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 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};
@@ -201,6 +202,7 @@ impl RbExpr {
                 descending,
                 nulls_last,
                 multithreaded: true,
+                maintain_order: false,
             })
             .into()
     }
@@ -212,16 +214,25 @@ impl RbExpr {
                 descending: reverse,
                 nulls_last,
                 multithreaded: true,
+                maintain_order: false,
             })
             .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: usize) -> Self {
-        self.inner.clone().bottom_k(k).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 peak_max(&self) -> Self {
+        self.inner.clone().peak_max().into()
     }
 
     pub fn arg_max(&self) -> Self {
@@ -239,8 +250,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> {
@@ -256,14 +267,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 {
@@ -319,26 +329,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) -> Self {
         self.clone()
             .inner
             .map(
-                move |s: Series| Ok(Some(s.take_every(n))),
+                move |s: Series| Ok(Some(s.gather_every(n))),
                 GetOutput::same_type(),
             )
-            .with_fmt("take_every")
+            .with_fmt("gather_every")
             .into()
     }
 
@@ -383,20 +397,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 {
@@ -488,20 +497,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 {
@@ -537,33 +546,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()
     }
@@ -572,185 +554,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),
-        };
-        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),
-        };
-        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),
-        };
-        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),
-        };
-
-        self.inner.clone().rolling_mean(options).into()
-    }
-
-    pub fn rolling_std(
-        &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),
-        };
-
-        self.inner.clone().rolling_std(options).into()
-    }
-
-    pub fn rolling_var(
-        &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),
-        };
-
-        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),
-        };
-        self.inner.clone().rolling_median(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),
-        };
-
-        self.inner
-            .clone()
-            .rolling_quantile(quantile, interpolation.0, 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()
     }
@@ -778,8 +581,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 {
@@ -794,8 +597,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 {
@@ -815,27 +618,27 @@ impl RbExpr {
 
     pub fn sample_n(
         &self,
-        n: usize,
+        n: &Self,
         with_replacement: bool,
         shuffle: bool,
         seed: Option<u64>,
     ) -> Self {
         self.inner
             .clone()
-            .sample_n(n, with_replacement, shuffle, 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>,
     ) -> Self {
         self.inner
             .clone()
-            .sample_frac(frac, with_replacement, shuffle, seed)
+            .sample_frac(frac.inner.clone(), with_replacement, shuffle, seed)
             .into()
     }
 
@@ -894,11 +697,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(),
@@ -907,12 +712,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 {
@@ -930,4 +735,13 @@ impl RbExpr {
     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 set_sorted_flag(&self, descending: bool) -> Self {
+        let is_sorted = if descending {
+            IsSorted::Descending
+        } else {
+            IsSorted::Ascending
+        };
+        self.inner.clone().set_sorted_flag(is_sorted).into()
+    }
 }

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

@@ -8,71 +8,84 @@ use crate::{RbExpr, RbResult};
 
 impl RbExpr {
     pub fn list_arg_max(&self) -> Self {
-        self.inner.clone().arr().arg_max().into()
+        self.inner.clone().list().arg_max().into()
     }
 
     pub fn list_arg_min(&self) -> Self {
-        self.inner.clone().arr().arg_min().into()
+        self.inner.clone().list().arg_min().into()
     }
 
     pub fn list_contains(&self, other: &RbExpr) -> Self {
         self.inner
             .clone()
-            .arr()
+            .list()
             .contains(other.inner.clone())
             .into()
     }
 
-    pub fn list_count_match(&self, expr: &RbExpr) -> Self {
+    pub fn list_count_matches(&self, expr: &RbExpr) -> Self {
         self.inner
             .clone()
-            .arr()
-            .count_match(expr.inner.clone())
+            .list()
+            .count_matches(expr.inner.clone())
             .into()
     }
 
     pub fn list_diff(&self, n: i64, null_behavior: Wrap<NullBehavior>) -> RbResult<Self> {
-        Ok(self.inner.clone().arr().diff(n, null_behavior.0).into())
+        Ok(self.inner.clone().list().diff(n, null_behavior.0).into())
     }
 
     pub fn list_eval(&self, expr: &RbExpr, parallel: bool) -> Self {
         self.inner
             .clone()
-            .arr()
+            .list()
             .eval(expr.inner.clone(), parallel)
             .into()
     }
 
     pub fn list_get(&self, index: &RbExpr) -> Self {
-        self.inner.clone().arr().get(index.inner.clone()).into()
+        self.inner.clone().list().get(index.inner.clone()).into()
     }
 
-    pub fn list_join(&self, separator: String) -> Self {
-        self.inner.clone().arr().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().arr().lengths().into()
+    pub fn list_len(&self) -> Self {
+        self.inner.clone().list().len().into()
     }
 
     pub fn list_max(&self) -> Self {
-        self.inner.clone().arr().max().into()
+        self.inner.clone().list().max().into()
     }
 
     pub fn list_mean(&self) -> Self {
-        self.inner.clone().arr().mean().with_fmt("arr.mean").into()
+        self.inner
+            .clone()
+            .list()
+            .mean()
+            .with_fmt("list.mean")
+            .into()
     }
 
     pub fn list_min(&self) -> Self {
-        self.inner.clone().arr().min().into()
+        self.inner.clone().list().min().into()
     }
 
     pub fn list_reverse(&self) -> Self {
-        self.inner.clone().arr().reverse().into()
+        self.inner.clone().list().reverse().into()
     }
 
-    pub fn list_shift(&self, periods: i64) -> Self {
-        self.inner.clone().arr().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 {
@@ -82,7 +95,7 @@ impl RbExpr {
         };
         self.inner
             .clone()
-            .arr()
+            .list()
             .slice(offset.inner.clone(), length)
             .into()
     }
@@ -90,23 +103,23 @@ impl RbExpr {
     pub fn list_sort(&self, reverse: bool) -> Self {
         self.inner
             .clone()
-            .arr()
+            .list()
             .sort(SortOptions {
                 descending: reverse,
                 ..Default::default()
             })
-            .with_fmt("arr.sort")
+            .with_fmt("list.sort")
             .into()
     }
 
     pub fn list_sum(&self) -> Self {
-        self.inner.clone().arr().sum().with_fmt("arr.sum").into()
+        self.inner.clone().list().sum().with_fmt("list.sum").into()
     }
 
     pub fn list_take(&self, index: &RbExpr, null_on_oob: bool) -> Self {
         self.inner
             .clone()
-            .arr()
+            .list()
             .take(index.inner.clone(), null_on_oob)
             .into()
     }
@@ -122,14 +135,14 @@ 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
         // });
 
         Ok(self
             .inner
             .clone()
-            .arr()
+            .list()
             .to_struct(width_strat.0, name_gen, upper_bound)
             .into())
     }
@@ -138,9 +151,9 @@ impl RbExpr {
         let e = self.inner.clone();
 
         if maintain_order {
-            e.arr().unique_stable().into()
+            e.list().unique_stable().into()
         } else {
-            e.arr().unique().into()
+            e.list().unique().into()
         }
     }
 }

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

@@ -1,13 +1,16 @@
+mod array;
 mod binary;
 mod categorical;
 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;
@@ -27,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)
 }