polars-df 0.7.0 → 0.9.0

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

@@ -1,3 +1,5 @@
+use std::ops::Neg;
+
 use magnus::{prelude::*, value::Opaque, IntoValue, RArray, Ruby, Value};
 use polars::lazy::dsl;
 use polars::prelude::*;
@@ -11,55 +13,67 @@ use crate::utils::reinterpret;
 use crate::{RbExpr, RbResult};
 
 impl RbExpr {
-    pub fn add(&self, rhs: &RbExpr) -> RbResult<Self> {
+    pub fn add(&self, rhs: &Self) -> RbResult<Self> {
         Ok(dsl::binary_expr(self.inner.clone(), Operator::Plus, rhs.inner.clone()).into())
     }
 
-    pub fn sub(&self, rhs: &RbExpr) -> RbResult<Self> {
+    pub fn sub(&self, rhs: &Self) -> RbResult<Self> {
         Ok(dsl::binary_expr(self.inner.clone(), Operator::Minus, rhs.inner.clone()).into())
     }
 
-    pub fn mul(&self, rhs: &RbExpr) -> RbResult<Self> {
+    pub fn mul(&self, rhs: &Self) -> RbResult<Self> {
         Ok(dsl::binary_expr(self.inner.clone(), Operator::Multiply, rhs.inner.clone()).into())
     }
 
-    pub fn truediv(&self, rhs: &RbExpr) -> RbResult<Self> {
+    pub fn truediv(&self, rhs: &Self) -> RbResult<Self> {
         Ok(dsl::binary_expr(self.inner.clone(), Operator::TrueDivide, rhs.inner.clone()).into())
     }
 
-    pub fn _mod(&self, rhs: &RbExpr) -> RbResult<Self> {
+    pub fn _mod(&self, rhs: &Self) -> RbResult<Self> {
         Ok(dsl::binary_expr(self.inner.clone(), Operator::Modulus, rhs.inner.clone()).into())
     }
 
-    pub fn floordiv(&self, rhs: &RbExpr) -> RbResult<Self> {
+    pub fn floordiv(&self, rhs: &Self) -> RbResult<Self> {
         Ok(dsl::binary_expr(self.inner.clone(), Operator::FloorDivide, rhs.inner.clone()).into())
     }
 
+    pub fn neg(&self) -> RbResult<Self> {
+        Ok(self.inner.clone().neg().into())
+    }
+
     pub fn to_str(&self) -> String {
         format!("{:?}", self.inner)
     }
 
-    pub fn eq(&self, other: &RbExpr) -> Self {
+    pub fn eq(&self, other: &Self) -> Self {
         self.clone().inner.eq(other.inner.clone()).into()
     }
 
-    pub fn neq(&self, other: &RbExpr) -> Self {
+    pub fn eq_missing(&self, other: &Self) -> Self {
+        self.inner.clone().eq_missing(other.inner.clone()).into()
+    }
+
+    pub fn neq(&self, other: &Self) -> Self {
         self.clone().inner.neq(other.inner.clone()).into()
     }
 
-    pub fn gt(&self, other: &RbExpr) -> Self {
+    pub fn neq_missing(&self, other: &Self) -> Self {
+        self.inner.clone().neq_missing(other.inner.clone()).into()
+    }
+
+    pub fn gt(&self, other: &Self) -> Self {
         self.clone().inner.gt(other.inner.clone()).into()
     }
 
-    pub fn gt_eq(&self, other: &RbExpr) -> Self {
+    pub fn gt_eq(&self, other: &Self) -> Self {
         self.clone().inner.gt_eq(other.inner.clone()).into()
     }
 
-    pub fn lt_eq(&self, other: &RbExpr) -> Self {
+    pub fn lt_eq(&self, other: &Self) -> Self {
         self.clone().inner.lt_eq(other.inner.clone()).into()
     }
 
-    pub fn lt(&self, other: &RbExpr) -> Self {
+    pub fn lt(&self, other: &Self) -> Self {
         self.clone().inner.lt(other.inner.clone()).into()
     }
 
@@ -67,7 +81,7 @@ impl RbExpr {
         self.clone().inner.alias(&name).into()
     }
 
-    pub fn is_not(&self) -> Self {
+    pub fn not_(&self) -> Self {
         self.clone().inner.not().into()
     }
 
@@ -151,23 +165,78 @@ impl RbExpr {
         self.clone().inner.implode().into()
     }
 
-    pub fn quantile(
-        &self,
-        quantile: &RbExpr,
-        interpolation: Wrap<QuantileInterpolOptions>,
-    ) -> Self {
+    pub fn quantile(&self, quantile: &Self, interpolation: Wrap<QuantileInterpolOptions>) -> Self {
         self.clone()
             .inner
             .quantile(quantile.inner.clone(), interpolation.0)
             .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 {
@@ -243,14 +312,14 @@ impl RbExpr {
         self.clone().inner.arg_min().into()
     }
 
-    pub fn search_sorted(&self, element: &RbExpr, side: Wrap<SearchSortedSide>) -> Self {
+    pub fn search_sorted(&self, element: &Self, side: Wrap<SearchSortedSide>) -> Self {
         self.inner
             .clone()
             .search_sorted(element.inner.clone(), side.0)
             .into()
     }
 
-    pub fn gather(&self, idx: &RbExpr) -> Self {
+    pub fn gather(&self, idx: &Self) -> Self {
         self.clone().inner.gather(idx.inner.clone()).into()
     }
 
@@ -276,7 +345,7 @@ impl RbExpr {
         out.into()
     }
 
-    pub fn fill_null(&self, expr: &RbExpr) -> Self {
+    pub fn fill_null(&self, expr: &Self) -> Self {
         self.clone().inner.fill_null(expr.inner.clone()).into()
     }
 
@@ -297,7 +366,7 @@ impl RbExpr {
             .into())
     }
 
-    pub fn fill_nan(&self, expr: &RbExpr) -> Self {
+    pub fn fill_nan(&self, expr: &Self) -> Self {
         self.inner.clone().fill_nan(expr.inner.clone()).into()
     }
 
@@ -309,7 +378,7 @@ impl RbExpr {
         self.inner.clone().drop_nans().into()
     }
 
-    pub fn filter(&self, predicate: &RbExpr) -> Self {
+    pub fn filter(&self, predicate: &Self) -> Self {
         self.clone().inner.filter(predicate.inner.clone()).into()
     }
 
@@ -345,11 +414,11 @@ impl RbExpr {
         self.clone().inner.explode().into()
     }
 
-    pub fn gather_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.gather_every(n))),
+                move |s: Series| Ok(Some(s.gather_every(n, offset))),
                 GetOutput::same_type(),
             )
             .with_fmt("gather_every")
@@ -364,14 +433,14 @@ impl RbExpr {
         self.clone().inner.head(n).into()
     }
 
-    pub fn slice(&self, offset: &RbExpr, length: &RbExpr) -> Self {
+    pub fn slice(&self, offset: &Self, length: &Self) -> Self {
         self.inner
             .clone()
             .slice(offset.inner.clone(), length.inner.clone())
             .into()
     }
 
-    pub fn append(&self, other: &RbExpr, upcast: bool) -> Self {
+    pub fn append(&self, other: &Self, upcast: bool) -> Self {
         self.inner
             .clone()
             .append(other.inner.clone(), upcast)
@@ -473,27 +542,27 @@ impl RbExpr {
         Ok(self.clone().inner.over(partition_by).into())
     }
 
-    pub fn _and(&self, expr: &RbExpr) -> Self {
+    pub fn _and(&self, expr: &Self) -> Self {
         self.clone().inner.and(expr.inner.clone()).into()
     }
 
-    pub fn _xor(&self, expr: &RbExpr) -> Self {
+    pub fn _xor(&self, expr: &Self) -> Self {
         self.clone().inner.xor(expr.inner.clone()).into()
     }
 
-    pub fn _or(&self, expr: &RbExpr) -> Self {
+    pub fn _or(&self, expr: &Self) -> Self {
         self.clone().inner.or(expr.inner.clone()).into()
     }
 
-    pub fn is_in(&self, expr: &RbExpr) -> Self {
+    pub fn is_in(&self, expr: &Self) -> Self {
         self.clone().inner.is_in(expr.inner.clone()).into()
     }
 
-    pub fn repeat_by(&self, by: &RbExpr) -> Self {
+    pub fn repeat_by(&self, by: &Self) -> Self {
         self.clone().inner.repeat_by(by.inner.clone()).into()
     }
 
-    pub fn pow(&self, exponent: &RbExpr) -> Self {
+    pub fn pow(&self, exponent: &Self) -> Self {
         self.clone().inner.pow(exponent.inner.clone()).into()
     }
 
@@ -525,7 +594,7 @@ impl RbExpr {
         map_single(self, lambda, output_type, agg_list)
     }
 
-    pub fn dot(&self, other: &RbExpr) -> Self {
+    pub fn dot(&self, other: &Self) -> Self {
         self.inner.clone().dot(other.inner.clone()).into()
     }
 
@@ -562,7 +631,7 @@ impl RbExpr {
         self.inner.clone().upper_bound().into()
     }
 
-    pub fn cumulative_eval(&self, expr: &RbExpr, min_periods: usize, parallel: bool) -> Self {
+    pub fn cumulative_eval(&self, expr: &Self, min_periods: usize, parallel: bool) -> Self {
         self.inner
             .clone()
             .cumulative_eval(expr.inner.clone(), min_periods, parallel)
@@ -744,4 +813,22 @@ impl RbExpr {
         };
         self.inner.clone().set_sorted_flag(is_sorted).into()
     }
+
+    pub fn replace(
+        &self,
+        old: &Self,
+        new: &Self,
+        default: Option<&Self>,
+        return_dtype: Option<Wrap<DataType>>,
+    ) -> Self {
+        self.inner
+            .clone()
+            .replace(
+                old.inner.clone(),
+                new.inner.clone(),
+                default.map(|e| e.inner.clone()),
+                return_dtype.map(|dt| dt.0),
+            )
+            .into()
+    }
 }

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

@@ -7,6 +7,14 @@ use crate::conversion::Wrap;
 use crate::{RbExpr, RbResult};
 
 impl RbExpr {
+    pub fn list_all(&self) -> Self {
+        self.inner.clone().list().all().into()
+    }
+
+    pub fn list_any(&self) -> Self {
+        self.inner.clone().list().any().into()
+    }
+
     pub fn list_arg_max(&self) -> Self {
         self.inner.clone().list().arg_max().into()
     }
@@ -47,11 +55,11 @@ impl RbExpr {
         self.inner.clone().list().get(index.inner.clone()).into()
     }
 
-    pub fn list_join(&self, separator: &RbExpr) -> Self {
+    pub fn list_join(&self, separator: &RbExpr, ignore_nulls: bool) -> Self {
         self.inner
             .clone()
             .list()
-            .join(separator.inner.clone())
+            .join(separator.inner.clone(), ignore_nulls)
             .into()
     }
 
@@ -100,6 +108,10 @@ impl RbExpr {
             .into()
     }
 
+    pub fn list_tail(&self, n: &RbExpr) -> Self {
+        self.inner.clone().list().tail(n.inner.clone()).into()
+    }
+
     pub fn list_sort(&self, reverse: bool) -> Self {
         self.inner
             .clone()
@@ -116,14 +128,50 @@ impl RbExpr {
         self.inner.clone().list().sum().with_fmt("list.sum").into()
     }
 
-    pub fn list_take(&self, index: &RbExpr, null_on_oob: bool) -> Self {
+    pub fn list_drop_nulls(&self) -> Self {
+        self.inner.clone().list().drop_nulls().into()
+    }
+
+    pub fn list_sample_n(
+        &self,
+        n: &RbExpr,
+        with_replacement: bool,
+        shuffle: bool,
+        seed: Option<u64>,
+    ) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .sample_n(n.inner.clone(), with_replacement, shuffle, seed)
+            .into()
+    }
+
+    pub fn list_sample_fraction(
+        &self,
+        fraction: &RbExpr,
+        with_replacement: bool,
+        shuffle: bool,
+        seed: Option<u64>,
+    ) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .sample_fraction(fraction.inner.clone(), with_replacement, shuffle, seed)
+            .into()
+    }
+
+    pub fn list_gather(&self, index: &RbExpr, null_on_oob: bool) -> Self {
         self.inner
             .clone()
             .list()
-            .take(index.inner.clone(), null_on_oob)
+            .gather(index.inner.clone(), null_on_oob)
             .into()
     }
 
+    pub fn list_to_array(&self, width: usize) -> Self {
+        self.inner.clone().list().to_array(width).into()
+    }
+
     pub fn list_to_struct(
         &self,
         width_strat: Wrap<ListToStructWidthStrategy>,

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

@@ -46,7 +46,55 @@ impl RbExpr {
         self.inner.clone().meta().has_multiple_outputs()
     }
 
+    pub fn meta_is_column(&self) -> bool {
+        self.inner.clone().meta().is_column()
+    }
+
     pub fn meta_is_regex_projection(&self) -> bool {
         self.inner.clone().meta().is_regex_projection()
     }
+
+    pub fn _meta_selector_add(&self, other: &RbExpr) -> RbResult<RbExpr> {
+        let out = self
+            .inner
+            .clone()
+            .meta()
+            ._selector_add(other.inner.clone())
+            .map_err(RbPolarsErr::from)?;
+        Ok(out.into())
+    }
+
+    pub fn _meta_selector_sub(&self, other: &RbExpr) -> RbResult<RbExpr> {
+        let out = self
+            .inner
+            .clone()
+            .meta()
+            ._selector_sub(other.inner.clone())
+            .map_err(RbPolarsErr::from)?;
+        Ok(out.into())
+    }
+
+    pub fn _meta_selector_and(&self, other: &RbExpr) -> RbResult<RbExpr> {
+        let out = self
+            .inner
+            .clone()
+            .meta()
+            ._selector_and(other.inner.clone())
+            .map_err(RbPolarsErr::from)?;
+        Ok(out.into())
+    }
+
+    pub fn _meta_as_selector(&self) -> RbExpr {
+        self.inner.clone().meta()._into_selector().into()
+    }
+
+    pub fn meta_tree_format(&self) -> RbResult<String> {
+        let e = self
+            .inner
+            .clone()
+            .meta()
+            .into_tree_formatter()
+            .map_err(RbPolarsErr::from)?;
+        Ok(format!("{e}"))
+    }
 }

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

@@ -100,6 +100,7 @@ impl RbExpr {
         by: Option<String>,
         closed: Option<Wrap<ClosedWindow>>,
         ddof: u8,
+        warn_if_unsorted: bool,
     ) -> Self {
         let options = RollingOptions {
             window_size: Duration::parse(&window_size),
@@ -109,6 +110,7 @@ impl RbExpr {
             by,
             closed_window: closed.map(|c| c.0),
             fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+            warn_if_unsorted,
         };
 
         self.inner.clone().rolling_std(options).into()
@@ -124,6 +126,7 @@ impl RbExpr {
         by: Option<String>,
         closed: Option<Wrap<ClosedWindow>>,
         ddof: u8,
+        warn_if_unsorted: bool,
     ) -> Self {
         let options = RollingOptions {
             window_size: Duration::parse(&window_size),
@@ -133,11 +136,13 @@ impl RbExpr {
             by,
             closed_window: closed.map(|c| c.0),
             fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+            warn_if_unsorted,
         };
 
         self.inner.clone().rolling_var(options).into()
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn rolling_median(
         &self,
         window_size: String,
@@ -146,6 +151,7 @@ impl RbExpr {
         center: bool,
         by: Option<String>,
         closed: Option<Wrap<ClosedWindow>>,
+        warn_if_unsorted: bool,
     ) -> Self {
         let options = RollingOptions {
             window_size: Duration::parse(&window_size),
@@ -154,12 +160,10 @@ impl RbExpr {
             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>),
+            fn_params: None,
+            warn_if_unsorted,
         };
-        self.inner.clone().rolling_quantile(options).into()
+        self.inner.clone().rolling_median(options).into()
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -173,6 +177,7 @@ impl RbExpr {
         center: bool,
         by: Option<String>,
         closed: Option<Wrap<ClosedWindow>>,
+        warn_if_unsorted: bool,
     ) -> Self {
         let options = RollingOptions {
             window_size: Duration::parse(&window_size),
@@ -181,13 +186,14 @@ impl RbExpr {
             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>),
+            fn_params: None,
+            warn_if_unsorted,
         };
 
-        self.inner.clone().rolling_quantile(options).into()
+        self.inner
+            .clone()
+            .rolling_quantile(interpolation.0, quantile, options)
+            .into()
     }
 
     pub fn rolling_skew(&self, window_size: usize, bias: bool) -> Self {