polars-df 0.23.0 → 0.24.0

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

@@ -1,19 +1,21 @@
-use magnus::{IntoValue, RArray, RHash, Ruby, TryConvert, Value, r_hash::ForEach, typed_data::Obj};
+use magnus::{
+    IntoValue, RArray, RHash, Ruby, TryConvert, Value, r_hash::ForEach,
+    try_convert::TryConvertOwned,
+};
 use polars::io::RowIndex;
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
 use polars_plan::dsl::ScanSources;
-use std::cell::RefCell;
-use std::io::BufWriter;
+use polars_plan::plans::{HintIR, Sorted};
 use std::num::NonZeroUsize;
 
-use super::SinkTarget;
+use super::{RbLazyFrame, RbOptFlags, SinkTarget};
 use crate::conversion::*;
-use crate::expr::rb_exprs_to_exprs;
+use crate::expr::ToExprs;
 use crate::expr::selector::RbSelector;
-use crate::file::get_file_like;
 use crate::io::RbScanOptions;
-use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbLazyGroupBy, RbPolarsErr, RbResult, RbValueError};
+use crate::utils::EnterPolarsExt;
+use crate::{RbDataFrame, RbExpr, RbLazyGroupBy, RbPolarsErr, RbResult, RbValueError};
 
 fn rbobject_to_first_path_and_scan_sources(obj: Value) -> RbResult<(Option<PlPath>, ScanSources)> {
     use crate::file::{RubyScanSourceInput, get_ruby_scan_source_input};
@@ -22,35 +24,58 @@ fn rbobject_to_first_path_and_scan_sources(obj: Value) -> RbResult<(Option<PlPat
             Some(path.clone()),
             ScanSources::Paths(FromIterator::from_iter([path])),
         ),
-        RubyScanSourceInput::File(file) => (None, ScanSources::Files([file].into())),
+        RubyScanSourceInput::File(file) => (None, ScanSources::Files([file.into()].into())),
         RubyScanSourceInput::Buffer(buff) => (None, ScanSources::Buffers([buff].into())),
     })
 }
 
 impl RbLazyFrame {
-    #[allow(clippy::too_many_arguments)]
-    pub fn new_from_ndjson(
-        source: Option<Value>,
-        sources: Wrap<ScanSources>,
-        infer_schema_length: Option<usize>,
-        batch_size: Option<NonZeroUsize>,
-        n_rows: Option<usize>,
-        low_memory: bool,
-        rechunk: bool,
-        row_index: Option<(String, IdxSize)>,
-    ) -> RbResult<Self> {
+    pub fn new_from_ndjson(arguments: &[Value]) -> RbResult<Self> {
+        let source = Option::<Value>::try_convert(arguments[0])?;
+        let sources = Wrap::<ScanSources>::try_convert(arguments[1])?;
+        let infer_schema_length = Option::<usize>::try_convert(arguments[2])?;
+        let schema = Option::<Wrap<Schema>>::try_convert(arguments[3])?;
+        let schema_overrides = Option::<Wrap<Schema>>::try_convert(arguments[4])?;
+        let batch_size = Option::<NonZeroUsize>::try_convert(arguments[5])?;
+        let n_rows = Option::<usize>::try_convert(arguments[6])?;
+        let low_memory = bool::try_convert(arguments[7])?;
+        let rechunk = bool::try_convert(arguments[8])?;
+        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[9])?;
+        let ignore_errors = bool::try_convert(arguments[10])?;
+        let include_file_paths = Option::<String>::try_convert(arguments[11])?;
+        let cloud_options = Option::<Vec<(String, String)>>::try_convert(arguments[12])?;
+        let credential_provider = Option::<Value>::try_convert(arguments[13])?;
+        let retries = usize::try_convert(arguments[14])?;
+        let file_cache_ttl = Option::<u64>::try_convert(arguments[15])?;
+
         let row_index = row_index.map(|(name, offset)| RowIndex {
             name: name.into(),
             offset,
         });
 
         let sources = sources.0;
-        let (_first_path, sources) = match source {
+        let (first_path, sources) = match source {
             None => (sources.first_path().map(|p| p.into_owned()), sources),
             Some(source) => rbobject_to_first_path_and_scan_sources(source)?,
         };
 
-        let r = LazyJsonLineReader::new_with_sources(sources);
+        let mut r = LazyJsonLineReader::new_with_sources(sources);
+
+        if let Some(first_path) = first_path {
+            let first_path_url = first_path.to_str();
+
+            let mut cloud_options =
+                parse_cloud_options(first_path_url, cloud_options.unwrap_or_default())?;
+            cloud_options = cloud_options
+                .with_max_retries(retries)
+                .with_credential_provider(credential_provider.map(|_| todo!()));
+
+            if let Some(file_cache_ttl) = file_cache_ttl {
+                cloud_options.file_cache_ttl = file_cache_ttl;
+            }
+
+            r = r.with_cloud_options(Some(cloud_options));
+        };
 
         let lf = r
             .with_infer_schema_length(infer_schema_length.and_then(NonZeroUsize::new))
@@ -58,11 +83,11 @@ impl RbLazyFrame {
             .with_n_rows(n_rows)
             .low_memory(low_memory)
             .with_rechunk(rechunk)
-            // .with_schema(schema.map(|schema| Arc::new(schema.0)))
-            // .with_schema_overwrite(schema_overrides.map(|x| Arc::new(x.0)))
+            .with_schema(schema.map(|schema| Arc::new(schema.0)))
+            .with_schema_overwrite(schema_overrides.map(|x| Arc::new(x.0)))
             .with_row_index(row_index)
-            // .with_ignore_errors(ignore_errors)
-            // .with_include_file_paths(include_file_paths.map(|x| x.into()))
+            .with_ignore_errors(ignore_errors)
+            .with_include_file_paths(include_file_paths.map(|x| x.into()))
             .finish()
             .map_err(RbPolarsErr::from)?;
 
@@ -73,27 +98,38 @@ impl RbLazyFrame {
         // start arguments
         // this pattern is needed for more than 16
         let source = Option::<Value>::try_convert(arguments[0])?;
-        let sources = Wrap::<ScanSources>::try_convert(arguments[21])?;
-        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_prefix = 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_index = Option::<(String, IdxSize)>::try_convert(arguments[17])?;
-        let try_parse_dates = bool::try_convert(arguments[18])?;
-        let eol_char = String::try_convert(arguments[19])?;
-        let truncate_ragged_lines = bool::try_convert(arguments[20])?;
+        let sources = Wrap::<ScanSources>::try_convert(arguments[1])?;
+        let separator = String::try_convert(arguments[2])?;
+        let has_header = bool::try_convert(arguments[3])?;
+        let ignore_errors = bool::try_convert(arguments[4])?;
+        let skip_rows = usize::try_convert(arguments[5])?;
+        let skip_lines = usize::try_convert(arguments[6])?;
+        let n_rows = Option::<usize>::try_convert(arguments[7])?;
+        let cache = bool::try_convert(arguments[8])?;
+        let overwrite_dtype = Option::<Vec<(String, Wrap<DataType>)>>::try_convert(arguments[9])?;
+        let low_memory = bool::try_convert(arguments[10])?;
+        let comment_prefix = Option::<String>::try_convert(arguments[11])?;
+        let quote_char = Option::<String>::try_convert(arguments[12])?;
+        let null_values = Option::<Wrap<NullValues>>::try_convert(arguments[13])?;
+        let missing_utf8_is_empty_string = bool::try_convert(arguments[14])?;
+        let infer_schema_length = Option::<usize>::try_convert(arguments[15])?;
+        let with_schema_modify = Option::<Value>::try_convert(arguments[16])?;
+        let rechunk = bool::try_convert(arguments[17])?;
+        let skip_rows_after_header = usize::try_convert(arguments[18])?;
+        let encoding = Wrap::<CsvEncoding>::try_convert(arguments[19])?;
+        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[20])?;
+        let try_parse_dates = bool::try_convert(arguments[21])?;
+        let eol_char = String::try_convert(arguments[22])?;
+        let raise_if_empty = bool::try_convert(arguments[23])?;
+        let truncate_ragged_lines = bool::try_convert(arguments[24])?;
+        let decimal_comma = bool::try_convert(arguments[25])?;
+        let glob = bool::try_convert(arguments[26])?;
+        let schema = Option::<Wrap<Schema>>::try_convert(arguments[27])?;
+        let cloud_options = Option::<Vec<(String, String)>>::try_convert(arguments[28])?;
+        let _credential_provider = Option::<Value>::try_convert(arguments[29])?;
+        let retries = usize::try_convert(arguments[30])?;
+        let file_cache_ttl = Option::<u64>::try_convert(arguments[31])?;
+        let include_file_paths = Option::<String>::try_convert(arguments[32])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
@@ -113,12 +149,24 @@ impl RbLazyFrame {
         });
 
         let sources = sources.0;
-        let (_first_path, sources) = match source {
+        let (first_path, sources) = match source {
             None => (sources.first_path().map(|p| p.into_owned()), sources),
             Some(source) => rbobject_to_first_path_and_scan_sources(source)?,
         };
 
-        let r = LazyCsvReader::new_with_sources(sources);
+        let mut r = LazyCsvReader::new_with_sources(sources);
+
+        if let Some(first_path) = first_path {
+            let first_path_url = first_path.to_str();
+
+            let mut cloud_options =
+                parse_cloud_options(first_path_url, cloud_options.unwrap_or_default())?;
+            if let Some(file_cache_ttl) = file_cache_ttl {
+                cloud_options.file_cache_ttl = file_cache_ttl;
+            }
+            cloud_options = cloud_options.with_max_retries(retries);
+            r = r.with_cloud_options(Some(cloud_options));
+        }
 
         let r = r
             .with_infer_schema_length(infer_schema_length)
@@ -126,10 +174,11 @@ impl RbLazyFrame {
             .with_has_header(has_header)
             .with_ignore_errors(ignore_errors)
             .with_skip_rows(skip_rows)
+            .with_skip_lines(skip_lines)
             .with_n_rows(n_rows)
             .with_cache(cache)
             .with_dtype_overwrite(overwrite_dtype.map(Arc::new))
-            // TODO add with_schema
+            .with_schema(schema.map(|schema| Arc::new(schema.0)))
             .with_low_memory(low_memory)
             .with_comment_prefix(comment_prefix.map(|x| x.into()))
             .with_quote_char(quote_char)
@@ -140,8 +189,12 @@ impl RbLazyFrame {
             .with_row_index(row_index)
             .with_try_parse_dates(try_parse_dates)
             .with_null_values(null_values)
-            // TODO add with_missing_is_null
-            .with_truncate_ragged_lines(truncate_ragged_lines);
+            .with_missing_is_null(!missing_utf8_is_empty_string)
+            .with_truncate_ragged_lines(truncate_ragged_lines)
+            .with_decimal_comma(decimal_comma)
+            .with_glob(glob)
+            .with_raise_if_empty(raise_if_empty)
+            .with_include_file_paths(include_file_paths.map(|x| x.into()));
 
         if let Some(_lambda) = with_schema_modify {
             todo!();
@@ -208,57 +261,28 @@ impl RbLazyFrame {
         Ok(lf.into())
     }
 
-    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.borrow().logical_plan)
-            .map_err(|err| RbValueError::new_err(format!("{err:?}")))?;
-        Ok(())
+    pub fn describe_plan(rb: &Ruby, self_: &Self) -> RbResult<String> {
+        rb.enter_polars(|| self_.ldf.read().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(rb: &Ruby, self_: &Self) -> RbResult<String> {
+        rb.enter_polars(|| self_.ldf.read().describe_optimized_plan())
     }
 
-    pub fn describe_optimized_plan(&self) -> RbResult<String> {
-        let result = self
-            .ldf
-            .borrow()
-            .describe_optimized_plan()
-            .map_err(RbPolarsErr::from)?;
-        Ok(result)
+    pub fn describe_plan_tree(rb: &Ruby, self_: &Self) -> RbResult<String> {
+        rb.enter_polars(|| self_.ldf.read().describe_plan_tree())
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub fn optimization_toggle(
-        &self,
-        type_coercion: bool,
-        predicate_pushdown: bool,
-        projection_pushdown: bool,
-        simplify_expr: bool,
-        slice_pushdown: bool,
-        comm_subplan_elim: bool,
-        comm_subexpr_elim: bool,
-        allow_streaming: bool,
-        _eager: bool,
-    ) -> RbLazyFrame {
-        let ldf = self.ldf.borrow().clone();
-        let mut ldf = ldf
-            .with_type_coercion(type_coercion)
-            .with_predicate_pushdown(predicate_pushdown)
-            .with_simplify_expr(simplify_expr)
-            .with_slice_pushdown(slice_pushdown)
-            .with_new_streaming(allow_streaming)
-            ._with_eager(_eager)
-            .with_projection_pushdown(projection_pushdown);
-
-        ldf = ldf.with_comm_subplan_elim(comm_subplan_elim);
-        ldf = ldf.with_comm_subexpr_elim(comm_subexpr_elim);
-
-        ldf.into()
+    pub fn describe_optimized_plan_tree(rb: &Ruby, self_: &Self) -> RbResult<String> {
+        rb.enter_polars(|| self_.ldf.read().describe_optimized_plan_tree())
+    }
+
+    pub fn to_dot(rb: &Ruby, self_: &Self, optimized: bool) -> RbResult<String> {
+        rb.enter_polars(|| self_.ldf.read().to_dot(optimized))
+    }
+
+    pub fn to_dot_streaming_phys(rb: &Ruby, self_: &Self, optimized: bool) -> RbResult<String> {
+        rb.enter_polars(|| self_.ldf.read().to_dot_streaming_phys(optimized))
     }
 
     pub fn sort(
@@ -269,7 +293,7 @@ impl RbLazyFrame {
         maintain_order: bool,
         multithreaded: bool,
     ) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.sort(
             [&by_column],
             SortMultipleOptions {
@@ -291,8 +315,8 @@ impl RbLazyFrame {
         maintain_order: bool,
         multithreaded: bool,
     ) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        let exprs = rb_exprs_to_exprs(by)?;
+        let ldf = self.ldf.read().clone();
+        let exprs = by.to_exprs()?;
         Ok(ldf
             .sort_by_exprs(
                 exprs,
@@ -308,8 +332,8 @@ impl RbLazyFrame {
     }
 
     pub fn top_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        let exprs = rb_exprs_to_exprs(by)?;
+        let ldf = self.ldf.read().clone();
+        let exprs = by.to_exprs()?;
         Ok(ldf
             .top_k(
                 k,
@@ -320,8 +344,8 @@ impl RbLazyFrame {
     }
 
     pub fn bottom_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        let exprs = rb_exprs_to_exprs(by)?;
+        let ldf = self.ldf.read().clone();
+        let exprs = by.to_exprs()?;
         Ok(ldf
             .bottom_k(
                 k,
@@ -332,19 +356,34 @@ impl RbLazyFrame {
     }
 
     pub fn cache(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.cache().into()
     }
 
-    pub fn collect(&self) -> RbResult<RbDataFrame> {
-        let ldf = self.ldf.borrow().clone();
-        let df = ldf.collect().map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn with_optimizations(&self, optflags: &RbOptFlags) -> Self {
+        let ldf = self.ldf.read().clone();
+        ldf.with_optimizations(optflags.clone().inner.into_inner())
+            .into()
+    }
+
+    pub fn profile(rb: &Ruby, self_: &Self) -> RbResult<(RbDataFrame, RbDataFrame)> {
+        let (df, time_df) = rb.enter_polars(|| {
+            let ldf = self_.ldf.read().clone();
+            ldf.profile()
+        })?;
+        Ok((df.into(), time_df.into()))
+    }
+
+    pub fn collect(rb: &Ruby, self_: &Self, engine: Wrap<Engine>) -> RbResult<RbDataFrame> {
+        rb.enter_polars_df(|| {
+            let ldf = self_.ldf.read().clone();
+            ldf.collect_with_engine(engine.0)
+        })
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn sink_parquet(
-        &self,
+        rb: &Ruby,
+        self_: &Self,
         target: SinkTarget,
         compression: String,
         compression_level: Option<i32>,
@@ -352,8 +391,11 @@ impl RbLazyFrame {
         row_group_size: Option<usize>,
         data_page_size: Option<usize>,
         cloud_options: Option<Vec<(String, String)>>,
+        credential_provider: Option<Value>,
         retries: usize,
         sink_options: Wrap<SinkOptions>,
+        metadata: Wrap<Option<KeyValueMetadata>>,
+        field_overwrites: Vec<Wrap<ParquetFieldOverwrites>>,
     ) -> RbResult<RbLazyFrame> {
         let compression = parse_parquet_compression(&compression, compression_level)?;
 
@@ -362,8 +404,8 @@ impl RbLazyFrame {
             statistics: statistics.0,
             row_group_size,
             data_page_size,
-            key_value_metadata: None,
-            field_overwrites: Vec::new(),
+            key_value_metadata: metadata.0,
+            field_overwrites: field_overwrites.into_iter().map(|f| f.0).collect(),
         };
 
         let cloud_options = match target.base_path() {
@@ -371,31 +413,39 @@ impl RbLazyFrame {
             Some(base_path) => {
                 let cloud_options =
                     parse_cloud_options(base_path.to_str(), cloud_options.unwrap_or_default())?;
-                Some(cloud_options.with_max_retries(retries))
+                Some(
+                    cloud_options
+                        .with_max_retries(retries)
+                        .with_credential_provider(credential_provider.map(|_| todo!())),
+                )
             }
         };
 
-        let ldf = self.ldf.borrow().clone();
-        match target {
-            SinkTarget::File(target) => {
-                ldf.sink_parquet(target, options, cloud_options, sink_options.0)
+        rb.enter_polars(|| {
+            let ldf = self_.ldf.read().clone();
+            match target {
+                SinkTarget::File(target) => {
+                    ldf.sink_parquet(target, options, cloud_options, sink_options.0)
+                }
             }
-        }
-        .map_err(RbPolarsErr::from)
+        })
         .map(Into::into)
-        .map_err(Into::into)
     }
 
     pub fn sink_ipc(
-        &self,
+        rb: &Ruby,
+        self_: &Self,
         target: SinkTarget,
         compression: Wrap<Option<IpcCompression>>,
+        compat_level: RbCompatLevel,
         cloud_options: Option<Vec<(String, String)>>,
+        credential_provider: Option<Value>,
         retries: usize,
         sink_options: Wrap<SinkOptions>,
     ) -> RbResult<RbLazyFrame> {
         let options = IpcWriterOptions {
             compression: compression.0,
+            compat_level: compat_level.0,
             ..Default::default()
         };
 
@@ -404,22 +454,26 @@ impl RbLazyFrame {
             Some(base_path) => {
                 let cloud_options =
                     parse_cloud_options(base_path.to_str(), cloud_options.unwrap_or_default())?;
-                Some(cloud_options.with_max_retries(retries))
+                Some(
+                    cloud_options
+                        .with_max_retries(retries)
+                        .with_credential_provider(credential_provider.map(|_| todo!())),
+                )
             }
         };
 
-        let ldf = self.ldf.borrow().clone();
-        match target {
-            SinkTarget::File(target) => {
-                ldf.sink_ipc(target, options, cloud_options, sink_options.0)
+        rb.enter_polars(|| {
+            let ldf = self_.ldf.read().clone();
+            match target {
+                SinkTarget::File(target) => {
+                    ldf.sink_ipc(target, options, cloud_options, sink_options.0)
+                }
             }
-        }
-        .map_err(RbPolarsErr::from)
+        })
         .map(Into::into)
-        .map_err(Into::into)
     }
 
-    pub fn sink_csv(&self, arguments: &[Value]) -> RbResult<RbLazyFrame> {
+    pub fn sink_csv(rb: &Ruby, self_: &Self, arguments: &[Value]) -> RbResult<RbLazyFrame> {
         let target = SinkTarget::try_convert(arguments[0])?;
         let include_bom = bool::try_convert(arguments[1])?;
         let include_header = bool::try_convert(arguments[2])?;
@@ -436,8 +490,9 @@ impl RbLazyFrame {
         let null_value = Option::<String>::try_convert(arguments[13])?;
         let quote_style = Option::<Wrap<QuoteStyle>>::try_convert(arguments[14])?;
         let cloud_options = Option::<Vec<(String, String)>>::try_convert(arguments[15])?;
-        let retries = usize::try_convert(arguments[16])?;
-        let sink_options = Wrap::<SinkOptions>::try_convert(arguments[17])?;
+        let credential_provider = Option::<Value>::try_convert(arguments[16])?;
+        let retries = usize::try_convert(arguments[17])?;
+        let sink_options = Wrap::<SinkOptions>::try_convert(arguments[18])?;
 
         let quote_style = quote_style.map_or(QuoteStyle::default(), |wrap| wrap.0);
         let null_value = null_value.unwrap_or(SerializeOptions::default().null);
@@ -468,25 +523,31 @@ impl RbLazyFrame {
             Some(base_path) => {
                 let cloud_options =
                     parse_cloud_options(base_path.to_str(), cloud_options.unwrap_or_default())?;
-                Some(cloud_options.with_max_retries(retries))
+                Some(
+                    cloud_options
+                        .with_max_retries(retries)
+                        .with_credential_provider(credential_provider.map(|_| todo!())),
+                )
             }
         };
 
-        let ldf = self.ldf.borrow().clone();
-        match target {
-            SinkTarget::File(target) => {
-                ldf.sink_csv(target, options, cloud_options, sink_options.0)
+        rb.enter_polars(|| {
+            let ldf = self_.ldf.read().clone();
+            match target {
+                SinkTarget::File(target) => {
+                    ldf.sink_csv(target, options, cloud_options, sink_options.0)
+                }
             }
-        }
-        .map_err(RbPolarsErr::from)
+        })
         .map(Into::into)
-        .map_err(Into::into)
     }
 
     pub fn sink_json(
-        &self,
+        rb: &Ruby,
+        self_: &Self,
         target: SinkTarget,
         cloud_options: Option<Vec<(String, String)>>,
+        credential_provider: Option<Value>,
         retries: usize,
         sink_options: Wrap<SinkOptions>,
     ) -> RbResult<RbLazyFrame> {
@@ -497,52 +558,56 @@ impl RbLazyFrame {
             Some(base_path) => {
                 let cloud_options =
                     parse_cloud_options(base_path.to_str(), cloud_options.unwrap_or_default())?;
-                Some(cloud_options.with_max_retries(retries))
+                Some(
+                    cloud_options
+                        .with_max_retries(retries)
+                        .with_credential_provider(credential_provider.map(|_| todo!())),
+                )
             }
         };
 
-        let ldf = self.ldf.borrow().clone();
-        match target {
-            SinkTarget::File(path) => ldf.sink_json(path, options, cloud_options, sink_options.0),
-        }
-        .map_err(RbPolarsErr::from)
+        rb.enter_polars(|| {
+            let ldf = self_.ldf.read().clone();
+            match target {
+                SinkTarget::File(path) => {
+                    ldf.sink_json(path, options, cloud_options, sink_options.0)
+                }
+            }
+        })
         .map(Into::into)
-        .map_err(Into::into)
     }
 
     pub fn filter(&self, predicate: &RbExpr) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.filter(predicate.inner.clone()).into()
     }
 
     pub fn remove(&self, predicate: &RbExpr) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.remove(predicate.inner.clone()).into()
     }
 
     pub fn select(&self, exprs: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        let exprs = rb_exprs_to_exprs(exprs)?;
+        let ldf = self.ldf.read().clone();
+        let exprs = exprs.to_exprs()?;
         Ok(ldf.select(exprs).into())
     }
 
     pub fn select_seq(&self, exprs: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        let exprs = rb_exprs_to_exprs(exprs)?;
+        let ldf = self.ldf.read().clone();
+        let exprs = exprs.to_exprs()?;
         Ok(ldf.select_seq(exprs).into())
     }
 
     pub fn group_by(&self, by: RArray, maintain_order: bool) -> RbResult<RbLazyGroupBy> {
-        let ldf = self.ldf.borrow().clone();
-        let by = rb_exprs_to_exprs(by)?;
+        let ldf = self.ldf.read().clone();
+        let by = by.to_exprs()?;
         let lazy_gb = if maintain_order {
             ldf.group_by_stable(by)
         } else {
             ldf.group_by(by)
         };
-        Ok(RbLazyGroupBy {
-            lgb: RefCell::new(Some(lazy_gb)),
-        })
+        Ok(RbLazyGroupBy { lgb: Some(lazy_gb) })
     }
 
     pub fn rolling(
@@ -554,8 +619,8 @@ impl RbLazyFrame {
         by: RArray,
     ) -> RbResult<RbLazyGroupBy> {
         let closed_window = closed.0;
-        let ldf = self.ldf.borrow().clone();
-        let by = rb_exprs_to_exprs(by)?;
+        let ldf = self.ldf.read().clone();
+        let by = by.to_exprs()?;
         let lazy_gb = ldf.rolling(
             index_column.inner.clone(),
             by,
@@ -567,12 +632,9 @@ impl RbLazyFrame {
             },
         );
 
-        Ok(RbLazyGroupBy {
-            lgb: RefCell::new(Some(lazy_gb)),
-        })
+        Ok(RbLazyGroupBy { lgb: Some(lazy_gb) })
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn group_by_dynamic(
         &self,
         index_column: &RbExpr,
@@ -586,8 +648,8 @@ impl RbLazyFrame {
         start_by: Wrap<StartBy>,
     ) -> RbResult<RbLazyGroupBy> {
         let closed_window = closed.0;
-        let by = rb_exprs_to_exprs(by)?;
-        let ldf = self.ldf.borrow().clone();
+        let by = by.to_exprs()?;
+        let ldf = self.ldf.read().clone();
         let lazy_gb = ldf.group_by_dynamic(
             index_column.inner.clone(),
             by,
@@ -603,21 +665,9 @@ impl RbLazyFrame {
             },
         );
 
-        Ok(RbLazyGroupBy {
-            lgb: RefCell::new(Some(lazy_gb)),
-        })
-    }
-
-    pub fn with_context(&self, contexts: RArray) -> RbResult<Self> {
-        let contexts = contexts.typecheck::<Obj<RbLazyFrame>>()?;
-        let contexts = contexts
-            .into_iter()
-            .map(|ldf| ldf.ldf.borrow().clone())
-            .collect::<Vec<_>>();
-        Ok(self.ldf.borrow().clone().with_context(contexts).into())
+        Ok(RbLazyGroupBy { lgb: Some(lazy_gb) })
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn join_asof(
         &self,
         other: &RbLazyFrame,
@@ -640,8 +690,8 @@ impl RbLazyFrame {
         } else {
             JoinCoalesce::KeepColumns
         };
-        let ldf = self.ldf.borrow().clone();
-        let other = other.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
+        let other = other.ldf.read().clone();
         let left_on = left_on.inner.clone();
         let right_on = right_on.inner.clone();
         Ok(ldf
@@ -670,7 +720,6 @@ impl RbLazyFrame {
             .into())
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn join(
         &self,
         other: &RbLazyFrame,
@@ -690,10 +739,10 @@ impl RbLazyFrame {
             Some(true) => JoinCoalesce::CoalesceColumns,
             Some(false) => JoinCoalesce::KeepColumns,
         };
-        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)?;
+        let ldf = self.ldf.read().clone();
+        let other = other.ldf.read().clone();
+        let left_on = left_on.to_exprs()?;
+        let right_on = right_on.to_exprs()?;
 
         Ok(ldf
             .join_builder()
@@ -713,10 +762,10 @@ impl RbLazyFrame {
     }
 
     pub fn join_where(&self, other: &Self, predicates: RArray, suffix: String) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        let other = other.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
+        let other = other.ldf.read().clone();
 
-        let predicates = rb_exprs_to_exprs(predicates)?;
+        let predicates = predicates.to_exprs()?;
 
         Ok(ldf
             .join_builder()
@@ -727,32 +776,32 @@ impl RbLazyFrame {
     }
 
     pub fn with_column(&self, expr: &RbExpr) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.with_column(expr.inner.clone()).into()
     }
 
     pub fn with_columns(&self, exprs: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        Ok(ldf.with_columns(rb_exprs_to_exprs(exprs)?).into())
+        let ldf = self.ldf.read().clone();
+        Ok(ldf.with_columns(exprs.to_exprs()?).into())
     }
 
     pub fn with_columns_seq(&self, exprs: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        Ok(ldf.with_columns_seq(rb_exprs_to_exprs(exprs)?).into())
+        let ldf = self.ldf.read().clone();
+        Ok(ldf.with_columns_seq(exprs.to_exprs()?).into())
     }
 
     pub fn rename(&self, existing: Vec<String>, new: Vec<String>, strict: bool) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.rename(existing, new, strict).into()
     }
 
     pub fn reverse(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.reverse().into()
     }
 
     pub fn shift(&self, n: &RbExpr, fill_value: Option<&RbExpr>) -> Self {
-        let lf = self.ldf.borrow().clone();
+        let lf = self.ldf.read().clone();
         let out = match fill_value {
             Some(v) => lf.shift_and_fill(n.inner.clone(), v.inner.clone()),
             None => lf.shift(n.inner.clone()),
@@ -761,68 +810,64 @@ impl RbLazyFrame {
     }
 
     pub fn fill_nan(&self, fill_value: &RbExpr) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.fill_nan(fill_value.inner.clone()).into()
     }
 
     pub fn min(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         let out = ldf.min();
         out.into()
     }
 
     pub fn max(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         let out = ldf.max();
         out.into()
     }
 
     pub fn sum(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         let out = ldf.sum();
         out.into()
     }
 
     pub fn mean(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         let out = ldf.mean();
         out.into()
     }
 
     pub fn std(&self, ddof: u8) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         let out = ldf.std(ddof);
         out.into()
     }
 
     pub fn var(&self, ddof: u8) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         let out = ldf.var(ddof);
         out.into()
     }
 
     pub fn median(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         let out = ldf.median();
         out.into()
     }
 
     pub fn quantile(&self, quantile: &RbExpr, interpolation: Wrap<QuantileMethod>) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         let out = ldf.quantile(quantile.inner.clone(), interpolation.0);
         out.into()
     }
 
     pub fn explode(&self, subset: &RbSelector) -> Self {
-        self.ldf
-            .borrow()
-            .clone()
-            .explode(subset.inner.clone())
-            .into()
+        self.ldf.read().clone().explode(subset.inner.clone()).into()
     }
 
     pub fn null_count(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.null_count().into()
     }
 
@@ -832,7 +877,7 @@ impl RbLazyFrame {
         subset: Option<&RbSelector>,
         keep: Wrap<UniqueKeepStrategy>,
     ) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         let subset = subset.map(|e| e.inner.clone());
         Ok(match maintain_order {
             true => ldf.unique_stable_generic(subset, keep.0),
@@ -843,7 +888,7 @@ impl RbLazyFrame {
 
     pub fn drop_nans(&self, subset: Option<&RbSelector>) -> Self {
         self.ldf
-            .borrow()
+            .read()
             .clone()
             .drop_nans(subset.map(|e| e.inner.clone()))
             .into()
@@ -851,19 +896,19 @@ impl RbLazyFrame {
 
     pub fn drop_nulls(&self, subset: Option<&RbSelector>) -> Self {
         self.ldf
-            .borrow()
+            .read()
             .clone()
             .drop_nulls(subset.map(|e| e.inner.clone()))
             .into()
     }
 
     pub fn slice(&self, offset: i64, len: Option<IdxSize>) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.slice(offset, len.unwrap_or(IdxSize::MAX)).into()
     }
 
     pub fn tail(&self, n: IdxSize) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.tail(n).into()
     }
 
@@ -881,17 +926,17 @@ impl RbLazyFrame {
             variable_name: variable_name.map(|s| s.into()),
         };
 
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         Ok(ldf.unpivot(args).into())
     }
 
     pub fn with_row_index(&self, name: String, offset: Option<IdxSize>) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.with_row_index(&name, offset).into()
     }
 
     pub fn drop(&self, columns: &RbSelector) -> Self {
-        self.ldf.borrow().clone().drop(columns.inner.clone()).into()
+        self.ldf.read().clone().drop(columns.inner.clone()).into()
     }
 
     pub fn cast(&self, rb_dtypes: RHash, strict: bool) -> RbResult<Self> {
@@ -902,30 +947,26 @@ impl RbLazyFrame {
         })?;
         let mut cast_map = PlHashMap::with_capacity(dtypes.len());
         cast_map.extend(dtypes.iter().map(|(k, v)| (k.as_ref(), v.clone())));
-        Ok(self.ldf.borrow().clone().cast(cast_map, strict).into())
+        Ok(self.ldf.read().clone().cast(cast_map, strict).into())
     }
 
     pub fn cast_all(&self, dtype: Wrap<DataType>, strict: bool) -> Self {
-        self.ldf.borrow().clone().cast_all(dtype.0, strict).into()
+        self.ldf.read().clone().cast_all(dtype.0, strict).into()
     }
 
     pub fn clone(&self) -> Self {
-        self.ldf.borrow().clone().into()
+        self.ldf.read().clone().into()
     }
 
-    pub fn collect_schema(ruby: &Ruby, rb_self: &Self) -> RbResult<RHash> {
-        let schema = rb_self
-            .ldf
-            .borrow_mut()
-            .collect_schema()
-            .map_err(RbPolarsErr::from)?;
+    pub fn collect_schema(rb: &Ruby, self_: &Self) -> RbResult<RHash> {
+        let schema = rb.enter_polars(|| self_.ldf.write().collect_schema())?;
 
-        let schema_dict = ruby.hash_new();
+        let schema_dict = rb.hash_new();
         schema.iter_fields().for_each(|fld| {
             schema_dict
                 .aset::<String, Value>(
                     fld.name().to_string(),
-                    Wrap(fld.dtype().clone()).into_value_with(ruby),
+                    Wrap(fld.dtype().clone()).into_value_with(rb),
                 )
                 .unwrap();
         });
@@ -934,7 +975,7 @@ impl RbLazyFrame {
 
     pub fn unnest(&self, columns: &RbSelector, separator: Option<String>) -> Self {
         self.ldf
-            .borrow()
+            .read()
             .clone()
             .unnest(
                 columns.inner.clone(),
@@ -944,17 +985,80 @@ impl RbLazyFrame {
     }
 
     pub fn count(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.count().into()
     }
 
     pub fn merge_sorted(&self, other: &Self, key: String) -> RbResult<Self> {
         let out = self
             .ldf
-            .borrow()
+            .read()
             .clone()
-            .merge_sorted(other.ldf.borrow().clone(), &key)
+            .merge_sorted(other.ldf.read().clone(), &key)
+            .map_err(RbPolarsErr::from)?;
+        Ok(out.into())
+    }
+
+    pub fn hint_sorted(
+        &self,
+        columns: Vec<String>,
+        descending: Vec<bool>,
+        nulls_last: Vec<bool>,
+    ) -> RbResult<Self> {
+        if columns.len() != descending.len() && descending.len() != 1 {
+            return Err(RbValueError::new_err(
+                "`set_sorted` expects the same amount of `columns` as `descending` values.",
+            ));
+        }
+        if columns.len() != nulls_last.len() && nulls_last.len() != 1 {
+            return Err(RbValueError::new_err(
+                "`set_sorted` expects the same amount of `columns` as `nulls_last` values.",
+            ));
+        }
+
+        let mut sorted = columns
+            .iter()
+            .map(|c| Sorted {
+                column: PlSmallStr::from_str(c.as_str()),
+                descending: false,
+                nulls_last: false,
+            })
+            .collect::<Vec<_>>();
+
+        if !columns.is_empty() {
+            if descending.len() != 1 {
+                sorted
+                    .iter_mut()
+                    .zip(descending)
+                    .for_each(|(s, d)| s.descending = d);
+            } else if descending[0] {
+                sorted.iter_mut().for_each(|s| s.descending = true);
+            }
+
+            if nulls_last.len() != 1 {
+                sorted
+                    .iter_mut()
+                    .zip(nulls_last)
+                    .for_each(|(s, d)| s.nulls_last = d);
+            } else if nulls_last[0] {
+                sorted.iter_mut().for_each(|s| s.nulls_last = true);
+            }
+        }
+
+        let out = self
+            .ldf
+            .read()
+            .clone()
+            .hint(HintIR::Sorted(sorted.into()))
             .map_err(RbPolarsErr::from)?;
         Ok(out.into())
     }
 }
+
+impl TryConvert for Wrap<polars_io::parquet::write::ParquetFieldOverwrites> {
+    fn try_convert(_ob: Value) -> RbResult<Self> {
+        todo!();
+    }
+}
+
+unsafe impl TryConvertOwned for Wrap<polars_io::parquet::write::ParquetFieldOverwrites> {}