polars-df 0.15.0 → 0.26.0

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

@@ -1,15 +1,17 @@
 use std::ops::Neg;
 
-use magnus::{prelude::*, value::Opaque, IntoValue, RArray, Ruby, Value};
+use magnus::RArray;
 use polars::lazy::dsl;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
-use polars_core::series::IsSorted;
+use polars_core::chunked_array::cast::CastOptions;
+use polars_plan::plans::AExprSorted;
 
-use crate::conversion::{parse_fill_null_strategy, Wrap};
-use crate::map::lazy::map_single;
-use crate::rb_exprs_to_exprs;
-use crate::{RbExpr, RbResult};
+use super::datatype::RbDataTypeExpr;
+use super::selector::RbSelector;
+use crate::conversion::{Wrap, parse_fill_null_strategy};
+use crate::expr::ToExprs;
+use crate::{RbDataType, RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
     pub fn add(&self, rhs: &Self) -> RbResult<Self> {
@@ -116,6 +118,14 @@ impl RbExpr {
         self.inner.clone().max().into()
     }
 
+    pub fn min_by(&self, by: &Self) -> Self {
+        self.inner.clone().min_by(by.inner.clone()).into()
+    }
+
+    pub fn max_by(&self, by: &Self) -> Self {
+        self.inner.clone().max_by(by.inner.clone()).into()
+    }
+
     pub fn nan_max(&self) -> Self {
         self.inner.clone().nan_max().into()
     }
@@ -152,16 +162,28 @@ impl RbExpr {
         self.inner.clone().unique_stable().into()
     }
 
-    pub fn first(&self) -> Self {
-        self.inner.clone().first().into()
+    pub fn first(&self, ignore_nulls: bool) -> Self {
+        if ignore_nulls {
+            self.inner.clone().first_non_null().into()
+        } else {
+            self.inner.clone().first().into()
+        }
     }
 
-    pub fn last(&self) -> Self {
-        self.inner.clone().last().into()
+    pub fn last(&self, ignore_nulls: bool) -> Self {
+        if ignore_nulls {
+            self.inner.clone().last_non_null().into()
+        } else {
+            self.inner.clone().last().into()
+        }
+    }
+
+    pub fn item(&self, allow_empty: bool) -> Self {
+        self.inner.clone().item(allow_empty).into()
     }
 
-    pub fn implode(&self) -> Self {
-        self.inner.clone().implode().into()
+    pub fn implode(&self, maintain_order: bool) -> Self {
+        self.inner.clone().implode(maintain_order).into()
     }
 
     pub fn quantile(&self, quantile: &Self, interpolation: Wrap<QuantileMethod>) -> Self {
@@ -253,14 +275,20 @@ impl RbExpr {
         self.inner.clone().null_count().into()
     }
 
-    pub fn cast(&self, data_type: Wrap<DataType>, strict: bool) -> RbResult<Self> {
-        let dt = data_type.0;
-        let expr = if strict {
-            self.inner.clone().strict_cast(dt)
+    pub fn cast(&self, dtype: &RbDataTypeExpr, strict: bool, wrap_numerical: bool) -> Self {
+        let options = if wrap_numerical {
+            CastOptions::Overflowing
+        } else if strict {
+            CastOptions::Strict
         } else {
-            self.inner.clone().cast(dt)
+            CastOptions::NonStrict
         };
-        Ok(expr.into())
+
+        let expr = self
+            .inner
+            .clone()
+            .cast_with_options(dtype.inner.clone(), options);
+        expr.into()
     }
 
     pub fn sort_with(&self, descending: bool, nulls_last: bool) -> Self {
@@ -271,20 +299,13 @@ impl RbExpr {
                 nulls_last,
                 multithreaded: true,
                 maintain_order: false,
+                limit: None,
             })
             .into()
     }
 
-    pub fn arg_sort(&self, reverse: bool, nulls_last: bool) -> Self {
-        self.clone()
-            .inner
-            .arg_sort(SortOptions {
-                descending: reverse,
-                nulls_last,
-                multithreaded: true,
-                maintain_order: false,
-            })
-            .into()
+    pub fn arg_sort(&self, descending: bool, nulls_last: bool) -> Self {
+        self.inner.clone().arg_sort(descending, nulls_last).into()
     }
 
     pub fn top_k(&self, k: &Self) -> Self {
@@ -292,7 +313,7 @@ impl RbExpr {
     }
 
     pub fn top_k_by(&self, by: RArray, k: &Self, reverse: Vec<bool>) -> RbResult<Self> {
-        let by = rb_exprs_to_exprs(by)?;
+        let by = by.to_exprs()?;
         Ok(self
             .inner
             .clone()
@@ -305,7 +326,7 @@ impl RbExpr {
     }
 
     pub fn bottom_k_by(&self, by: RArray, k: &Self, reverse: Vec<bool>) -> RbResult<Self> {
-        let by = rb_exprs_to_exprs(by)?;
+        let by = by.to_exprs()?;
         Ok(self
             .inner
             .clone()
@@ -329,19 +350,34 @@ impl RbExpr {
         self.inner.clone().arg_min().into()
     }
 
-    pub fn search_sorted(&self, element: &Self, side: Wrap<SearchSortedSide>) -> Self {
+    pub fn index_of(&self, element: &Self) -> Self {
+        self.inner.clone().index_of(element.inner.clone()).into()
+    }
+
+    pub fn search_sorted(
+        &self,
+        element: &Self,
+        side: Wrap<SearchSortedSide>,
+        descending: bool,
+    ) -> Self {
         self.inner
             .clone()
-            .search_sorted(element.inner.clone(), side.0)
+            .search_sorted(element.inner.clone(), side.0, descending)
             .into()
     }
 
-    pub fn gather(&self, idx: &Self) -> Self {
-        self.inner.clone().gather(idx.inner.clone()).into()
+    pub fn gather(&self, idx: &Self, null_on_oob: bool) -> Self {
+        self.inner
+            .clone()
+            .gather(idx.inner.clone(), null_on_oob)
+            .into()
     }
 
-    pub fn get(&self, idx: &Self) -> Self {
-        self.inner.clone().get(idx.inner.clone()).into()
+    pub fn get(&self, idx: &Self, null_on_oob: bool) -> Self {
+        self.inner
+            .clone()
+            .get(idx.inner.clone(), null_on_oob)
+            .into()
     }
 
     pub fn sort_by(
@@ -352,7 +388,7 @@ impl RbExpr {
         multithreaded: bool,
         maintain_order: bool,
     ) -> RbResult<Self> {
-        let by = rb_exprs_to_exprs(by)?;
+        let by = by.to_exprs()?;
         Ok(self
             .inner
             .clone()
@@ -363,19 +399,12 @@ impl RbExpr {
                     nulls_last,
                     multithreaded,
                     maintain_order,
+                    limit: None,
                 },
             )
             .into())
     }
 
-    pub fn backward_fill(&self, limit: FillNullLimit) -> Self {
-        self.inner.clone().backward_fill(limit).into()
-    }
-
-    pub fn forward_fill(&self, limit: FillNullLimit) -> Self {
-        self.inner.clone().forward_fill(limit).into()
-    }
-
     pub fn shift(&self, n: &Self, fill_value: Option<&Self>) -> Self {
         let expr = self.inner.clone();
         let out = match fill_value {
@@ -394,16 +423,8 @@ impl RbExpr {
         strategy: String,
         limit: FillNullLimit,
     ) -> RbResult<Self> {
-        let strat = parse_fill_null_strategy(&strategy, limit)?;
-        Ok(self
-            .inner
-            .clone()
-            .apply(
-                move |s| s.fill_null(strat).map(Some),
-                GetOutput::same_type(),
-            )
-            .with_fmt("fill_null_with_strategy")
-            .into())
+        let strategy = parse_fill_null_strategy(&strategy, limit)?;
+        Ok(self.inner.clone().fill_null_with_strategy(strategy).into())
     }
 
     pub fn fill_nan(&self, expr: &Self) -> Self {
@@ -445,6 +466,13 @@ impl RbExpr {
             .into()
     }
 
+    pub fn is_close(&self, other: &Self, abs_tol: f64, rel_tol: f64, nans_equal: bool) -> Self {
+        self.inner
+            .clone()
+            .is_close(other.inner.clone(), abs_tol, rel_tol, nans_equal)
+            .into()
+    }
+
     pub fn approx_n_unique(&self) -> Self {
         self.inner.clone().approx_n_unique().into()
     }
@@ -457,8 +485,14 @@ impl RbExpr {
         self.inner.clone().is_last_distinct().into()
     }
 
-    pub fn explode(&self) -> Self {
-        self.inner.clone().explode().into()
+    pub fn explode(&self, empty_as_null: bool, keep_nulls: bool) -> Self {
+        self.inner
+            .clone()
+            .explode(ExplodeOptions {
+                empty_as_null,
+                keep_nulls,
+            })
+            .into()
     }
 
     pub fn gather_every(&self, n: usize, offset: usize) -> Self {
@@ -488,14 +522,19 @@ impl RbExpr {
     }
 
     pub fn rechunk(&self) -> Self {
-        self.inner
-            .clone()
-            .map(|s| Ok(Some(s.rechunk())), GetOutput::same_type())
-            .into()
+        self.inner.clone().rechunk().into()
+    }
+
+    pub fn round(&self, decimals: u32, mode: Wrap<RoundMode>) -> Self {
+        self.inner.clone().round(decimals, mode.0).into()
     }
 
-    pub fn round(&self, decimals: u32) -> Self {
-        self.inner.clone().round(decimals).into()
+    pub fn round_sig_figs(&self, digits: i32) -> Self {
+        self.clone().inner.round_sig_figs(digits).into()
+    }
+
+    pub fn truncate(&self, decimals: u32) -> Self {
+        self.inner.clone().truncate(decimals).into()
     }
 
     pub fn floor(&self) -> Self {
@@ -533,6 +572,10 @@ impl RbExpr {
         self.inner.clone().tan().into()
     }
 
+    pub fn cot(&self) -> Self {
+        self.inner.clone().cot().into()
+    }
+
     pub fn arcsin(&self) -> Self {
         self.inner.clone().arcsin().into()
     }
@@ -569,6 +612,14 @@ impl RbExpr {
         self.inner.clone().arctanh().into()
     }
 
+    pub fn degrees(&self) -> Self {
+        self.inner.clone().degrees().into()
+    }
+
+    pub fn radians(&self) -> Self {
+        self.inner.clone().radians().into()
+    }
+
     pub fn sign(&self) -> Self {
         self.inner.clone().sign().into()
     }
@@ -577,25 +628,70 @@ impl RbExpr {
         self.inner.clone().is_duplicated().into()
     }
 
-    pub fn over(&self, partition_by: RArray) -> RbResult<Self> {
-        let partition_by = rb_exprs_to_exprs(partition_by)?;
-        Ok(self.inner.clone().over(partition_by).into())
+    pub fn over(
+        &self,
+        partition_by: Option<RArray>,
+        order_by: Option<RArray>,
+        order_by_descending: bool,
+        order_by_nulls_last: bool,
+        mapping_strategy: Wrap<WindowMapping>,
+    ) -> RbResult<Self> {
+        let partition_by = partition_by.map(|v| v.to_exprs()).transpose()?;
+
+        let order_by = order_by.map(|v| v.to_exprs()).transpose()?.map(|order_by| {
+            (
+                order_by,
+                SortOptions {
+                    descending: order_by_descending,
+                    nulls_last: order_by_nulls_last,
+                    ..Default::default()
+                },
+            )
+        });
+
+        Ok(self
+            .inner
+            .clone()
+            .over_with_options(partition_by, order_by, mapping_strategy.0)
+            .map_err(RbPolarsErr::from)?
+            .into())
+    }
+
+    pub fn rolling(
+        &self,
+        index_column: &RbExpr,
+        period: String,
+        offset: String,
+        closed: Wrap<ClosedWindow>,
+    ) -> RbResult<Self> {
+        let period = Duration::try_parse(&period).map_err(RbPolarsErr::from)?;
+        let offset = Duration::try_parse(&offset).map_err(RbPolarsErr::from)?;
+        let closed = closed.0;
+
+        Ok(self
+            .inner
+            .clone()
+            .rolling(index_column.inner.clone(), period, offset, closed)
+            .into())
     }
 
-    pub fn _and(&self, expr: &Self) -> Self {
+    pub fn and_(&self, expr: &Self) -> Self {
         self.inner.clone().and(expr.inner.clone()).into()
     }
 
-    pub fn _xor(&self, expr: &Self) -> Self {
-        self.inner.clone().xor(expr.inner.clone()).into()
+    pub fn or_(&self, expr: &Self) -> Self {
+        self.inner.clone().or(expr.inner.clone()).into()
     }
 
-    pub fn _or(&self, expr: &Self) -> Self {
-        self.inner.clone().or(expr.inner.clone()).into()
+    pub fn xor_(&self, expr: &Self) -> Self {
+        self.inner.clone().xor(expr.inner.clone()).into()
     }
 
-    pub fn is_in(&self, expr: &Self) -> Self {
-        self.inner.clone().is_in(expr.inner.clone()).into()
+    pub fn is_in(&self, expr: &Self, nulls_equal: bool) -> Self {
+        self.inner
+            .clone()
+            .is_in(expr.inner.clone(), nulls_equal)
+            .into()
     }
 
     pub fn repeat_by(&self, by: &Self) -> Self {
@@ -606,6 +702,14 @@ impl RbExpr {
         self.inner.clone().pow(exponent.inner.clone()).into()
     }
 
+    pub fn sqrt(&self) -> Self {
+        self.inner.clone().sqrt().into()
+    }
+
+    pub fn cbrt(&self) -> Self {
+        self.inner.clone().cbrt().into()
+    }
+
     pub fn cum_sum(&self, reverse: bool) -> Self {
         self.inner.clone().cum_sum(reverse).into()
     }
@@ -626,48 +730,29 @@ impl RbExpr {
         self.inner.clone().product().into()
     }
 
-    pub fn shrink_dtype(&self) -> Self {
-        self.inner.clone().shrink_dtype().into()
-    }
-
-    pub fn map_batches(
-        &self,
-        lambda: Value,
-        output_type: Option<Wrap<DataType>>,
-        agg_list: bool,
-        is_elementwise: bool,
-        returns_scalar: bool,
-    ) -> Self {
-        map_single(
-            self,
-            lambda,
-            output_type,
-            agg_list,
-            is_elementwise,
-            returns_scalar,
-        )
-    }
-
     pub fn dot(&self, other: &Self) -> Self {
         self.inner.clone().dot(other.inner.clone()).into()
     }
 
-    pub fn reinterpret(&self, signed: bool) -> Self {
-        self.inner.clone().reinterpret(signed).into()
-    }
-
-    pub fn mode(&self) -> Self {
-        self.inner.clone().mode().into()
+    pub fn reinterpret(&self, signed: Option<bool>, dtype: Option<RbDataType>) -> Self {
+        self.inner
+            .clone()
+            .reinterpret(signed, dtype.map(|dt| dt.0))
+            .into()
     }
 
-    pub fn exclude(&self, columns: Vec<String>) -> Self {
-        self.inner.clone().exclude(columns).into()
+    pub fn mode(&self, maintain_order: bool) -> Self {
+        self.inner.clone().mode(maintain_order).into()
     }
 
     pub fn interpolate(&self, method: Wrap<InterpolationMethod>) -> Self {
         self.inner.clone().interpolate(method.0).into()
     }
 
+    pub fn interpolate_by(&self, by: &Self) -> Self {
+        self.inner.clone().interpolate_by(by.inner.clone()).into()
+    }
+
     pub fn lower_bound(&self) -> Self {
         self.inner.clone().lower_bound().into()
     }
@@ -676,10 +761,10 @@ impl RbExpr {
         self.inner.clone().upper_bound().into()
     }
 
-    pub fn cumulative_eval(&self, expr: &Self, min_periods: usize, parallel: bool) -> Self {
+    pub fn cumulative_eval(&self, expr: &Self, min_samples: usize) -> Self {
         self.inner
             .clone()
-            .cumulative_eval(expr.inner.clone(), min_periods, parallel)
+            .cumulative_eval(expr.inner.clone(), min_samples)
             .into()
     }
 
@@ -691,8 +776,11 @@ impl RbExpr {
         self.inner.clone().rank(options, seed).into()
     }
 
-    pub fn diff(&self, n: i64, null_behavior: Wrap<NullBehavior>) -> Self {
-        self.inner.clone().diff(n, null_behavior.0).into()
+    pub fn diff(&self, n: &Self, null_behavior: Wrap<NullBehavior>) -> Self {
+        self.inner
+            .clone()
+            .diff(n.inner.clone(), null_behavior.0)
+            .into()
     }
 
     pub fn pct_change(&self, n: &Self) -> Self {
@@ -766,6 +854,15 @@ impl RbExpr {
         self.inner.clone().ewm_mean(options).into()
     }
 
+    pub fn ewm_mean_by(&self, times: &RbExpr, half_life: String) -> RbResult<Self> {
+        let half_life = Duration::try_parse(&half_life).map_err(RbPolarsErr::from)?;
+        Ok(self
+            .inner
+            .clone()
+            .ewm_mean_by(times.inner.clone(), half_life)
+            .into())
+    }
+
     pub fn ewm_std(
         &self,
         alpha: f64,
@@ -802,20 +899,10 @@ impl RbExpr {
         self.inner.clone().ewm_var(options).into()
     }
 
-    pub fn extend_constant(&self, value: Wrap<AnyValue>, n: usize) -> Self {
-        let value = value.into_value();
-        let value = Opaque::from(value);
+    pub fn extend_constant(&self, value: &Self, n: &Self) -> Self {
         self.inner
             .clone()
-            .apply(
-                move |s| {
-                    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(),
-            )
-            .with_fmt("extend")
+            .extend_constant(value.inner.clone(), n.inner.clone())
             .into()
     }
 
@@ -827,8 +914,16 @@ impl RbExpr {
         self.inner.clone().all(drop_nulls).into()
     }
 
-    pub fn log(&self, base: f64) -> Self {
-        self.inner.clone().log(base).into()
+    pub fn is_empty(&self, ignore_nulls: bool) -> Self {
+        self.inner.clone().is_empty(ignore_nulls).into()
+    }
+
+    pub fn log(&self, base: &RbExpr) -> Self {
+        self.inner.clone().log(base.inner.clone()).into()
+    }
+
+    pub fn log1p(&self) -> Self {
+        self.inner.clone().log1p().into()
     }
 
     pub fn exp(&self) -> Self {
@@ -843,13 +938,11 @@ impl RbExpr {
         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()
+    pub fn set_sorted_flag(&self, descending: bool, nulls_last: bool) -> Self {
+        let sortedness = AExprSorted::default()
+            .with_desc(Some(descending))
+            .with_nulls_last(Some(nulls_last));
+        self.inner.clone().set_sorted_flag(sortedness).into()
     }
 
     pub fn replace(&self, old: &Self, new: &Self) -> Self {
@@ -876,4 +969,35 @@ impl RbExpr {
             )
             .into()
     }
+
+    pub fn hist(
+        &self,
+        bins: Option<&RbExpr>,
+        bin_count: Option<usize>,
+        include_category: bool,
+        include_breakpoint: bool,
+    ) -> Self {
+        let bins = bins.map(|e| e.inner.clone());
+        self.inner
+            .clone()
+            .hist(bins, bin_count, include_category, include_breakpoint)
+            .into()
+    }
+
+    #[allow(clippy::wrong_self_convention)]
+    pub fn into_selector(&self) -> RbResult<RbSelector> {
+        Ok(self
+            .inner
+            .clone()
+            .into_selector()
+            .ok_or_else(
+                || polars_err!(InvalidOperation: "expr `{}` is not a selector", &self.inner),
+            )
+            .map_err(RbPolarsErr::from)?
+            .into())
+    }
+
+    pub fn new_selector(selector: &RbSelector) -> Self {
+        Expr::Selector(selector.inner.clone()).into()
+    }
 }