polars-df 0.10.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),
+        }
     }
 }
 
@@ -48,7 +43,7 @@ impl RbLazyFrame {
         // in this scope
         let json = unsafe { std::mem::transmute::<&'_ str, &'static str>(json.as_str()) };
 
-        let lp = serde_json::from_str::<LogicalPlan>(json)
+        let lp = serde_json::from_str::<DslPlan>(json)
             .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
         Ok(LazyFrame::from(lp).into())
     }
@@ -63,10 +58,13 @@ impl RbLazyFrame {
         row_index: Option<(String, IdxSize)>,
     ) -> RbResult<Self> {
         let batch_size = batch_size.map(|v| v.0);
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
 
         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)
@@ -107,7 +105,10 @@ impl RbLazyFrame {
         let quote_char = quote_char.map(|s| s.as_bytes()[0]);
         let separator = separator.as_bytes()[0];
         let eol_char = eol_char.as_bytes()[0];
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
 
         let overwrite_dtype = overwrite_dtype.map(|overwrite_dtype| {
             overwrite_dtype
@@ -119,17 +120,17 @@ impl RbLazyFrame {
         let r = LazyCsvReader::new(path)
             .with_infer_schema_length(infer_schema_length)
             .with_separator(separator)
-            .has_header(has_header)
+            .with_has_header(has_header)
             .with_ignore_errors(ignore_errors)
             .with_skip_rows(skip_rows)
             .with_n_rows(n_rows)
             .with_cache(cache)
-            .with_dtype_overwrite(overwrite_dtype.as_ref())
+            .with_dtype_overwrite(overwrite_dtype.map(Arc::new))
             // TODO add with_schema
-            .low_memory(low_memory)
+            .with_low_memory(low_memory)
             .with_comment_prefix(comment_prefix.as_deref())
             .with_quote_char(quote_char)
-            .with_end_of_line_char(eol_char)
+            .with_eol_char(eol_char)
             .with_rechunk(rechunk)
             .with_skip_rows_after_header(skip_rows_after_header)
             .with_encoding(encoding.0)
@@ -137,7 +138,7 @@ impl RbLazyFrame {
             .with_try_parse_dates(try_parse_dates)
             .with_null_values(null_values)
             // TODO add with_missing_is_null
-            .truncate_ragged_lines(truncate_ragged_lines);
+            .with_truncate_ragged_lines(truncate_ragged_lines);
 
         if let Some(_lambda) = with_schema_modify {
             todo!();
@@ -157,8 +158,10 @@ 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;
         let hive_schema = hive_schema.map(|s| Arc::new(s.0));
@@ -171,10 +174,15 @@ impl RbLazyFrame {
                 .ok_or_else(|| RbValueError::new_err("expected a path argument".to_string()))?
         };
 
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
         let hive_options = HiveOptions {
             enabled: hive_partitioning,
+            hive_start_idx: 0,
             schema: hive_schema,
+            try_parse_dates: try_parse_hive_dates,
         };
 
         let args = ScanArgsParquet {
@@ -187,6 +195,7 @@ impl RbLazyFrame {
             cloud_options: None,
             use_statistics,
             hive_options,
+            glob,
         };
 
         let lf = if path.is_some() {
@@ -206,7 +215,11 @@ impl RbLazyFrame {
         row_index: Option<(String, IdxSize)>,
         memory_map: bool,
     ) -> RbResult<Self> {
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
+
         let args = ScanArgsIpc {
             n_rows,
             cache,
@@ -221,18 +234,23 @@ 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) -> String {
-        self.ldf.describe_plan()
+    pub fn describe_plan(&self) -> RbResult<String> {
+        self.ldf
+            .borrow()
+            .describe_plan()
+            .map_err(RbPolarsErr::from)
+            .map_err(Into::into)
     }
 
     pub fn describe_optimized_plan(&self) -> RbResult<String> {
         let result = self
             .ldf
+            .borrow()
             .describe_optimized_plan()
             .map_err(RbPolarsErr::from)?;
         Ok(result)
@@ -251,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)
@@ -275,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,
             },
@@ -292,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(
@@ -312,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())
     }
@@ -328,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,
@@ -337,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(())
     }
@@ -359,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(())
     }
@@ -377,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>>,
@@ -389,6 +408,7 @@ impl RbLazyFrame {
             date_format,
             time_format,
             datetime_format,
+            float_scientific,
             float_precision,
             separator,
             quote_char,
@@ -405,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(())
     }
@@ -413,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)
@@ -461,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(),
@@ -474,7 +493,6 @@ impl RbLazyFrame {
                 period: Duration::parse(&period),
                 offset: Duration::parse(&offset),
                 closed_window,
-                check_sorted,
             },
         );
 
@@ -495,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,
@@ -511,7 +528,6 @@ impl RbLazyFrame {
                 include_boundaries,
                 closed_window,
                 start_by: start_by.0,
-                check_sorted,
                 ..Default::default()
             },
         );
@@ -523,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)]
@@ -548,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
@@ -583,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)?;
 
@@ -603,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()),
@@ -637,72 +653,70 @@ 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) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let out = ldf.min().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+    pub fn min(&self) -> Self {
+        let ldf = self.ldf.borrow().clone();
+        let out = ldf.min();
+        out.into()
     }
 
-    pub fn max(&self) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let out = ldf.max().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+    pub fn max(&self) -> Self {
+        let ldf = self.ldf.borrow().clone();
+        let out = ldf.max();
+        out.into()
     }
 
-    pub fn sum(&self) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let out = ldf.sum().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+    pub fn sum(&self) -> Self {
+        let ldf = self.ldf.borrow().clone();
+        let out = ldf.sum();
+        out.into()
     }
 
-    pub fn mean(&self) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let out = ldf.mean().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+    pub fn mean(&self) -> Self {
+        let ldf = self.ldf.borrow().clone();
+        let out = ldf.mean();
+        out.into()
     }
 
-    pub fn std(&self, ddof: u8) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let out = ldf.std(ddof).map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+    pub fn std(&self, ddof: u8) -> Self {
+        let ldf = self.ldf.borrow().clone();
+        let out = ldf.std(ddof);
+        out.into()
     }
 
-    pub fn var(&self, ddof: u8) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let out = ldf.var(ddof).map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+    pub fn var(&self, ddof: u8) -> Self {
+        let ldf = self.ldf.borrow().clone();
+        let out = ldf.var(ddof);
+        out.into()
     }
 
-    pub fn median(&self) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let out = ldf.median().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+    pub fn median(&self) -> Self {
+        let ldf = self.ldf.borrow().clone();
+        let out = ldf.median();
+        out.into()
     }
 
     pub fn quantile(
         &self,
         quantile: &RbExpr,
         interpolation: Wrap<QuantileInterpolOptions>,
-    ) -> RbResult<Self> {
-        let ldf = self.ldf.clone();
-        let out = ldf
-            .quantile(quantile.inner.clone(), interpolation.0)
-            .map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+    ) -> Self {
+        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()
     }
 
@@ -712,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),
@@ -721,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()
+        self.ldf.borrow().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))
-    }
+    pub fn collect_schema(&self) -> RbResult<RHash> {
+        let schema = self.ldf.borrow_mut().schema().map_err(RbPolarsErr::from)?;
 
-    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 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
@@ -803,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())
     }