polars-df 0.5.0 → 0.7.0

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,13 +227,20 @@ 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()
     }
 
-    pub fn sort(&self, by_column: String, reverse: bool, nulls_last: bool) -> Self {
+    pub fn sort(
+        &self,
+        by_column: String,
+        reverse: bool,
+        nulls_last: bool,
+        maintain_order: bool,
+    ) -> Self {
         let ldf = self.ldf.clone();
         ldf.sort(
             &by_column,
@@ -238,6 +248,7 @@ impl RbLazyFrame {
                 descending: reverse,
                 nulls_last,
                 multithreaded: true,
+                maintain_order,
             },
         )
         .into()
@@ -248,10 +259,13 @@ impl RbLazyFrame {
         by_column: RArray,
         reverse: Vec<bool>,
         nulls_last: bool,
+        maintain_order: bool,
     ) -> RbResult<Self> {
         let ldf = self.ldf.clone();
         let exprs = rb_exprs_to_exprs(by_column)?;
-        Ok(ldf.sort_by_exprs(exprs, reverse, nulls_last).into())
+        Ok(ldf
+            .sort_by_exprs(exprs, reverse, nulls_last, maintain_order)
+            .into())
     }
 
     pub fn cache(&self) -> Self {
@@ -308,31 +322,32 @@ 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,
         offset: String,
         closed: Wrap<ClosedWindow>,
         by: RArray,
+        check_sorted: bool,
     ) -> RbResult<RbLazyGroupBy> {
         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 {
@@ -340,6 +355,7 @@ impl RbLazyFrame {
                 period: Duration::parse(&period),
                 offset: Duration::parse(&offset),
                 closed_window,
+                check_sorted,
             },
         );
 
@@ -349,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()
             },
         );
@@ -387,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()
@@ -478,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 {