polars-df 0.6.0 → 0.7.0

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

@@ -1,10 +1,12 @@
 use magnus::encoding::{self, EncodingCapable};
-use magnus::{class, Float, Integer, RArray, RString, Value};
+use magnus::{
+    class, prelude::*, typed_data::Obj, value::Opaque, Float, Integer, RArray, RString, Ruby, Value,
+};
 use polars::lazy::dsl;
 use polars::prelude::*;
 
-use crate::apply::lazy::binary_lambda;
 use crate::conversion::{get_lf, get_rbseq, Wrap};
+use crate::map::lazy::binary_lambda;
 use crate::prelude::vec_extract_wrapped;
 use crate::rb_exprs_to_exprs;
 use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
@@ -15,10 +17,6 @@ macro_rules! set_unwrapped_or_0 {
     };
 }
 
-pub fn arange(low: &RbExpr, high: &RbExpr, step: i64) -> RbExpr {
-    dsl::arange(low.inner.clone(), high.inner.clone(), step).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())
@@ -30,7 +28,7 @@ pub fn arg_where(condition: &RbExpr) -> RbExpr {
 
 pub fn as_struct(exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
-    Ok(dsl::as_struct(&exprs).into())
+    Ok(dsl::as_struct(exprs).into())
 }
 
 pub fn coalesce(exprs: RArray) -> RbResult<RbExpr> {
@@ -45,7 +43,7 @@ pub fn col(name: String) -> RbExpr {
 pub fn collect_all(lfs: RArray) -> RbResult<RArray> {
     let lfs = lfs
         .each()
-        .map(|v| v?.try_convert::<&RbLazyFrame>())
+        .map(|v| <&RbLazyFrame>::try_convert(v?))
         .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
 
     Ok(RArray::from_iter(lfs.iter().map(|lf| {
@@ -85,36 +83,65 @@ pub fn concat_lf(
     Ok(lf.into())
 }
 
+pub fn concat_lf_diagonal(
+    lfs: RArray,
+    rechunk: bool,
+    parallel: bool,
+    to_supertypes: bool,
+) -> RbResult<RbLazyFrame> {
+    let iter = lfs.each();
+
+    let lfs = iter
+        .map(|item| {
+            let item = item?;
+            get_lf(item)
+        })
+        .collect::<RbResult<Vec<_>>>()?;
+
+    let lf = dsl::functions::concat_lf_diagonal(
+        lfs,
+        UnionArgs {
+            rechunk,
+            parallel,
+            to_supertypes,
+        },
+    )
+    .map_err(RbPolarsErr::from)?;
+    Ok(lf.into())
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn duration(
+    weeks: Option<&RbExpr>,
     days: Option<&RbExpr>,
+    hours: Option<&RbExpr>,
+    minutes: Option<&RbExpr>,
     seconds: Option<&RbExpr>,
-    nanoseconds: Option<&RbExpr>,
-    microseconds: Option<&RbExpr>,
     milliseconds: Option<&RbExpr>,
-    minutes: Option<&RbExpr>,
-    hours: Option<&RbExpr>,
-    weeks: Option<&RbExpr>,
+    microseconds: Option<&RbExpr>,
+    nanoseconds: Option<&RbExpr>,
+    time_unit: Wrap<TimeUnit>,
 ) -> RbExpr {
     set_unwrapped_or_0!(
+        weeks,
         days,
+        hours,
+        minutes,
         seconds,
-        nanoseconds,
-        microseconds,
         milliseconds,
-        minutes,
-        hours,
-        weeks,
+        microseconds,
+        nanoseconds,
     );
     let args = DurationArgs {
+        weeks,
         days,
+        hours,
+        minutes,
         seconds,
-        nanoseconds,
-        microseconds,
         milliseconds,
-        minutes,
-        hours,
-        weeks,
+        microseconds,
+        nanoseconds,
+        time_unit: time_unit.0,
     };
     dsl::duration(args).into()
 }
@@ -137,23 +164,27 @@ pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
 
 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(lambda, a, b);
+    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())
 }
 
 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(lambda, a, b);
-    Ok(polars::lazy::dsl::cumfold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+    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 lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
     if value.is_kind_of(class::true_class()) || value.is_kind_of(class::false_class()) {
-        Ok(dsl::lit(value.try_convert::<bool>()?).into())
+        Ok(dsl::lit(bool::try_convert(value)?).into())
     } else if let Some(v) = Integer::from_value(value) {
-        match v.try_convert::<i64>() {
+        match v.to_i64() {
             Ok(val) => {
                 if val > 0 && val < i32::MAX as i64 || val < 0 && val > i32::MIN as i64 {
                     Ok(dsl::lit(val as i32).into())
@@ -162,19 +193,19 @@ pub fn lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
                 }
             }
             _ => {
-                let val = value.try_convert::<u64>()?;
+                let val = v.to_u64()?;
                 Ok(dsl::lit(val).into())
             }
         }
     } else if let Some(v) = Float::from_value(value) {
-        Ok(dsl::lit(v.try_convert::<f64>()?).into())
+        Ok(dsl::lit(v.to_f64()).into())
     } else if let Some(v) = RString::from_value(value) {
         if v.enc_get() == encoding::Index::utf8() {
-            Ok(dsl::lit(v.try_convert::<String>()?).into())
+            Ok(dsl::lit(v.to_string()?).into())
         } else {
             Ok(dsl::lit(unsafe { v.as_slice() }).into())
         }
-    } else if let Ok(series) = value.try_convert::<&RbSeries>() {
+    } else if let Ok(series) = Obj::<RbSeries>::try_convert(value) {
         Ok(dsl::lit(series.series.borrow().clone()).into())
     } else if value.is_nil() {
         Ok(dsl::lit(Null {}).into())
@@ -218,8 +249,8 @@ pub fn spearman_rank_corr(a: &RbExpr, b: &RbExpr, ddof: u8, propagate_nans: bool
         .into()
 }
 
-pub fn cov(a: &RbExpr, b: &RbExpr) -> RbExpr {
-    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone()).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> {
@@ -236,14 +267,8 @@ pub fn concat_lst(s: RArray) -> RbResult<RbExpr> {
 pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
     let dtypes = dtypes
         .each()
-        .map(|v| v?.try_convert::<Wrap<DataType>>())
+        .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))
 }
-
-// TODO rename to sum_horizontal
-pub fn sum_exprs(exprs: RArray) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
-    Ok(polars::lazy::dsl::sum_horizontal(exprs).into())
-}

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

@@ -1,8 +1,38 @@
 use magnus::{IntoValue, Value};
+use polars_core;
+use polars_core::fmt::FloatFmt;
 use polars_core::prelude::IDX_DTYPE;
+use polars_core::POOL;
 
 use crate::conversion::Wrap;
+use crate::{RbResult, RbValueError};
 
 pub fn get_idx_type() -> Value {
     Wrap(IDX_DTYPE).into_value()
 }
+
+pub fn threadpool_size() -> usize {
+    POOL.current_num_threads()
+}
+
+pub fn set_float_fmt(fmt: String) -> RbResult<()> {
+    let fmt = match fmt.as_str() {
+        "full" => FloatFmt::Full,
+        "mixed" => FloatFmt::Mixed,
+        e => {
+            return Err(RbValueError::new_err(format!(
+                "fmt must be one of {{'full', 'mixed'}}, got {e}",
+            )))
+        }
+    };
+    polars_core::fmt::set_float_fmt(fmt);
+    Ok(())
+}
+
+pub fn get_float_fmt() -> RbResult<String> {
+    let strfmt = match polars_core::fmt::get_float_fmt() {
+        FloatFmt::Full => "full",
+        FloatFmt::Mixed => "mixed",
+    };
+    Ok(strfmt.to_string())
+}

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

@@ -0,0 +1,8 @@
+use crate::conversion::Wrap;
+use crate::prelude::DataType;
+use crate::RbResult;
+
+pub fn dtype_str_repr(dtype: Wrap<DataType>) -> RbResult<String> {
+    let dtype = dtype.0;
+    Ok(dtype.to_string())
+}

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

@@ -1,5 +1,10 @@
+pub mod aggregation;
 pub mod eager;
 pub mod io;
 pub mod lazy;
 pub mod meta;
+pub mod misc;
+pub mod random;
+pub mod range;
+pub mod string_cache;
 pub mod whenthen;

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

@@ -0,0 +1,6 @@
+use crate::RbResult;
+
+pub fn set_random_seed(seed: u64) -> RbResult<()> {
+    polars_core::random::set_global_random_seed(seed);
+    Ok(())
+}

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

@@ -0,0 +1,46 @@
+use polars::lazy::dsl;
+use polars_core::datatypes::{TimeUnit, TimeZone};
+
+use crate::conversion::Wrap;
+use crate::prelude::*;
+use crate::RbExpr;
+
+pub fn int_range(start: &RbExpr, end: &RbExpr, step: i64, dtype: Wrap<DataType>) -> RbExpr {
+    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()
+}
+
+pub fn int_ranges(start: &RbExpr, end: &RbExpr, step: i64, dtype: Wrap<DataType>) -> RbExpr {
+    let dtype = dtype.0;
+
+    let mut result = dsl::int_ranges(start.inner.clone(), end.inner.clone(), step);
+
+    if dtype != DataType::Int64 {
+        result = result.cast(DataType::List(Box::new(dtype)))
+    }
+
+    result.into()
+}
+
+pub fn date_range(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    time_unit: Option<Wrap<TimeUnit>>,
+    time_zone: Option<TimeZone>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    let time_unit = time_unit.map(|x| x.0);
+    dsl::date_range(start, end, every, closed, time_unit, time_zone).into()
+}

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

@@ -0,0 +1,11 @@
+pub fn enable_string_cache() {
+    polars_core::enable_string_cache()
+}
+
+pub fn disable_string_cache() {
+    polars_core::disable_string_cache()
+}
+
+pub fn using_string_cache() -> bool {
+    polars_core::using_string_cache()
+}

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

@@ -16,23 +16,23 @@ impl From<dsl::When> for RbWhen {
 
 #[magnus::wrap(class = "Polars::RbWhenThen")]
 #[derive(Clone)]
-pub struct RbWhenThen {
-    pub inner: dsl::WhenThen,
+pub struct RbThen {
+    pub inner: dsl::Then,
 }
 
-impl From<dsl::WhenThen> for RbWhenThen {
-    fn from(inner: dsl::WhenThen) -> Self {
-        RbWhenThen { inner }
+impl From<dsl::Then> for RbThen {
+    fn from(inner: dsl::Then) -> Self {
+        RbThen { inner }
     }
 }
 
 impl RbWhen {
-    pub fn then(&self, expr: &RbExpr) -> RbWhenThen {
+    pub fn then(&self, expr: &RbExpr) -> RbThen {
         self.inner.clone().then(expr.inner.clone()).into()
     }
 }
 
-impl RbWhenThen {
+impl RbThen {
     pub fn overwise(&self, expr: &RbExpr) -> RbExpr {
         self.inner.clone().otherwise(expr.inner.clone()).into()
     }

data/ext/polars/src/lazyframe.rs

@@ -1,4 +1,4 @@
-use magnus::{IntoValue, RArray, RHash, Value};
+use magnus::{IntoValue, RArray, RHash, TryConvert, Value};
 use polars::io::RowCount;
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
@@ -78,32 +78,32 @@ impl RbLazyFrame {
     pub fn new_from_csv(arguments: &[Value]) -> RbResult<Self> {
         // start arguments
         // this pattern is needed for more than 16
-        let path: String = arguments[0].try_convert()?;
-        let sep: String = arguments[1].try_convert()?;
-        let has_header: bool = arguments[2].try_convert()?;
-        let ignore_errors: bool = arguments[3].try_convert()?;
-        let skip_rows: usize = arguments[4].try_convert()?;
-        let n_rows: Option<usize> = arguments[5].try_convert()?;
-        let cache: bool = arguments[6].try_convert()?;
-        let overwrite_dtype: Option<Vec<(String, Wrap<DataType>)>> = arguments[7].try_convert()?;
-        let low_memory: bool = arguments[8].try_convert()?;
-        let comment_char: Option<String> = arguments[9].try_convert()?;
-        let quote_char: Option<String> = arguments[10].try_convert()?;
-        let null_values: Option<Wrap<NullValues>> = arguments[11].try_convert()?;
-        let infer_schema_length: Option<usize> = arguments[12].try_convert()?;
-        let with_schema_modify: Option<Value> = arguments[13].try_convert()?;
-        let rechunk: bool = arguments[14].try_convert()?;
-        let skip_rows_after_header: usize = arguments[15].try_convert()?;
-        let encoding: Wrap<CsvEncoding> = arguments[16].try_convert()?;
-        let row_count: Option<(String, IdxSize)> = arguments[17].try_convert()?;
-        let try_parse_dates: bool = arguments[18].try_convert()?;
-        let eol_char: String = arguments[19].try_convert()?;
+        let path = String::try_convert(arguments[0])?;
+        let separator = String::try_convert(arguments[1])?;
+        let has_header = bool::try_convert(arguments[2])?;
+        let ignore_errors = bool::try_convert(arguments[3])?;
+        let skip_rows = usize::try_convert(arguments[4])?;
+        let n_rows = Option::<usize>::try_convert(arguments[5])?;
+        let cache = bool::try_convert(arguments[6])?;
+        let overwrite_dtype = Option::<Vec<(String, Wrap<DataType>)>>::try_convert(arguments[7])?;
+        let low_memory = bool::try_convert(arguments[8])?;
+        let comment_char = Option::<String>::try_convert(arguments[9])?;
+        let quote_char = Option::<String>::try_convert(arguments[10])?;
+        let null_values = Option::<Wrap<NullValues>>::try_convert(arguments[11])?;
+        let infer_schema_length = Option::<usize>::try_convert(arguments[12])?;
+        let with_schema_modify = Option::<Value>::try_convert(arguments[13])?;
+        let rechunk = bool::try_convert(arguments[14])?;
+        let skip_rows_after_header = usize::try_convert(arguments[15])?;
+        let encoding = Wrap::<CsvEncoding>::try_convert(arguments[16])?;
+        let row_count = Option::<(String, IdxSize)>::try_convert(arguments[17])?;
+        let try_parse_dates = bool::try_convert(arguments[18])?;
+        let eol_char = String::try_convert(arguments[19])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
         let comment_char = comment_char.map(|s| s.as_bytes()[0]);
         let quote_char = quote_char.map(|s| s.as_bytes()[0]);
-        let delimiter = sep.as_bytes()[0];
+        let separator = separator.as_bytes()[0];
         let eol_char = eol_char.as_bytes()[0];
 
         let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
@@ -116,7 +116,7 @@ impl RbLazyFrame {
         });
         let r = LazyCsvReader::new(path)
             .with_infer_schema_length(infer_schema_length)
-            .with_delimiter(delimiter)
+            .with_separator(separator)
             .has_header(has_header)
             .with_ignore_errors(ignore_errors)
             .with_skip_rows(skip_rows)
@@ -151,6 +151,7 @@ impl RbLazyFrame {
         row_count: Option<(String, IdxSize)>,
         low_memory: bool,
         use_statistics: bool,
+        hive_partitioning: bool,
     ) -> RbResult<Self> {
         let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
         let args = ScanArgsParquet {
@@ -163,6 +164,7 @@ impl RbLazyFrame {
             // TODO support cloud options
             cloud_options: None,
             use_statistics,
+            hive_partitioning,
         };
         let lf = LazyFrame::scan_parquet(path, args).map_err(RbPolarsErr::from)?;
         Ok(lf.into())
@@ -217,6 +219,7 @@ impl RbLazyFrame {
         slice_pushdown: bool,
         cse: bool,
         allow_streaming: bool,
+        _eager: bool,
     ) -> RbLazyFrame {
         let ldf = self.ldf.clone();
         let ldf = ldf
@@ -224,8 +227,9 @@ impl RbLazyFrame {
             .with_predicate_pushdown(predicate_pushdown)
             .with_simplify_expr(simplify_expr)
             .with_slice_pushdown(slice_pushdown)
-            .with_common_subplan_elimination(cse)
+            .with_comm_subplan_elim(cse)
             .with_streaming(allow_streaming)
+            ._with_eager(_eager)
             .with_projection_pushdown(projection_pushdown);
         ldf.into()
     }
@@ -318,20 +322,20 @@ impl RbLazyFrame {
         Ok(ldf.select(exprs).into())
     }
 
-    pub fn groupby(&self, by: RArray, maintain_order: bool) -> RbResult<RbLazyGroupBy> {
+    pub fn group_by(&self, by: RArray, maintain_order: bool) -> RbResult<RbLazyGroupBy> {
         let ldf = self.ldf.clone();
         let by = rb_exprs_to_exprs(by)?;
         let lazy_gb = if maintain_order {
-            ldf.groupby_stable(by)
+            ldf.group_by_stable(by)
         } else {
-            ldf.groupby(by)
+            ldf.group_by(by)
         };
         Ok(RbLazyGroupBy {
             lgb: RefCell::new(Some(lazy_gb)),
         })
     }
 
-    pub fn groupby_rolling(
+    pub fn group_by_rolling(
         &self,
         index_column: &RbExpr,
         period: String,
@@ -343,7 +347,7 @@ impl RbLazyFrame {
         let closed_window = closed.0;
         let ldf = self.ldf.clone();
         let by = rb_exprs_to_exprs(by)?;
-        let lazy_gb = ldf.groupby_rolling(
+        let lazy_gb = ldf.group_by_rolling(
             index_column.inner.clone(),
             by,
             RollingGroupOptions {
@@ -361,32 +365,34 @@ impl RbLazyFrame {
     }
 
     #[allow(clippy::too_many_arguments)]
-    pub fn groupby_dynamic(
+    pub fn group_by_dynamic(
         &self,
         index_column: &RbExpr,
         every: String,
         period: String,
         offset: String,
-        truncate: bool,
+        label: Wrap<Label>,
         include_boundaries: bool,
         closed: Wrap<ClosedWindow>,
         by: RArray,
         start_by: Wrap<StartBy>,
+        check_sorted: bool,
     ) -> RbResult<RbLazyGroupBy> {
         let closed_window = closed.0;
         let by = rb_exprs_to_exprs(by)?;
         let ldf = self.ldf.clone();
-        let lazy_gb = ldf.groupby_dynamic(
+        let lazy_gb = ldf.group_by_dynamic(
             index_column.inner.clone(),
             by,
             DynamicGroupOptions {
                 every: Duration::parse(&every),
                 period: Duration::parse(&period),
                 offset: Duration::parse(&offset),
-                truncate,
+                label: label.0,
                 include_boundaries,
                 closed_window,
                 start_by: start_by.0,
+                check_sorted,
                 ..Default::default()
             },
         );
@@ -399,7 +405,7 @@ impl RbLazyFrame {
     pub fn with_context(&self, contexts: RArray) -> RbResult<Self> {
         let contexts = contexts
             .each()
-            .map(|v| v.unwrap().try_convert())
+            .map(|v| TryConvert::try_convert(v.unwrap()))
             .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
         let contexts = contexts
             .into_iter()
@@ -490,14 +496,13 @@ impl RbLazyFrame {
         ldf.reverse().into()
     }
 
-    pub fn shift(&self, periods: i64) -> Self {
-        let ldf = self.ldf.clone();
-        ldf.shift(periods).into()
-    }
-
-    pub fn shift_and_fill(&self, periods: i64, fill_value: &RbExpr) -> Self {
-        let ldf = self.ldf.clone();
-        ldf.shift_and_fill(periods, fill_value.inner.clone()).into()
+    pub fn shift(&self, n: &RbExpr, fill_value: Option<&RbExpr>) -> Self {
+        let lf = self.ldf.clone();
+        let out = match fill_value {
+            Some(v) => lf.shift_and_fill(n.inner.clone(), v.inner.clone()),
+            None => lf.shift(n.inner.clone()),
+        };
+        out.into()
     }
 
     pub fn fill_nan(&self, fill_value: &RbExpr) -> Self {