polars-df 0.11.0 → 0.12.0

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

@@ -15,19 +15,14 @@ use crate::{RbDataFrame, RbExpr, RbLazyGroupBy, RbPolarsErr, RbResult, RbValueEr
 #[magnus::wrap(class = "Polars::RbLazyFrame")]
 #[derive(Clone)]
 pub struct RbLazyFrame {
-    pub ldf: LazyFrame,
-}
-
-impl RbLazyFrame {
-    fn get_schema(&self) -> RbResult<SchemaRef> {
-        let schema = self.ldf.schema().map_err(RbPolarsErr::from)?;
-        Ok(schema)
-    }
+    pub ldf: RefCell<LazyFrame>,
 }
 
 impl From<LazyFrame> for RbLazyFrame {
     fn from(ldf: LazyFrame) -> Self {
-        RbLazyFrame { ldf }
+        RbLazyFrame {
+            ldf: RefCell::new(ldf),
+        }
     }
 }
 
@@ -69,7 +64,7 @@ impl RbLazyFrame {
         });
 
         let lf = LazyJsonLineReader::new(path)
-            .with_infer_schema_length(infer_schema_length)
+            .with_infer_schema_length(infer_schema_length.and_then(NonZeroUsize::new))
             .with_batch_size(batch_size)
             .with_n_rows(n_rows)
             .low_memory(low_memory)
@@ -163,8 +158,9 @@ impl RbLazyFrame {
         row_index: Option<(String, IdxSize)>,
         low_memory: bool,
         use_statistics: bool,
-        hive_partitioning: bool,
+        hive_partitioning: Option<bool>,
         hive_schema: Option<Wrap<Schema>>,
+        try_parse_hive_dates: bool,
         glob: bool,
     ) -> RbResult<Self> {
         let parallel = parallel.0;
@@ -184,7 +180,9 @@ impl RbLazyFrame {
         });
         let hive_options = HiveOptions {
             enabled: hive_partitioning,
+            hive_start_idx: 0,
             schema: hive_schema,
+            try_parse_dates: try_parse_hive_dates,
         };
 
         let args = ScanArgsParquet {
@@ -236,13 +234,14 @@ impl RbLazyFrame {
 
     pub fn write_json(&self, rb_f: Value) -> RbResult<()> {
         let file = BufWriter::new(get_file_like(rb_f, true)?);
-        serde_json::to_writer(file, &self.ldf.logical_plan)
+        serde_json::to_writer(file, &self.ldf.borrow().logical_plan)
             .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
         Ok(())
     }
 
     pub fn describe_plan(&self) -> RbResult<String> {
         self.ldf
+            .borrow()
             .describe_plan()
             .map_err(RbPolarsErr::from)
             .map_err(Into::into)
@@ -251,6 +250,7 @@ impl RbLazyFrame {
     pub fn describe_optimized_plan(&self) -> RbResult<String> {
         let result = self
             .ldf
+            .borrow()
             .describe_optimized_plan()
             .map_err(RbPolarsErr::from)?;
         Ok(result)
@@ -269,7 +269,7 @@ impl RbLazyFrame {
         allow_streaming: bool,
         _eager: bool,
     ) -> RbLazyFrame {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let mut ldf = ldf
             .with_type_coercion(type_coercion)
             .with_predicate_pushdown(predicate_pushdown)
@@ -293,12 +293,12 @@ impl RbLazyFrame {
         maintain_order: bool,
         multithreaded: bool,
     ) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.sort(
             [&by_column],
             SortMultipleOptions {
                 descending: vec![descending],
-                nulls_last,
+                nulls_last: vec![nulls_last],
                 multithreaded,
                 maintain_order,
             },
@@ -310,11 +310,11 @@ impl RbLazyFrame {
         &self,
         by: RArray,
         descending: Vec<bool>,
-        nulls_last: bool,
+        nulls_last: Vec<bool>,
         maintain_order: bool,
         multithreaded: bool,
     ) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let exprs = rb_exprs_to_exprs(by)?;
         Ok(ldf
             .sort_by_exprs(
@@ -330,12 +330,12 @@ impl RbLazyFrame {
     }
 
     pub fn cache(&self) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.cache().into()
     }
 
     pub fn collect(&self) -> RbResult<RbDataFrame> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let df = ldf.collect().map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }
@@ -346,7 +346,7 @@ impl RbLazyFrame {
         path: PathBuf,
         compression: String,
         compression_level: Option<i32>,
-        statistics: bool,
+        statistics: Wrap<StatisticsOptions>,
         row_group_size: Option<usize>,
         data_pagesize_limit: Option<usize>,
         maintain_order: bool,
@@ -355,13 +355,13 @@ impl RbLazyFrame {
 
         let options = ParquetWriteOptions {
             compression,
-            statistics,
+            statistics: statistics.0,
             row_group_size,
             data_pagesize_limit,
             maintain_order,
         };
 
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.sink_parquet(path, options).map_err(RbPolarsErr::from)?;
         Ok(())
     }
@@ -377,7 +377,7 @@ impl RbLazyFrame {
             maintain_order,
         };
 
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.sink_ipc(path, options).map_err(RbPolarsErr::from)?;
         Ok(())
     }
@@ -395,6 +395,7 @@ impl RbLazyFrame {
         datetime_format: Option<String>,
         date_format: Option<String>,
         time_format: Option<String>,
+        float_scientific: Option<bool>,
         float_precision: Option<usize>,
         null_value: Option<String>,
         quote_style: Option<Wrap<QuoteStyle>>,
@@ -407,6 +408,7 @@ impl RbLazyFrame {
             date_format,
             time_format,
             datetime_format,
+            float_scientific,
             float_precision,
             separator,
             quote_char,
@@ -423,7 +425,7 @@ impl RbLazyFrame {
             serialize_options,
         };
 
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.sink_csv(path, options).map_err(RbPolarsErr::from)?;
         Ok(())
     }
@@ -431,36 +433,36 @@ impl RbLazyFrame {
     pub fn sink_json(&self, path: PathBuf, maintain_order: bool) -> RbResult<()> {
         let options = JsonWriterOptions { maintain_order };
 
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.sink_json(path, options).map_err(RbPolarsErr::from)?;
         Ok(())
     }
 
     pub fn fetch(&self, n_rows: usize) -> RbResult<RbDataFrame> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let df = ldf.fetch(n_rows).map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }
 
     pub fn filter(&self, predicate: &RbExpr) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.filter(predicate.inner.clone()).into()
     }
 
     pub fn select(&self, exprs: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let exprs = rb_exprs_to_exprs(exprs)?;
         Ok(ldf.select(exprs).into())
     }
 
     pub fn select_seq(&self, exprs: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let exprs = rb_exprs_to_exprs(exprs)?;
         Ok(ldf.select_seq(exprs).into())
     }
 
     pub fn group_by(&self, by: RArray, maintain_order: bool) -> RbResult<RbLazyGroupBy> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let by = rb_exprs_to_exprs(by)?;
         let lazy_gb = if maintain_order {
             ldf.group_by_stable(by)
@@ -479,10 +481,9 @@ impl RbLazyFrame {
         offset: String,
         closed: Wrap<ClosedWindow>,
         by: RArray,
-        check_sorted: bool,
     ) -> RbResult<RbLazyGroupBy> {
         let closed_window = closed.0;
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let by = rb_exprs_to_exprs(by)?;
         let lazy_gb = ldf.rolling(
             index_column.inner.clone(),
@@ -492,7 +493,6 @@ impl RbLazyFrame {
                 period: Duration::parse(&period),
                 offset: Duration::parse(&offset),
                 closed_window,
-                check_sorted,
             },
         );
 
@@ -513,11 +513,10 @@ impl RbLazyFrame {
         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 ldf = self.ldf.borrow().clone();
         let lazy_gb = ldf.group_by_dynamic(
             index_column.inner.clone(),
             by,
@@ -529,7 +528,6 @@ impl RbLazyFrame {
                 include_boundaries,
                 closed_window,
                 start_by: start_by.0,
-                check_sorted,
                 ..Default::default()
             },
         );
@@ -541,14 +539,14 @@ impl RbLazyFrame {
 
     pub fn with_context(&self, contexts: RArray) -> RbResult<Self> {
         let contexts = contexts
-            .each()
-            .map(|v| TryConvert::try_convert(v.unwrap()))
+            .into_iter()
+            .map(TryConvert::try_convert)
             .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
         let contexts = contexts
             .into_iter()
-            .map(|ldf| ldf.ldf.clone())
+            .map(|ldf| ldf.ldf.borrow().clone())
             .collect::<Vec<_>>();
-        Ok(self.ldf.clone().with_context(contexts).into())
+        Ok(self.ldf.borrow().clone().with_context(contexts).into())
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -566,8 +564,8 @@ impl RbLazyFrame {
         tolerance: Option<Wrap<AnyValue<'_>>>,
         tolerance_str: Option<String>,
     ) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let other = other.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
+        let other = other.ldf.borrow().clone();
         let left_on = left_on.inner.clone();
         let right_on = right_on.inner.clone();
         Ok(ldf
@@ -601,8 +599,8 @@ impl RbLazyFrame {
         how: Wrap<JoinType>,
         suffix: String,
     ) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let other = other.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
+        let other = other.ldf.borrow().clone();
         let left_on = rb_exprs_to_exprs(left_on)?;
         let right_on = rb_exprs_to_exprs(right_on)?;
 
@@ -621,32 +619,32 @@ impl RbLazyFrame {
     }
 
     pub fn with_column(&self, expr: &RbExpr) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.with_column(expr.inner.clone()).into()
     }
 
     pub fn with_columns(&self, exprs: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         Ok(ldf.with_columns(rb_exprs_to_exprs(exprs)?).into())
     }
 
     pub fn with_columns_seq(&self, exprs: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         Ok(ldf.with_columns_seq(rb_exprs_to_exprs(exprs)?).into())
     }
 
     pub fn rename(&self, existing: Vec<String>, new: Vec<String>) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.rename(existing, new).into()
     }
 
     pub fn reverse(&self) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.reverse().into()
     }
 
     pub fn shift(&self, n: &RbExpr, fill_value: Option<&RbExpr>) -> Self {
-        let lf = self.ldf.clone();
+        let lf = self.ldf.borrow().clone();
         let out = match fill_value {
             Some(v) => lf.shift_and_fill(n.inner.clone(), v.inner.clone()),
             None => lf.shift(n.inner.clone()),
@@ -655,48 +653,48 @@ impl RbLazyFrame {
     }
 
     pub fn fill_nan(&self, fill_value: &RbExpr) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.fill_nan(fill_value.inner.clone()).into()
     }
 
     pub fn min(&self) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let out = ldf.min();
         out.into()
     }
 
     pub fn max(&self) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let out = ldf.max();
         out.into()
     }
 
     pub fn sum(&self) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let out = ldf.sum();
         out.into()
     }
 
     pub fn mean(&self) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let out = ldf.mean();
         out.into()
     }
 
     pub fn std(&self, ddof: u8) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let out = ldf.std(ddof);
         out.into()
     }
 
     pub fn var(&self, ddof: u8) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let out = ldf.var(ddof);
         out.into()
     }
 
     pub fn median(&self) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let out = ldf.median();
         out.into()
     }
@@ -706,19 +704,19 @@ impl RbLazyFrame {
         quantile: &RbExpr,
         interpolation: Wrap<QuantileInterpolOptions>,
     ) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let out = ldf.quantile(quantile.inner.clone(), interpolation.0);
         out.into()
     }
 
     pub fn explode(&self, column: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         let column = rb_exprs_to_exprs(column)?;
         Ok(ldf.explode(column).into())
     }
 
     pub fn null_count(&self) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.null_count().into()
     }
 
@@ -728,7 +726,7 @@ impl RbLazyFrame {
         subset: Option<Vec<String>>,
         keep: Wrap<UniqueKeepStrategy>,
     ) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         Ok(match maintain_order {
             true => ldf.unique_stable(subset, keep.0),
             false => ldf.unique(subset, keep.0),
@@ -737,75 +735,63 @@ impl RbLazyFrame {
     }
 
     pub fn drop_nulls(&self, subset: Option<Vec<String>>) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.drop_nulls(subset.map(|v| v.into_iter().map(|s| col(&s)).collect()))
             .into()
     }
 
     pub fn slice(&self, offset: i64, len: Option<IdxSize>) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.slice(offset, len.unwrap_or(IdxSize::MAX)).into()
     }
 
     pub fn tail(&self, n: IdxSize) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.tail(n).into()
     }
 
-    pub fn melt(
+    pub fn unpivot(
         &self,
-        id_vars: Vec<String>,
-        value_vars: Vec<String>,
+        on: Vec<String>,
+        index: Vec<String>,
         value_name: Option<String>,
         variable_name: Option<String>,
         streamable: bool,
     ) -> Self {
-        let args = MeltArgs {
-            id_vars: strings_to_smartstrings(id_vars),
-            value_vars: strings_to_smartstrings(value_vars),
+        let args = UnpivotArgs {
+            on: strings_to_smartstrings(on),
+            index: strings_to_smartstrings(index),
             value_name: value_name.map(|s| s.into()),
             variable_name: variable_name.map(|s| s.into()),
             streamable,
         };
 
-        let ldf = self.ldf.clone();
-        ldf.melt(args).into()
+        let ldf = self.ldf.borrow().clone();
+        ldf.unpivot(args).into()
     }
 
     pub fn with_row_index(&self, name: String, offset: Option<IdxSize>) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.with_row_index(&name, offset).into()
     }
 
     pub fn drop(&self, cols: Vec<String>) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.drop(cols).into()
     }
 
     pub fn cast_all(&self, dtype: Wrap<DataType>, strict: bool) -> Self {
-        self.ldf.clone().cast_all(dtype.0, strict).into()
+        self.ldf.borrow().clone().cast_all(dtype.0, strict).into()
     }
 
     pub fn clone(&self) -> Self {
-        self.ldf.clone().into()
-    }
-
-    pub fn columns(&self) -> RbResult<RArray> {
-        let schema = self.get_schema()?;
-        let iter = schema.iter_names().map(|s| s.as_str());
-        Ok(RArray::from_iter(iter))
+        self.ldf.borrow().clone().into()
     }
 
-    pub fn dtypes(&self) -> RbResult<RArray> {
-        let schema = self.get_schema()?;
-        let iter = schema.iter_dtypes().map(|dt| Wrap(dt.clone()).into_value());
-        Ok(RArray::from_iter(iter))
-    }
+    pub fn collect_schema(&self) -> RbResult<RHash> {
+        let schema = self.ldf.borrow_mut().schema().map_err(RbPolarsErr::from)?;
 
-    pub fn schema(&self) -> RbResult<RHash> {
-        let schema = self.get_schema()?;
         let schema_dict = RHash::new();
-
         schema.iter_fields().for_each(|fld| {
             // TODO remove unwrap
             schema_dict
@@ -819,23 +805,20 @@ impl RbLazyFrame {
     }
 
     pub fn unnest(&self, cols: Vec<String>) -> Self {
-        self.ldf.clone().unnest(cols).into()
-    }
-
-    pub fn width(&self) -> RbResult<usize> {
-        Ok(self.get_schema()?.len())
+        self.ldf.borrow().clone().unnest(cols).into()
     }
 
     pub fn count(&self) -> Self {
-        let ldf = self.ldf.clone();
+        let ldf = self.ldf.borrow().clone();
         ldf.count().into()
     }
 
     pub fn merge_sorted(&self, other: &Self, key: String) -> RbResult<Self> {
         let out = self
             .ldf
+            .borrow()
             .clone()
-            .merge_sorted(other.ldf.clone(), &key)
+            .merge_sorted(other.ldf.borrow().clone(), &key)
             .map_err(RbPolarsErr::from)?;
         Ok(out.into())
     }