polars-df 0.21.0 → 0.21.1

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

@@ -323,6 +323,10 @@ impl RbExpr {
         self.inner.clone().arg_min().into()
     }
 
+    pub fn index_of(&self, element: &Self) -> Self {
+        self.inner.clone().index_of(element.inner.clone()).into()
+    }
+
     pub fn search_sorted(
         &self,
         element: &Self,
@@ -429,6 +433,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()
     }
@@ -482,6 +493,10 @@ impl RbExpr {
         self.inner.clone().round(decimals, mode.0).into()
     }
 
+    pub fn round_sig_figs(&self, digits: i32) -> Self {
+        self.clone().inner.round_sig_figs(digits).into()
+    }
+
     pub fn floor(&self) -> Self {
         self.inner.clone().floor().into()
     }
@@ -517,6 +532,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()
     }
@@ -553,6 +572,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()
     }
@@ -566,18 +593,35 @@ impl RbExpr {
         Ok(self.inner.clone().over(partition_by).into())
     }
 
-    pub fn _and(&self, expr: &Self) -> Self {
-        self.inner.clone().and(expr.inner.clone()).into()
+    pub fn rolling(
+        &self,
+        index_column: String,
+        period: String,
+        offset: String,
+        closed: Wrap<ClosedWindow>,
+    ) -> RbResult<Self> {
+        let options = RollingGroupOptions {
+            index_column: index_column.into(),
+            period: Duration::try_parse(&period).map_err(RbPolarsErr::from)?,
+            offset: Duration::try_parse(&offset).map_err(RbPolarsErr::from)?,
+            closed_window: closed.0,
+        };
+
+        Ok(self.inner.clone().rolling(options).into())
     }
 
-    pub fn _xor(&self, expr: &Self) -> Self {
-        self.inner.clone().xor(expr.inner.clone()).into()
+    pub fn and_(&self, expr: &Self) -> Self {
+        self.inner.clone().and(expr.inner.clone()).into()
     }
 
-    pub fn _or(&self, expr: &Self) -> Self {
+    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, nulls_equal: bool) -> Self {
         self.inner
             .clone()
@@ -593,6 +637,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()
     }
@@ -756,6 +808,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,
@@ -811,6 +872,10 @@ impl RbExpr {
         self.inner.clone().log(base).into()
     }
 
+    pub fn log1p(&self) -> Self {
+        self.inner.clone().log1p().into()
+    }
+
     pub fn exp(&self) -> Self {
         self.inner.clone().exp().into()
     }
@@ -857,6 +922,20 @@ 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

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

@@ -83,6 +83,18 @@ impl RbExpr {
         self.inner.clone().list().mean().into()
     }
 
+    pub fn list_median(&self) -> Self {
+        self.inner.clone().list().median().into()
+    }
+
+    pub fn list_std(&self, ddof: u8) -> Self {
+        self.inner.clone().list().std(ddof).into()
+    }
+
+    pub fn list_var(&self, ddof: u8) -> Self {
+        self.inner.clone().list().var(ddof).into()
+    }
+
     pub fn list_min(&self) -> Self {
         self.inner.clone().list().min().into()
     }
@@ -171,6 +183,14 @@ impl RbExpr {
             .into()
     }
 
+    pub fn list_gather_every(&self, n: &RbExpr, offset: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .gather_every(n.inner.clone(), offset.inner.clone())
+            .into()
+    }
+
     pub fn list_to_array(&self, width: usize) -> Self {
         self.inner.clone().list().to_array(width).into()
     }
@@ -205,6 +225,10 @@ impl RbExpr {
             .into())
     }
 
+    pub fn list_n_unique(&self) -> Self {
+        self.inner.clone().list().n_unique().into()
+    }
+
     pub fn list_unique(&self, maintain_order: bool) -> Self {
         let e = self.inner.clone();
 

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

@@ -57,6 +57,17 @@ impl RbExpr {
         self.inner.clone().meta().is_regex_projection()
     }
 
+    pub fn meta_is_column_selection(&self, allow_aliasing: bool) -> bool {
+        self.inner
+            .clone()
+            .meta()
+            .is_column_selection(allow_aliasing)
+    }
+
+    pub fn meta_is_literal(&self, allow_aliasing: bool) -> bool {
+        self.inner.clone().meta().is_literal(allow_aliasing)
+    }
+
     fn compute_tree_format(
         &self,
         display_as_dot: bool,

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

@@ -1,5 +1,6 @@
 mod array;
 mod binary;
+mod bitwise;
 mod categorical;
 pub mod datatype;
 mod datetime;

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

@@ -42,4 +42,12 @@ impl RbExpr {
     pub fn name_to_uppercase(&self) -> Self {
         self.inner.clone().name().to_uppercase().into()
     }
+
+    pub fn name_prefix_fields(&self, prefix: String) -> Self {
+        self.inner.clone().name().prefix_fields(&prefix).into()
+    }
+
+    pub fn name_suffix_fields(&self, suffix: String) -> Self {
+        self.inner.clone().name().suffix_fields(&suffix).into()
+    }
 }

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

@@ -337,4 +337,24 @@ impl RbExpr {
 
         self.inner.clone().rolling_skew(options).into()
     }
+
+    pub fn rolling_kurtosis(
+        &self,
+        window_size: usize,
+        fisher: bool,
+        bias: bool,
+        min_periods: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            window_size,
+            weights: None,
+            min_periods,
+            center,
+            fn_params: Some(RollingFnParams::Kurtosis { fisher, bias }),
+        };
+
+        self.inner.clone().rolling_kurtosis(options).into()
+    }
 }

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

@@ -122,6 +122,14 @@ impl RbExpr {
             .into()
     }
 
+    pub fn str_head(&self, n: &Self) -> Self {
+        self.inner.clone().str().head(n.inner.clone()).into()
+    }
+
+    pub fn str_tail(&self, n: &Self) -> Self {
+        self.inner.clone().str().tail(n.inner.clone()).into()
+    }
+
     pub fn str_to_uppercase(&self) -> Self {
         self.inner.clone().str().to_uppercase().into()
     }
@@ -159,6 +167,10 @@ impl RbExpr {
             .into()
     }
 
+    pub fn str_normalize(&self, form: Wrap<UnicodeForm>) -> Self {
+        self.inner.clone().str().normalize(form.0).into()
+    }
+
     pub fn str_reverse(&self) -> Self {
         self.inner.clone().str().reverse().into()
     }
@@ -200,6 +212,23 @@ impl RbExpr {
         }
     }
 
+    pub fn str_find(&self, pat: &Self, literal: Option<bool>, strict: bool) -> Self {
+        match literal {
+            Some(true) => self
+                .inner
+                .clone()
+                .str()
+                .find_literal(pat.inner.clone())
+                .into(),
+            _ => self
+                .inner
+                .clone()
+                .str()
+                .find(pat.inner.clone(), strict)
+                .into(),
+        }
+    }
+
     pub fn str_ends_with(&self, sub: &RbExpr) -> Self {
         self.inner.clone().str().ends_with(sub.inner.clone()).into()
     }
@@ -351,4 +380,34 @@ impl RbExpr {
             )
             .into()
     }
+
+    pub fn str_extract_many(
+        &self,
+        patterns: &RbExpr,
+        ascii_case_insensitive: bool,
+        overlapping: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .extract_many(patterns.inner.clone(), ascii_case_insensitive, overlapping)
+            .into()
+    }
+
+    pub fn str_find_many(
+        &self,
+        patterns: &RbExpr,
+        ascii_case_insensitive: bool,
+        overlapping: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .find_many(patterns.inner.clone(), ascii_case_insensitive, overlapping)
+            .into()
+    }
+
+    pub fn str_escape_regex(&self) -> Self {
+        self.inner.clone().str().escape_regex().into()
+    }
 }

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

@@ -1,4 +1,6 @@
-use crate::RbExpr;
+use magnus::RArray;
+
+use crate::{RbExpr, RbResult, rb_exprs_to_exprs};
 
 impl RbExpr {
     pub fn struct_field_by_index(&self, index: i64) -> Self {
@@ -16,4 +18,10 @@ impl RbExpr {
     pub fn struct_json_encode(&self) -> Self {
         self.inner.clone().struct_().json_encode().into()
     }
+
+    pub fn struct_with_fields(&self, fields: RArray) -> RbResult<Self> {
+        let fields = rb_exprs_to_exprs(fields)?;
+        let e = self.inner.clone().struct_().with_fields(fields);
+        Ok(e.into())
+    }
 }

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

@@ -21,6 +21,25 @@ pub fn read_ipc_schema(rb_f: Value) -> RbResult<RHash> {
     Ok(dict)
 }
 
+pub fn read_parquet_metadata(rb_f: Value) -> RbResult<RHash> {
+    use polars_parquet::read::read_metadata;
+    use polars_parquet::read::schema::read_custom_key_value_metadata;
+
+    let metadata = match get_either_file(rb_f, false)? {
+        EitherRustRubyFile::Rust(r) => {
+            read_metadata(&mut BufReader::new(r)).map_err(RbPolarsErr::from)?
+        }
+        EitherRustRubyFile::Rb(mut r) => read_metadata(&mut r).map_err(RbPolarsErr::from)?,
+    };
+
+    let key_value_metadata = read_custom_key_value_metadata(metadata.key_value_metadata());
+    let dict = RHash::new();
+    for (key, value) in key_value_metadata.into_iter() {
+        dict.aset(key.as_str(), value.as_str())?;
+    }
+    Ok(dict)
+}
+
 pub fn read_parquet_schema(rb_f: Value) -> RbResult<RHash> {
     use polars_parquet::read::{infer_schema, read_metadata};
 

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

@@ -85,6 +85,10 @@ pub fn as_struct(exprs: RArray) -> RbResult<RbExpr> {
     Ok(dsl::as_struct(exprs).into())
 }
 
+pub fn field(names: Vec<String>) -> RbExpr {
+    dsl::Expr::Field(names.into_iter().map(|x| x.into()).collect()).into()
+}
+
 pub fn coalesce(exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
     Ok(dsl::coalesce(&exprs).into())

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

@@ -325,6 +325,30 @@ impl RbLazyFrame {
             .into())
     }
 
+    pub fn top_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
+        let ldf = self.ldf.borrow().clone();
+        let exprs = rb_exprs_to_exprs(by)?;
+        Ok(ldf
+            .top_k(
+                k,
+                exprs,
+                SortMultipleOptions::new().with_order_descending_multi(reverse),
+            )
+            .into())
+    }
+
+    pub fn bottom_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
+        let ldf = self.ldf.borrow().clone();
+        let exprs = rb_exprs_to_exprs(by)?;
+        Ok(ldf
+            .bottom_k(
+                k,
+                exprs,
+                SortMultipleOptions::new().with_order_descending_multi(reverse),
+            )
+            .into())
+    }
+
     pub fn cache(&self) -> Self {
         let ldf = self.ldf.borrow().clone();
         ldf.cache().into()
@@ -509,6 +533,11 @@ impl RbLazyFrame {
         ldf.filter(predicate.inner.clone()).into()
     }
 
+    pub fn remove(&self, predicate: &RbExpr) -> Self {
+        let ldf = self.ldf.borrow().clone();
+        ldf.remove(predicate.inner.clone()).into()
+    }
+
     pub fn select(&self, exprs: RArray) -> RbResult<Self> {
         let ldf = self.ldf.borrow().clone();
         let exprs = rb_exprs_to_exprs(exprs)?;
@@ -701,6 +730,20 @@ impl RbLazyFrame {
             .into())
     }
 
+    pub fn join_where(&self, other: &Self, predicates: RArray, suffix: String) -> RbResult<Self> {
+        let ldf = self.ldf.borrow().clone();
+        let other = other.ldf.borrow().clone();
+
+        let predicates = rb_exprs_to_exprs(predicates)?;
+
+        Ok(ldf
+            .join_builder()
+            .with(other)
+            .suffix(suffix)
+            .join_where(predicates)
+            .into())
+    }
+
     pub fn with_column(&self, expr: &RbExpr) -> Self {
         let ldf = self.ldf.borrow().clone();
         ldf.with_column(expr.inner.clone()).into()
@@ -816,6 +859,14 @@ impl RbLazyFrame {
         .into())
     }
 
+    pub fn drop_nans(&self, subset: Option<&RbSelector>) -> Self {
+        self.ldf
+            .borrow()
+            .clone()
+            .drop_nans(subset.map(|e| e.inner.clone()))
+            .into()
+    }
+
     pub fn drop_nulls(&self, subset: Option<&RbSelector>) -> Self {
         self.ldf
             .borrow()