polars-df 0.7.0 → 0.9.0

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 {
@@ -192,7 +200,7 @@ impl RbExpr {
         self.clone()
             .inner
             .map(
-                move |s| s.utf8().map(|s| Some(s.hex_encode().into_series())),
+                move |s| s.str().map(|s| Some(s.hex_encode().into_series())),
                 GetOutput::same_type(),
             )
             .with_fmt("str.hex_encode")
@@ -203,7 +211,7 @@ impl RbExpr {
         self.clone()
             .inner
             .map(
-                move |s| s.utf8()?.hex_decode(strict).map(|s| Some(s.into_series())),
+                move |s| s.str()?.hex_decode(strict).map(|s| Some(s.into_series())),
                 GetOutput::same_type(),
             )
             .with_fmt("str.hex_decode")
@@ -214,7 +222,7 @@ impl RbExpr {
         self.clone()
             .inner
             .map(
-                move |s| s.utf8().map(|s| Some(s.base64_encode().into_series())),
+                move |s| s.str().map(|s| Some(s.base64_encode().into_series())),
                 GetOutput::same_type(),
             )
             .with_fmt("str.base64_encode")
@@ -226,7 +234,7 @@ impl RbExpr {
             .inner
             .map(
                 move |s| {
-                    s.utf8()?
+                    s.str()?
                         .base64_decode(strict)
                         .map(|s| Some(s.into_series()))
                 },
@@ -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>,
@@ -258,8 +266,8 @@ impl RbExpr {
         };
 
         let function = move |s: Series| {
-            let ca = s.utf8()?;
-            match ca.json_extract(dtype.clone(), infer_schema_len) {
+            let ca = s.str()?;
+            match ca.json_decode(dtype.clone(), infer_schema_len) {
                 Ok(ca) => Ok(Some(ca.into_series())),
                 Err(e) => Err(PolarsError::ComputeError(format!("{e:?}").into())),
             }
@@ -268,13 +276,13 @@ impl RbExpr {
         self.clone()
             .inner
             .map(function, output_type)
-            .with_fmt("str.json_extract")
+            .with_fmt("str.json_decode")
             .into()
     }
 
     pub fn str_json_path_match(&self, pat: String) -> Self {
         let function = move |s: Series| {
-            let ca = s.utf8()?;
+            let ca = s.str()?;
             match ca.json_path_match(&pat) {
                 Ok(ca) => Ok(Some(ca.into_series())),
                 Err(e) => Err(PolarsError::ComputeError(format!("{:?}", e).into())),
@@ -282,13 +290,17 @@ impl RbExpr {
         };
         self.clone()
             .inner
-            .map(function, GetOutput::from_type(DataType::Utf8))
+            .map(function, GetOutput::from_type(DataType::String))
             .with_fmt("str.json_path_match")
             .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())
+}

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

@@ -17,6 +17,44 @@ macro_rules! set_unwrapped_or_0 {
     };
 }
 
+pub fn rolling_corr(
+    x: &RbExpr,
+    y: &RbExpr,
+    window_size: IdxSize,
+    min_periods: IdxSize,
+    ddof: u8,
+) -> RbExpr {
+    dsl::rolling_corr(
+        x.inner.clone(),
+        y.inner.clone(),
+        RollingCovOptions {
+            min_periods,
+            window_size,
+            ddof,
+        },
+    )
+    .into()
+}
+
+pub fn rolling_cov(
+    x: &RbExpr,
+    y: &RbExpr,
+    window_size: IdxSize,
+    min_periods: IdxSize,
+    ddof: u8,
+) -> RbExpr {
+    dsl::rolling_cov(
+        x.inner.clone(),
+        y.inner.clone(),
+        RollingCovOptions {
+            min_periods,
+            window_size,
+            ddof,
+        },
+    )
+    .into()
+}
+
 pub fn arg_sort_by(by: RArray, descending: Vec<bool>) -> RbResult<RbExpr> {
     let by = rb_exprs_to_exprs(by)?;
     Ok(dsl::arg_sort_by(by, &descending).into())
@@ -83,6 +121,47 @@ pub fn concat_lf(
     Ok(lf.into())
 }
 
+pub fn concat_list(s: RArray) -> RbResult<RbExpr> {
+    let s = rb_exprs_to_exprs(s)?;
+    let expr = dsl::concat_list(s).map_err(RbPolarsErr::from)?;
+    Ok(expr.into())
+}
+
+pub fn concat_str(s: RArray, separator: String, ignore_nulls: bool) -> RbResult<RbExpr> {
+    let s = rb_exprs_to_exprs(s)?;
+    Ok(dsl::concat_str(s, &separator, ignore_nulls).into())
+}
+
+pub fn len() -> RbExpr {
+    dsl::len().into()
+}
+
+pub fn cov(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
+    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone(), ddof).into()
+}
+
+pub fn arctan2(y: &RbExpr, x: &RbExpr) -> RbExpr {
+    y.inner.clone().arctan2(x.inner.clone()).into()
+}
+
+pub fn arctan2d(y: &RbExpr, x: &RbExpr) -> RbExpr {
+    y.inner.clone().arctan2(x.inner.clone()).degrees().into()
+}
+
+pub fn cum_fold(
+    acc: &RbExpr,
+    lambda: Value,
+    exprs: RArray,
+    include_init: bool,
+) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let lambda = Opaque::from(lambda);
+
+    let func =
+        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
+    Ok(dsl::cum_fold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+}
+
 pub fn concat_lf_diagonal(
     lfs: RArray,
     rechunk: bool,
@@ -110,6 +189,19 @@ pub fn concat_lf_diagonal(
     Ok(lf.into())
 }
 
+pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
+    dsl::dtype_cols(dtypes).into()
+}
+
+pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
+    let dtypes = dtypes
+        .each()
+        .map(|v| Wrap::<DataType>::try_convert(v?))
+        .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
+    let dtypes = vec_extract_wrapped(dtypes);
+    Ok(crate::functions::lazy::dtype_cols(dtypes))
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn duration(
     weeks: Option<&RbExpr>,
@@ -146,38 +238,21 @@ pub fn duration(
     dsl::duration(args).into()
 }
 
-pub fn count() -> RbExpr {
-    dsl::count().into()
-}
-
 pub fn first() -> RbExpr {
     dsl::first().into()
 }
 
-pub fn last() -> RbExpr {
-    dsl::last().into()
-}
-
-pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
-    dsl::dtype_cols(dtypes).into()
-}
-
 pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
     let lambda = Opaque::from(lambda);
 
     let func =
         move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
-    Ok(polars::lazy::dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
+    Ok(dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
 }
 
-pub fn cumfold(acc: &RbExpr, lambda: Value, exprs: RArray, include_init: bool) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
-    let lambda = Opaque::from(lambda);
-
-    let func =
-        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
-    Ok(polars::lazy::dsl::cum_fold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+pub fn last() -> RbExpr {
+    dsl::last().into()
 }
 
 pub fn lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
@@ -219,6 +294,10 @@ pub fn lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
     }
 }
 
+pub fn pearson_corr(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
+    dsl::pearson_corr(a.inner.clone(), b.inner.clone(), ddof).into()
+}
+
 pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbResult<RbExpr> {
     let mut value = value.inner.clone();
     let n = n.inner.clone();
@@ -228,7 +307,7 @@ pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbRe
     }
 
     if let Expr::Literal(lv) = &value {
-        let av = lv.to_anyvalue().unwrap();
+        let av = lv.to_any_value().unwrap();
         // Integer inputs that fit in Int32 are parsed as such
         if let DataType::Int64 = av.dtype() {
             let int_value = av.try_extract::<i64>().unwrap();
@@ -240,35 +319,11 @@ pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbRe
     Ok(dsl::repeat(value, n).into())
 }
 
-pub fn pearson_corr(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
-    polars::lazy::dsl::pearson_corr(a.inner.clone(), b.inner.clone(), ddof).into()
-}
-
 pub fn spearman_rank_corr(a: &RbExpr, b: &RbExpr, ddof: u8, propagate_nans: bool) -> RbExpr {
-    polars::lazy::dsl::spearman_rank_corr(a.inner.clone(), b.inner.clone(), ddof, propagate_nans)
-        .into()
-}
-
-pub fn cov(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
-    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone(), ddof).into()
-}
-
-pub fn concat_str(s: RArray, sep: String) -> RbResult<RbExpr> {
-    let s = rb_exprs_to_exprs(s)?;
-    Ok(dsl::concat_str(s, &sep).into())
+    dsl::spearman_rank_corr(a.inner.clone(), b.inner.clone(), ddof, propagate_nans).into()
 }
 
-pub fn concat_lst(s: RArray) -> RbResult<RbExpr> {
-    let s = rb_exprs_to_exprs(s)?;
-    let expr = dsl::concat_list(s).map_err(RbPolarsErr::from)?;
+pub fn sql_expr(sql: String) -> RbResult<RbExpr> {
+    let expr = polars::sql::sql_expr(&sql).map_err(RbPolarsErr::from)?;
     Ok(expr.into())
 }
-
-pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
-    let dtypes = dtypes
-        .each()
-        .map(|v| Wrap::<DataType>::try_convert(v?))
-        .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
-    let dtypes = vec_extract_wrapped(dtypes);
-    Ok(crate::functions::lazy::dtype_cols(dtypes))
-}

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

@@ -7,7 +7,7 @@ use polars_core::POOL;
 use crate::conversion::Wrap;
 use crate::{RbResult, RbValueError};
 
-pub fn get_idx_type() -> Value {
+pub fn get_index_type() -> Value {
     Wrap(IDX_DTYPE).into_value()
 }
 
@@ -36,3 +36,47 @@ pub fn get_float_fmt() -> RbResult<String> {
     };
     Ok(strfmt.to_string())
 }
+
+pub fn set_float_precision(precision: Option<usize>) -> RbResult<()> {
+    use polars_core::fmt::set_float_precision;
+    set_float_precision(precision);
+    Ok(())
+}
+
+pub fn get_float_precision() -> RbResult<Option<usize>> {
+    use polars_core::fmt::get_float_precision;
+    Ok(get_float_precision())
+}
+
+pub fn set_thousands_separator(sep: Option<char>) -> RbResult<()> {
+    use polars_core::fmt::set_thousands_separator;
+    set_thousands_separator(sep);
+    Ok(())
+}
+
+pub fn get_thousands_separator() -> RbResult<Option<String>> {
+    use polars_core::fmt::get_thousands_separator;
+    Ok(Some(get_thousands_separator()))
+}
+
+pub fn set_decimal_separator(sep: Option<char>) -> RbResult<()> {
+    use polars_core::fmt::set_decimal_separator;
+    set_decimal_separator(sep);
+    Ok(())
+}
+
+pub fn get_decimal_separator() -> RbResult<Option<char>> {
+    use polars_core::fmt::get_decimal_separator;
+    Ok(Some(get_decimal_separator()))
+}
+
+pub fn set_trim_decimal_zeros(trim: Option<bool>) -> RbResult<()> {
+    use polars_core::fmt::set_trim_decimal_zeros;
+    set_trim_decimal_zeros(trim);
+    Ok(())
+}
+
+pub fn get_trim_decimal_zeros() -> RbResult<Option<bool>> {
+    use polars_core::fmt::get_trim_decimal_zeros;
+    Ok(Some(get_trim_decimal_zeros()))
+}

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

@@ -6,21 +6,16 @@ use crate::prelude::*;
 use crate::RbExpr;
 
 pub fn int_range(start: &RbExpr, end: &RbExpr, step: i64, dtype: Wrap<DataType>) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
     let dtype = dtype.0;
-
-    let mut result = dsl::int_range(start.inner.clone(), end.inner.clone(), step);
-
-    if dtype != DataType::Int64 {
-        result = result.cast(dtype)
-    }
-
-    result.into()
+    dsl::int_range(start, end, step, dtype).into()
 }
 
-pub fn int_ranges(start: &RbExpr, end: &RbExpr, step: i64, dtype: Wrap<DataType>) -> RbExpr {
+pub fn int_ranges(start: &RbExpr, end: &RbExpr, step: &RbExpr, dtype: Wrap<DataType>) -> RbExpr {
     let dtype = dtype.0;
 
-    let mut result = dsl::int_ranges(start.inner.clone(), end.inner.clone(), step);
+    let mut result = dsl::int_ranges(start.inner.clone(), end.inner.clone(), step.inner.clone());
 
     if dtype != DataType::Int64 {
         result = result.cast(DataType::List(Box::new(dtype)))

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

@@ -1,3 +1,7 @@
+use crate::RbResult;
+use magnus::{RArray, Ruby, Value};
+use polars_core::StringCacheHolder;
+
 pub fn enable_string_cache() {
     polars_core::enable_string_cache()
 }
@@ -9,3 +13,13 @@ pub fn disable_string_cache() {
 pub fn using_string_cache() -> bool {
     polars_core::using_string_cache()
 }
+
+#[magnus::wrap(class = "Polars::RbStringCacheHolder")]
+pub struct RbStringCacheHolder {}
+
+impl RbStringCacheHolder {
+    pub fn hold() -> RbResult<Value> {
+        let _hold = StringCacheHolder::hold();
+        Ruby::get().unwrap().yield_splat(RArray::new())
+    }
+}