polars-df 0.8.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,11 +165,7 @@ 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)
@@ -302,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()
     }
 
@@ -335,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()
     }
 
@@ -356,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()
     }
 
@@ -368,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()
     }
 
@@ -423,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)
@@ -532,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()
     }
 
@@ -584,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()
     }
 
@@ -621,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)
@@ -803,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

@@ -142,6 +142,7 @@ impl RbExpr {
         self.inner.clone().rolling_var(options).into()
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn rolling_median(
         &self,
         window_size: String,

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

@@ -1,7 +1,7 @@
 use polars::prelude::*;
 
 use crate::conversion::Wrap;
-use crate::RbExpr;
+use crate::{RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
     pub fn str_concat(&self, delimiter: String, ignore_nulls: bool) -> Self {
@@ -107,8 +107,12 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_slice(&self, start: i64, length: Option<u64>) -> Self {
-        self.inner.clone().str().slice(start, length).into()
+    pub fn str_slice(&self, start: &Self, length: &Self) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .slice(start.inner.clone(), length.inner.clone())
+            .into()
     }
 
     pub fn str_explode(&self) -> Self {
@@ -147,6 +151,10 @@ impl RbExpr {
             .into()
     }
 
+    pub fn str_reverse(&self) -> Self {
+        self.inner.clone().str().reverse().into()
+    }
+
     pub fn str_pad_start(&self, length: usize, fillchar: char) -> Self {
         self.clone().inner.str().pad_start(length, fillchar).into()
     }
@@ -155,8 +163,8 @@ impl RbExpr {
         self.clone().inner.str().pad_end(length, fillchar).into()
     }
 
-    pub fn str_zfill(&self, alignment: usize) -> Self {
-        self.clone().inner.str().zfill(alignment).into()
+    pub fn str_zfill(&self, length: &Self) -> Self {
+        self.clone().inner.str().zfill(length.inner.clone()).into()
     }
 
     pub fn str_contains(&self, pat: &RbExpr, literal: Option<bool>, strict: bool) -> Self {
@@ -245,7 +253,7 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_json_extract(
+    pub fn str_json_decode(
         &self,
         dtype: Option<Wrap<DataType>>,
         infer_schema_len: Option<usize>,
@@ -287,8 +295,12 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_extract(&self, pat: String, group_index: usize) -> Self {
-        self.inner.clone().str().extract(&pat, group_index).into()
+    pub fn str_extract(&self, pat: &Self, group_index: usize) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .extract(pat.inner.clone(), group_index)
+            .into()
     }
 
     pub fn str_extract_all(&self, pat: &RbExpr) -> Self {
@@ -299,6 +311,16 @@ impl RbExpr {
             .into()
     }
 
+    pub fn str_extract_groups(&self, pat: String) -> RbResult<Self> {
+        Ok(self
+            .inner
+            .clone()
+            .str()
+            .extract_groups(&pat)
+            .map_err(RbPolarsErr::from)?
+            .into())
+    }
+
     pub fn str_count_matches(&self, pat: &Self, literal: bool) -> Self {
         self.inner
             .clone()
@@ -338,4 +360,33 @@ impl RbExpr {
     pub fn str_splitn(&self, by: &Self, n: usize) -> Self {
         self.inner.clone().str().splitn(by.inner.clone(), n).into()
     }
+
+    pub fn str_to_decimal(&self, infer_len: usize) -> Self {
+        self.inner.clone().str().to_decimal(infer_len).into()
+    }
+
+    pub fn str_contains_any(&self, patterns: &RbExpr, ascii_case_insensitive: bool) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .contains_any(patterns.inner.clone(), ascii_case_insensitive)
+            .into()
+    }
+
+    pub fn str_replace_many(
+        &self,
+        patterns: &RbExpr,
+        replace_with: &RbExpr,
+        ascii_case_insensitive: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .replace_many(
+                patterns.inner.clone(),
+                replace_with.inner.clone(),
+                ascii_case_insensitive,
+            )
+            .into()
+    }
 }

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

@@ -1,15 +1,19 @@
 use crate::RbExpr;
 
 impl RbExpr {
-    pub fn struct_field_by_name(&self, name: String) -> Self {
-        self.inner.clone().struct_().field_by_name(&name).into()
-    }
-
     pub fn struct_field_by_index(&self, index: i64) -> Self {
         self.inner.clone().struct_().field_by_index(index).into()
     }
 
+    pub fn struct_field_by_name(&self, name: String) -> Self {
+        self.inner.clone().struct_().field_by_name(&name).into()
+    }
+
     pub fn struct_rename_fields(&self, names: Vec<String>) -> Self {
         self.inner.clone().struct_().rename_fields(names).into()
     }
+
+    pub fn struct_json_encode(&self) -> Self {
+        self.inner.clone().struct_().json_encode().into()
+    }
 }

data/ext/polars/src/functions/aggregation.rs

@@ -33,3 +33,9 @@ pub fn sum_horizontal(exprs: RArray) -> RbResult<RbExpr> {
     let e = dsl::sum_horizontal(exprs).map_err(RbPolarsErr::from)?;
     Ok(e.into())
 }
+
+pub fn mean_horizontal(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let e = dsl::mean_horizontal(exprs).map_err(RbPolarsErr::from)?;
+    Ok(e.into())
+}