polars-df 0.15.0 → 0.26.0

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

@@ -1,52 +1,90 @@
-use magnus::{r_hash::ForEach, typed_data::Obj, IntoValue, RArray, RHash, TryConvert, Value};
-use polars::io::{HiveOptions, RowIndex};
+use arrow::ffi::export_iterator;
+use magnus::{
+    IntoValue, RArray, RHash, Ruby, TryConvert, Value, r_hash::ForEach, value::ReprValue,
+};
+use parking_lot::Mutex;
+use polars::frame::PivotColumnNaming;
+use polars::io::RowIndex;
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
-use polars_plan::plans::ScanSources;
-use std::cell::RefCell;
-use std::io::BufWriter;
+use polars_core::query_result::QueryResult;
+use polars_plan::dsl::ScanSources;
+use polars_plan::plans::{HintIR, Sorted};
 use std::num::NonZeroUsize;
-use std::path::PathBuf;
 
+use super::{RbLazyFrame, RbOptFlags};
 use crate::conversion::*;
-use crate::expr::rb_exprs_to_exprs;
-use crate::file::get_file_like;
-use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbLazyGroupBy, RbPolarsErr, RbResult, RbValueError};
-
-fn rbobject_to_first_path_and_scan_sources(obj: Value) -> RbResult<(Option<PathBuf>, ScanSources)> {
-    use crate::file::{get_ruby_scan_source_input, RubyScanSourceInput};
+use crate::expr::ToExprs;
+use crate::expr::selector::RbSelector;
+use crate::io::cloud_options::OptRbCloudOptions;
+use crate::io::scan_options::RbScanOptions;
+use crate::io::sink_options::RbSinkOptions;
+use crate::io::sink_output::RbFileSinkDestination;
+use crate::ruby::gvl::GvlExt;
+use crate::ruby::lazy::RubyUdfLazyFrameExt;
+use crate::ruby::plan_callback::PlanCallbackExt;
+use crate::ruby::ruby_function::RubyObject;
+use crate::ruby::utils::TryIntoValue;
+use crate::utils::{EnterPolarsExt, to_rb_err};
+use crate::{
+    RbArrowArrayStream, RbDataFrame, RbExpr, RbLazyGroupBy, RbPolarsErr, RbResult, RbTypeError,
+    RbValueError,
+};
+
+fn rbobject_to_first_path_and_scan_sources(
+    obj: Value,
+) -> RbResult<(Option<PlRefPath>, ScanSources)> {
+    use crate::file::{RubyScanSourceInput, get_ruby_scan_source_input};
     Ok(match get_ruby_scan_source_input(obj, false)? {
-        RubyScanSourceInput::Path(path) => (Some(path.clone()), ScanSources::Paths([path].into())),
+        RubyScanSourceInput::Path(path) => (
+            Some(path.clone()),
+            ScanSources::Paths(FromIterator::from_iter([path])),
+        ),
         RubyScanSourceInput::File(file) => (None, ScanSources::Files([file].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<Wrap<NonZeroUsize>>,
-        n_rows: Option<usize>,
-        low_memory: bool,
-        rechunk: bool,
-        row_index: Option<(String, IdxSize)>,
-    ) -> RbResult<Self> {
-        let batch_size = batch_size.map(|v| v.0);
+    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 = OptRbCloudOptions::try_convert(arguments[12])?;
+        let credential_provider = Option::<Value>::try_convert(arguments[13])?;
+
         let row_index = row_index.map(|(name, offset)| RowIndex {
             name: name.into(),
             offset,
         });
 
         let sources = sources.0;
-        let (_first_path, sources) = match source {
-            None => (sources.first_path().map(|p| p.to_path_buf()), sources),
+        let (first_path, sources) = match source {
+            None => (sources.first_path().cloned(), 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.as_str();
+
+            let cloud_options = cloud_options.extract_opt_cloud_options(
+                CloudScheme::from_path(first_path_url),
+                credential_provider,
+            )?;
+
+            r = r.with_cloud_options(cloud_options);
+        };
 
         let lf = r
             .with_infer_schema_length(infer_schema_length.and_then(NonZeroUsize::new))
@@ -54,11 +92,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)?;
 
@@ -69,27 +107,37 @@ 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 = OptRbCloudOptions::try_convert(arguments[28])?;
+        let credential_provider = Option::<Value>::try_convert(arguments[29])?;
+        let include_file_paths = Option::<String>::try_convert(arguments[30])?;
+        let missing_columns = Option::<Wrap<MissingColumnsPolicy>>::try_convert(arguments[31])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
@@ -109,23 +157,33 @@ impl RbLazyFrame {
         });
 
         let sources = sources.0;
-        let (_first_path, sources) = match source {
-            None => (sources.first_path().map(|p| p.to_path_buf()), sources),
+        let (first_path, sources) = match source {
+            None => (sources.first_path().cloned(), 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);
 
-        let r = r
+        if let Some(first_path) = first_path {
+            let first_path_url = first_path.as_str();
+            let cloud_options = cloud_options.extract_opt_cloud_options(
+                CloudScheme::from_path(first_path_url),
+                credential_provider,
+            )?;
+            r = r.with_cloud_options(cloud_options);
+        }
+
+        let mut r = r
             .with_infer_schema_length(infer_schema_length)
             .with_separator(separator)
             .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)
@@ -136,182 +194,136 @@ 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);
-
-        if let Some(_lambda) = with_schema_modify {
-            todo!();
+            .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()))
+            .with_missing_columns_policy(missing_columns.map(|x| x.0));
+
+        if let Some(lambda) = with_schema_modify {
+            let f = |schema: Schema| {
+                let iter = schema.iter_names().map(|s| s.as_str());
+                Ruby::attach(|rb| {
+                    let names = rb.ary_from_iter(iter);
+
+                    let out = lambda
+                        .funcall("call", (names,))
+                        .expect("ruby function failed");
+                    let new_names = Vec::<String>::try_convert(out)
+                        .expect("ruby function should return Array[String]");
+                    polars_ensure!(new_names.len() == schema.len(),
+                        ShapeMismatch: "The length of the new names list should be equal to or less than the original column length",
+                    );
+                    Ok(schema
+                        .iter_values()
+                        .zip(new_names)
+                        .map(|(dtype, name)| Field::new(name.into(), dtype.clone()))
+                        .collect())
+                })
+            };
+            r = r.with_schema_modify(f).map_err(RbPolarsErr::from)?
         }
 
         Ok(r.finish().map_err(RbPolarsErr::from)?.into())
     }
 
-    pub fn new_from_parquet(arguments: &[Value]) -> RbResult<Self> {
-        let source = Option::<Value>::try_convert(arguments[0])?;
-        let sources = Wrap::<ScanSources>::try_convert(arguments[1])?;
-        let n_rows = Option::<usize>::try_convert(arguments[2])?;
-        let cache = bool::try_convert(arguments[3])?;
-        let parallel = Wrap::<ParallelStrategy>::try_convert(arguments[4])?;
-        let rechunk = bool::try_convert(arguments[5])?;
-        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[6])?;
-        let low_memory = bool::try_convert(arguments[7])?;
-        let cloud_options = Option::<Vec<(String, String)>>::try_convert(arguments[8])?;
-        let _credential_provider = Option::<Value>::try_convert(arguments[9])?;
-        let use_statistics = bool::try_convert(arguments[10])?;
-        let hive_partitioning = Option::<bool>::try_convert(arguments[11])?;
-        let schema = Option::<Wrap<Schema>>::try_convert(arguments[12])?;
-        let hive_schema = Option::<Wrap<Schema>>::try_convert(arguments[13])?;
-        let try_parse_hive_dates = bool::try_convert(arguments[14])?;
-        let retries = usize::try_convert(arguments[15])?;
-        let glob = bool::try_convert(arguments[16])?;
-        let include_file_paths = Option::<String>::try_convert(arguments[17])?;
-        let allow_missing_columns = bool::try_convert(arguments[18])?;
+    pub fn new_from_parquet(
+        sources: Wrap<ScanSources>,
+        schema: Option<Wrap<Schema>>,
+        scan_options: RbScanOptions,
+        parallel: Wrap<ParallelStrategy>,
+        low_memory: bool,
+        use_statistics: bool,
+    ) -> RbResult<Self> {
+        use crate::utils::to_rb_err;
 
         let parallel = parallel.0;
-        let hive_schema = hive_schema.map(|s| Arc::new(s.0));
-
-        let row_index = row_index.map(|(name, offset)| RowIndex {
-            name: name.into(),
-            offset,
-        });
-
-        let hive_options = HiveOptions {
-            enabled: hive_partitioning,
-            hive_start_idx: 0,
-            schema: hive_schema,
-            try_parse_dates: try_parse_hive_dates,
-        };
 
-        let mut args = ScanArgsParquet {
-            n_rows,
-            cache,
+        let options = ParquetOptions {
+            schema: schema.map(|x| Arc::new(x.0)),
             parallel,
-            rechunk,
-            row_index,
             low_memory,
-            cloud_options: None,
             use_statistics,
-            schema: schema.map(|x| Arc::new(x.0)),
-            hive_options,
-            glob,
-            include_file_paths: include_file_paths.map(|x| x.into()),
-            allow_missing_columns,
         };
 
         let sources = sources.0;
-        let (first_path, sources) = match source {
-            None => (sources.first_path().map(|p| p.to_path_buf()), sources),
-            Some(source) => rbobject_to_first_path_and_scan_sources(source)?,
-        };
+        let first_path = sources.first_path();
 
-        if let Some(first_path) = first_path {
-            let first_path_url = first_path.to_string_lossy();
-            let cloud_options =
-                parse_cloud_options(&first_path_url, cloud_options.unwrap_or_default())?;
-            args.cloud_options = Some(cloud_options.with_max_retries(retries));
-        }
+        let unified_scan_args =
+            scan_options.extract_unified_scan_args(first_path.and_then(|x| x.scheme()))?;
 
-        let lf = LazyFrame::scan_parquet_sources(sources, args).map_err(RbPolarsErr::from)?;
+        let lf: LazyFrame = DslBuilder::scan_parquet(sources, options, unified_scan_args)
+            .map_err(to_rb_err)?
+            .build()
+            .into();
 
         Ok(lf.into())
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn new_from_ipc(
-        source: Option<Value>,
         sources: Wrap<ScanSources>,
-        n_rows: Option<usize>,
-        cache: bool,
-        rechunk: bool,
-        row_index: Option<(String, IdxSize)>,
-        hive_partitioning: Option<bool>,
-        hive_schema: Option<Wrap<Schema>>,
-        try_parse_hive_dates: bool,
-        include_file_paths: Option<String>,
+        record_batch_statistics: bool,
+        scan_options: RbScanOptions,
     ) -> RbResult<Self> {
-        let row_index = row_index.map(|(name, offset)| RowIndex {
-            name: name.into(),
-            offset,
-        });
-
-        let hive_options = HiveOptions {
-            enabled: hive_partitioning,
-            hive_start_idx: 0,
-            schema: hive_schema.map(|x| Arc::new(x.0)),
-            try_parse_dates: try_parse_hive_dates,
+        let options = IpcScanOptions {
+            record_batch_statistics,
+            checked: Default::default(),
         };
 
-        let args = ScanArgsIpc {
-            n_rows,
-            cache,
-            rechunk,
-            row_index,
-            cloud_options: None,
-            hive_options,
-            include_file_paths: include_file_paths.map(|x| x.into()),
-        };
+        let sources = sources.0;
+        let first_path = sources.first_path().cloned();
+
+        let unified_scan_args =
+            scan_options.extract_unified_scan_args(first_path.as_ref().and_then(|x| x.scheme()))?;
+
+        let lf = LazyFrame::scan_ipc_sources(sources, options, unified_scan_args)
+            .map_err(RbPolarsErr::from)?;
+        Ok(lf.into())
+    }
 
+    pub fn new_from_scan_lines(
+        sources: Wrap<ScanSources>,
+        scan_options: RbScanOptions,
+        name: String,
+    ) -> RbResult<Self> {
         let sources = sources.0;
-        let (_first_path, sources) = match source {
-            None => (sources.first_path().map(|p| p.to_path_buf()), sources),
-            Some(source) => rbobject_to_first_path_and_scan_sources(source)?,
-        };
+        let first_path = sources.first_path();
+
+        let unified_scan_args =
+            scan_options.extract_unified_scan_args(first_path.and_then(|x| x.scheme()))?;
+
+        let dsl: DslPlan = DslBuilder::scan_lines(sources, unified_scan_args, (&*name).into())
+            .map_err(to_rb_err)?
+            .build();
+        let lf: LazyFrame = dsl.into();
 
-        let lf = LazyFrame::scan_ipc_sources(sources, args).map_err(RbPolarsErr::from)?;
         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_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(
@@ -322,7 +334,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 {
@@ -330,6 +342,7 @@ impl RbLazyFrame {
                 nulls_last: vec![nulls_last],
                 multithreaded,
                 maintain_order,
+                limit: None,
             },
         )
         .into()
@@ -343,8 +356,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,
@@ -353,33 +366,101 @@ impl RbLazyFrame {
                     nulls_last,
                     maintain_order,
                     multithreaded,
+                    limit: None,
                 },
             )
             .into())
     }
 
+    pub fn top_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
+        let ldf = self.ldf.read().clone();
+        let exprs = by.to_exprs()?;
+        Ok(ldf
+            .top_k(
+                k,
+                exprs,
+                SortMultipleOptions::new().with_order_descending_multi(reverse),
+            )
+            .into())
+    }
+
+    pub fn bottom_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
+        let ldf = self.ldf.read().clone();
+        let exprs = by.to_exprs()?;
+        Ok(ldf
+            .bottom_k(
+                k,
+                exprs,
+                SortMultipleOptions::new().with_order_descending_multi(reverse),
+            )
+            .into())
+    }
+
     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).map(|r| match r {
+                QueryResult::Single(df) => df,
+                // TODO: Should return query results
+                QueryResult::Multiple(_) => DataFrame::empty(),
+            })
+        })
+    }
+
+    pub fn collect_batches(
+        rb: &Ruby,
+        self_: &Self,
+        engine: Wrap<Engine>,
+        maintain_order: bool,
+        chunk_size: Option<NonZeroUsize>,
+        lazy: bool,
+    ) -> RbResult<RbCollectBatches> {
+        rb.enter_polars(|| {
+            let ldf = self_.ldf.read().clone();
+
+            let collect_batches =
+                ldf.clone()
+                    .collect_batches(engine.0, maintain_order, chunk_size, lazy)?;
+
+            PolarsResult::Ok(RbCollectBatches {
+                inner: Arc::new(Mutex::new(collect_batches)),
+                ldf,
+            })
+        })
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn sink_parquet(
-        &self,
-        path: PathBuf,
+        rb: &Ruby,
+        self_: &Self,
+        target: RbFileSinkDestination,
+        sink_options: RbSinkOptions,
         compression: String,
         compression_level: Option<i32>,
         statistics: Wrap<StatisticsOptions>,
         row_group_size: Option<usize>,
         data_page_size: Option<usize>,
-        maintain_order: bool,
-    ) -> RbResult<()> {
+        metadata: Wrap<Option<KeyValueMetadata>>,
+        arrow_schema: Option<Wrap<ArrowSchema>>,
+    ) -> RbResult<RbLazyFrame> {
         let compression = parse_parquet_compression(&compression, compression_level)?;
 
         let options = ParquetWriteOptions {
@@ -387,120 +468,190 @@ impl RbLazyFrame {
             statistics: statistics.0,
             row_group_size,
             data_page_size,
-            maintain_order,
+            key_value_metadata: metadata.0,
+            arrow_schema: arrow_schema.map(|x| Arc::new(x.0)),
+            compat_level: None,
         };
 
-        let ldf = self.ldf.borrow().clone();
-        ldf.sink_parquet(path, options).map_err(RbPolarsErr::from)?;
-        Ok(())
+        let target = target.extract_file_sink_destination()?;
+        let unified_sink_args = sink_options.extract_unified_sink_args(target.cloud_scheme())?;
+
+        rb.enter_polars(|| {
+            self_.ldf.read().clone().sink(
+                target,
+                FileWriteFormat::Parquet(Arc::new(options)),
+                unified_sink_args,
+            )
+        })
+        .map(Into::into)
     }
 
     pub fn sink_ipc(
-        &self,
-        path: PathBuf,
-        compression: Option<Wrap<IpcCompression>>,
-        maintain_order: bool,
-    ) -> RbResult<()> {
+        rb: &Ruby,
+        self_: &Self,
+        target: RbFileSinkDestination,
+        sink_options: RbSinkOptions,
+        compression: Wrap<Option<IpcCompression>>,
+        compat_level: RbCompatLevel,
+        record_batch_size: Option<usize>,
+        record_batch_statistics: bool,
+    ) -> RbResult<RbLazyFrame> {
         let options = IpcWriterOptions {
-            compression: compression.map(|c| c.0),
-            maintain_order,
+            compression: compression.0,
+            compat_level: compat_level.0,
+            record_batch_size,
+            record_batch_statistics,
         };
 
-        let ldf = self.ldf.borrow().clone();
-        ldf.sink_ipc(path, options).map_err(RbPolarsErr::from)?;
-        Ok(())
+        let target = target.extract_file_sink_destination()?;
+        let unified_sink_args = sink_options.extract_unified_sink_args(target.cloud_scheme())?;
+
+        rb.enter_polars(|| {
+            self_
+                .ldf
+                .read()
+                .clone()
+                .sink(target, FileWriteFormat::Ipc(options), unified_sink_args)
+        })
+        .map(Into::into)
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub fn sink_csv(
-        &self,
-        path: PathBuf,
-        include_bom: bool,
-        include_header: bool,
-        separator: u8,
-        line_terminator: String,
-        quote_char: u8,
-        batch_size: Wrap<NonZeroUsize>,
-        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>>,
-        maintain_order: bool,
-    ) -> RbResult<()> {
+    pub fn sink_csv(rb: &Ruby, self_: &Self, arguments: &[Value]) -> RbResult<RbLazyFrame> {
+        let target = RbFileSinkDestination::try_convert(arguments[0])?;
+        let sink_options = RbSinkOptions::try_convert(arguments[1])?;
+        let include_bom = bool::try_convert(arguments[2])?;
+        let compression = String::try_convert(arguments[3])?;
+        let compression_level = Option::<u32>::try_convert(arguments[4])?;
+        let check_extension = bool::try_convert(arguments[5])?;
+        let include_header = bool::try_convert(arguments[6])?;
+        let separator = u8::try_convert(arguments[7])?;
+        let line_terminator = Wrap::<PlSmallStr>::try_convert(arguments[8])?;
+        let quote_char = u8::try_convert(arguments[9])?;
+        let batch_size = NonZeroUsize::try_convert(arguments[10])?;
+        let datetime_format = Option::<Wrap<PlSmallStr>>::try_convert(arguments[11])?;
+        let date_format = Option::<Wrap<PlSmallStr>>::try_convert(arguments[12])?;
+        let time_format = Option::<Wrap<PlSmallStr>>::try_convert(arguments[13])?;
+        let float_scientific = Option::<bool>::try_convert(arguments[14])?;
+        let float_precision = Option::<usize>::try_convert(arguments[15])?;
+        let decimal_comma = bool::try_convert(arguments[16])?;
+        let null_value = Option::<Wrap<PlSmallStr>>::try_convert(arguments[17])?;
+        let quote_style = Option::<Wrap<QuoteStyle>>::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);
+        let null_value = null_value
+            .map(|x| x.0)
+            .unwrap_or(SerializeOptions::default().null);
 
         let serialize_options = SerializeOptions {
-            date_format,
-            time_format,
-            datetime_format,
+            date_format: date_format.map(|x| x.0),
+            time_format: time_format.map(|x| x.0),
+            datetime_format: datetime_format.map(|x| x.0),
             float_scientific,
             float_precision,
+            decimal_comma,
             separator,
             quote_char,
             null: null_value,
-            line_terminator,
+            line_terminator: line_terminator.0,
             quote_style,
         };
 
         let options = CsvWriterOptions {
             include_bom,
+            compression: ExternalCompression::try_from(&compression, compression_level)
+                .map_err(RbPolarsErr::from)?,
+            check_extension,
             include_header,
-            maintain_order,
-            batch_size: batch_size.0,
-            serialize_options,
+            batch_size,
+            serialize_options: serialize_options.into(),
         };
 
-        let ldf = self.ldf.borrow().clone();
-        ldf.sink_csv(path, options).map_err(RbPolarsErr::from)?;
-        Ok(())
+        let target = target.extract_file_sink_destination()?;
+        let unified_sink_args = sink_options.extract_unified_sink_args(target.cloud_scheme())?;
+
+        rb.enter_polars(|| {
+            self_
+                .ldf
+                .read()
+                .clone()
+                .sink(target, FileWriteFormat::Csv(options), unified_sink_args)
+        })
+        .map(Into::into)
     }
 
-    pub fn sink_json(&self, path: PathBuf, maintain_order: bool) -> RbResult<()> {
-        let options = JsonWriterOptions { maintain_order };
+    pub fn sink_ndjson(
+        rb: &Ruby,
+        self_: &Self,
+        target: RbFileSinkDestination,
+        compression: String,
+        compression_level: Option<u32>,
+        check_extension: bool,
+        sink_options: RbSinkOptions,
+    ) -> RbResult<RbLazyFrame> {
+        let options = NDJsonWriterOptions {
+            compression: ExternalCompression::try_from(&compression, compression_level)
+                .map_err(RbPolarsErr::from)?,
+            check_extension,
+        };
 
-        let ldf = self.ldf.borrow().clone();
-        ldf.sink_json(path, options).map_err(RbPolarsErr::from)?;
-        Ok(())
+        let target = target.extract_file_sink_destination()?;
+        let unified_sink_args = sink_options.extract_unified_sink_args(target.cloud_scheme())?;
+
+        rb.enter_polars(|| {
+            self_.ldf.read().clone().sink(
+                target,
+                FileWriteFormat::NDJson(options),
+                unified_sink_args,
+            )
+        })
+        .map(Into::into)
     }
 
-    pub fn fetch(&self, n_rows: usize) -> RbResult<RbDataFrame> {
-        let ldf = self.ldf.borrow().clone();
-        let df = ldf.fetch(n_rows).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn sink_batches(
+        rb: &Ruby,
+        self_: &Self,
+        function: Value,
+        maintain_order: bool,
+        chunk_size: Option<NonZeroUsize>,
+    ) -> RbResult<RbLazyFrame> {
+        let ldf = self_.ldf.read().clone();
+        // ensure new_ruby is called with GVL
+        let callback = PlanCallback::new_ruby(RubyObject::from(function));
+        rb.enter_polars(|| ldf.sink_batches(callback, maintain_order, chunk_size))
+            .map(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.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(
@@ -512,8 +663,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,
@@ -525,12 +676,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,
@@ -544,8 +692,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,
@@ -561,21 +709,9 @@ impl RbLazyFrame {
             },
         );
 
-        Ok(RbLazyGroupBy {
-            lgb: RefCell::new(Some(lazy_gb)),
-        })
+        Ok(RbLazyGroupBy { lgb: 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())
-    }
-
-    #[allow(clippy::too_many_arguments)]
     pub fn join_asof(
         &self,
         other: &RbLazyFrame,
@@ -590,14 +726,16 @@ impl RbLazyFrame {
         tolerance: Option<Wrap<AnyValue<'_>>>,
         tolerance_str: Option<String>,
         coalesce: bool,
+        allow_eq: bool,
+        check_sortedness: bool,
     ) -> RbResult<Self> {
         let coalesce = if coalesce {
             JoinCoalesce::CoalesceColumns
         } 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
@@ -608,19 +746,24 @@ impl RbLazyFrame {
             .allow_parallel(allow_parallel)
             .force_parallel(force_parallel)
             .coalesce(coalesce)
-            .how(JoinType::AsOf(AsOfOptions {
+            .how(JoinType::AsOf(Box::new(AsOfOptions {
                 strategy: strategy.0,
                 left_by: left_by.map(strings_to_pl_smallstr),
                 right_by: right_by.map(strings_to_pl_smallstr),
-                tolerance: tolerance.map(|t| t.0.into_static()),
+                tolerance: tolerance.map(|t| {
+                    let av = t.0.into_static();
+                    let dtype = av.dtype();
+                    Scalar::new(dtype, av)
+                }),
                 tolerance_str: tolerance_str.map(|s| s.into()),
-            }))
+                allow_eq,
+                check_sortedness,
+            })))
             .suffix(suffix)
             .finish()
             .into())
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn join(
         &self,
         other: &RbLazyFrame,
@@ -632,6 +775,7 @@ impl RbLazyFrame {
         how: Wrap<JoinType>,
         suffix: String,
         validate: Wrap<JoinValidation>,
+        maintain_order: Wrap<MaintainOrderJoin>,
         coalesce: Option<bool>,
     ) -> RbResult<Self> {
         let coalesce = match coalesce {
@@ -639,10 +783,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()
@@ -655,38 +799,180 @@ impl RbLazyFrame {
             .how(how.0)
             .validate(validate.0)
             .coalesce(coalesce)
+            .maintain_order(maintain_order.0)
             .suffix(suffix)
             .finish()
             .into())
     }
 
+    pub fn join_where(&self, other: &Self, predicates: RArray, suffix: String) -> RbResult<Self> {
+        let ldf = self.ldf.read().clone();
+        let other = other.ldf.read().clone();
+
+        let predicates = predicates.to_exprs()?;
+
+        Ok(ldf
+            .join_builder()
+            .with(other)
+            .suffix(suffix)
+            .join_where(predicates)
+            .into())
+    }
+
+    pub fn gather(&self, idxs: &Self, null_on_oob: bool) -> Self {
+        let ldf = self.ldf.read().clone();
+        let idxs = idxs.clone().ldf.into_inner();
+        ldf.gather(idxs, null_on_oob).into()
+    }
+
     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 match_to_schema(
+        &self,
+        schema: Wrap<Schema>,
+        missing_columns: Value,
+        missing_struct_fields: Value,
+        extra_columns: Wrap<ExtraColumnsPolicy>,
+        extra_struct_fields: Value,
+        integer_cast: Value,
+        float_cast: Value,
+    ) -> RbResult<Self> {
+        fn parse_missing_columns(
+            schema: &Schema,
+            missing_columns: Value,
+        ) -> RbResult<Vec<MissingColumnsPolicyOrExpr>> {
+            let mut out = Vec::with_capacity(schema.len());
+            if let Ok(policy) = Wrap::<MissingColumnsPolicyOrExpr>::try_convert(missing_columns) {
+                out.extend(std::iter::repeat_n(policy.0, schema.len()));
+            } else if let Ok(dict) = RHash::try_convert(missing_columns) {
+                out.extend(std::iter::repeat_n(
+                    MissingColumnsPolicyOrExpr::Raise,
+                    schema.len(),
+                ));
+                dict.foreach(|key: String, value: Wrap<MissingColumnsPolicyOrExpr>| {
+                    out[schema.try_index_of(&key).map_err(to_rb_err)?] = value.0;
+                    Ok(ForEach::Continue)
+                })?;
+            } else {
+                return Err(RbTypeError::new_err("Invalid value for `missing_columns`"));
+            }
+            Ok(out)
+        }
+        fn parse_missing_struct_fields(
+            schema: &Schema,
+            missing_struct_fields: Value,
+        ) -> RbResult<Vec<MissingColumnsPolicy>> {
+            let mut out = Vec::with_capacity(schema.len());
+            if let Ok(policy) = Wrap::<MissingColumnsPolicy>::try_convert(missing_struct_fields) {
+                out.extend(std::iter::repeat_n(policy.0, schema.len()));
+            } else if let Ok(dict) = RHash::try_convert(missing_struct_fields) {
+                out.extend(std::iter::repeat_n(
+                    MissingColumnsPolicy::Raise,
+                    schema.len(),
+                ));
+                dict.foreach(|key: String, value: Wrap<MissingColumnsPolicy>| {
+                    out[schema.try_index_of(&key).map_err(to_rb_err)?] = value.0;
+                    Ok(ForEach::Continue)
+                })?;
+            } else {
+                return Err(RbTypeError::new_err(
+                    "Invalid value for `missing_struct_fields`",
+                ));
+            }
+            Ok(out)
+        }
+        fn parse_extra_struct_fields(
+            schema: &Schema,
+            extra_struct_fields: Value,
+        ) -> RbResult<Vec<ExtraColumnsPolicy>> {
+            let mut out = Vec::with_capacity(schema.len());
+            if let Ok(policy) = Wrap::<ExtraColumnsPolicy>::try_convert(extra_struct_fields) {
+                out.extend(std::iter::repeat_n(policy.0, schema.len()));
+            } else if let Ok(dict) = RHash::try_convert(extra_struct_fields) {
+                out.extend(std::iter::repeat_n(ExtraColumnsPolicy::Raise, schema.len()));
+                dict.foreach(|key: String, value: Wrap<ExtraColumnsPolicy>| {
+                    out[schema.try_index_of(&key).map_err(to_rb_err)?] = value.0;
+                    Ok(ForEach::Continue)
+                })?;
+            } else {
+                return Err(RbTypeError::new_err(
+                    "Invalid value for `extra_struct_fields`",
+                ));
+            }
+            Ok(out)
+        }
+        fn parse_cast(schema: &Schema, cast: Value) -> RbResult<Vec<UpcastOrForbid>> {
+            let mut out = Vec::with_capacity(schema.len());
+            if let Ok(policy) = Wrap::<UpcastOrForbid>::try_convert(cast) {
+                out.extend(std::iter::repeat_n(policy.0, schema.len()));
+            } else if let Ok(dict) = RHash::try_convert(cast) {
+                out.extend(std::iter::repeat_n(UpcastOrForbid::Forbid, schema.len()));
+                dict.foreach(|key: String, value: Wrap<UpcastOrForbid>| {
+                    out[schema.try_index_of(&key).map_err(to_rb_err)?] = value.0;
+                    Ok(ForEach::Continue)
+                })?;
+            } else {
+                return Err(RbTypeError::new_err(
+                    "Invalid value for `integer_cast` / `float_cast`",
+                ));
+            }
+            Ok(out)
+        }
+
+        let missing_columns = parse_missing_columns(&schema.0, missing_columns)?;
+        let missing_struct_fields = parse_missing_struct_fields(&schema.0, missing_struct_fields)?;
+        let extra_struct_fields = parse_extra_struct_fields(&schema.0, extra_struct_fields)?;
+        let integer_cast = parse_cast(&schema.0, integer_cast)?;
+        let float_cast = parse_cast(&schema.0, float_cast)?;
+
+        let per_column = (0..schema.0.len())
+            .map(|i| MatchToSchemaPerColumn {
+                missing_columns: missing_columns[i].clone(),
+                missing_struct_fields: missing_struct_fields[i],
+                extra_struct_fields: extra_struct_fields[i],
+                integer_cast: integer_cast[i],
+                float_cast: float_cast[i],
+            })
+            .collect();
+
+        let ldf = self.ldf.read().clone();
+        Ok(ldf
+            .match_to_schema(Arc::new(schema.0), per_column, extra_columns.0)
+            .into())
+    }
+
+    pub fn pipe_with_schema(&self, callback: Value) -> Self {
+        let ldf = self.ldf.read().clone();
+        let function = RubyObject::from(callback);
+        ldf.pipe_with_schema(PlanCallback::new_ruby(function))
+            .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()),
@@ -695,76 +981,88 @@ 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, column: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        let column = rb_exprs_to_exprs(column)?;
-        Ok(ldf.explode(column).into())
+    pub fn explode(&self, subset: &RbSelector, empty_as_null: bool, keep_nulls: bool) -> Self {
+        self.ldf
+            .read()
+            .clone()
+            .explode(
+                subset.inner.clone(),
+                ExplodeOptions {
+                    empty_as_null,
+                    keep_nulls,
+                },
+            )
+            .into()
     }
 
     pub fn null_count(&self) -> Self {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
         ldf.null_count().into()
     }
 
     pub fn unique(
         &self,
         maintain_order: bool,
-        subset: Option<Vec<String>>,
+        subset: Option<RArray>,
         keep: Wrap<UniqueKeepStrategy>,
     ) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
+        let ldf = self.ldf.read().clone();
+        let subset = match subset {
+            Some(e) => Some(e.to_exprs()?),
+            None => None,
+        };
         Ok(match maintain_order {
             true => ldf.unique_stable_generic(subset, keep.0),
             false => ldf.unique_generic(subset, keep.0),
@@ -772,50 +1070,110 @@ impl RbLazyFrame {
         .into())
     }
 
-    pub fn drop_nulls(&self, subset: Option<Vec<String>>) -> Self {
-        let ldf = self.ldf.borrow().clone();
-        ldf.drop_nulls(subset.map(|v| v.into_iter().map(|s| col(&s)).collect()))
+    pub fn drop_nans(&self, subset: Option<&RbSelector>) -> Self {
+        self.ldf
+            .read()
+            .clone()
+            .drop_nans(subset.map(|e| e.inner.clone()))
+            .into()
+    }
+
+    pub fn drop_nulls(&self, subset: Option<&RbSelector>) -> Self {
+        self.ldf
+            .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()
     }
 
+    pub fn pivot(
+        &self,
+        on: &RbSelector,
+        on_columns: &RbDataFrame,
+        index: &RbSelector,
+        values: &RbSelector,
+        agg: &RbExpr,
+        maintain_order: bool,
+        separator: String,
+        column_naming: Wrap<PivotColumnNaming>,
+    ) -> Self {
+        let ldf = self.ldf.read().clone();
+        ldf.pivot(
+            on.inner.clone(),
+            Arc::new(on_columns.df.read().clone()),
+            index.inner.clone(),
+            values.inner.clone(),
+            agg.inner.clone(),
+            maintain_order,
+            separator.into(),
+            column_naming.0,
+        )
+        .into()
+    }
+
     pub fn unpivot(
         &self,
-        on: RArray,
-        index: RArray,
+        on: Option<&RbSelector>,
+        index: &RbSelector,
         value_name: Option<String>,
         variable_name: Option<String>,
     ) -> RbResult<Self> {
-        let on = rb_exprs_to_exprs(on)?;
-        let index = rb_exprs_to_exprs(index)?;
         let args = UnpivotArgsDSL {
-            on: on.into_iter().map(|e| e.into()).collect(),
-            index: index.into_iter().map(|e| e.into()).collect(),
+            on: on.map(|on| on.inner.clone()),
+            index: index.inner.clone(),
             value_name: value_name.map(|s| s.into()),
             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, cols: Vec<String>) -> Self {
-        let ldf = self.ldf.borrow().clone();
-        ldf.drop(cols).into()
+    pub fn map_batches(
+        &self,
+        function: Value,
+        predicate_pushdown: bool,
+        projection_pushdown: bool,
+        slice_pushdown: bool,
+        streamable: bool,
+        schema: Option<Wrap<Schema>>,
+        validate_output: bool,
+    ) -> Self {
+        let mut opt = OptFlags::default();
+        opt.set(OptFlags::PREDICATE_PUSHDOWN, predicate_pushdown);
+        opt.set(OptFlags::PROJECTION_PUSHDOWN, projection_pushdown);
+        opt.set(OptFlags::SLICE_PUSHDOWN, slice_pushdown);
+        opt.set(OptFlags::STREAMING, streamable);
+
+        self.ldf
+            .read()
+            .clone()
+            .map_ruby(
+                function.into(),
+                opt,
+                schema.map(|s| Arc::new(s.0)),
+                validate_output,
+            )
+            .into()
+    }
+
+    pub fn drop(&self, columns: &RbSelector) -> Self {
+        self.ldf.read().clone().drop(columns.inner.clone()).into()
     }
 
     pub fn cast(&self, rb_dtypes: RHash, strict: bool) -> RbResult<Self> {
@@ -826,52 +1184,172 @@ 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(&self) -> RbResult<RHash> {
-        let schema = 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 = RHash::new();
-        schema.iter_fields().for_each(|fld| {
-            schema_dict
-                .aset::<String, Value>(
-                    fld.name().to_string(),
-                    Wrap(fld.dtype().clone()).into_value(),
-                )
-                .unwrap();
-        });
+        let schema_dict = rb.hash_new();
+        for fld in schema.iter_fields() {
+            schema_dict.aset(
+                fld.name().to_string(),
+                Wrap(fld.dtype().clone()).try_into_value_with(rb)?,
+            )?;
+        }
         Ok(schema_dict)
     }
 
-    pub fn unnest(&self, cols: Vec<String>) -> Self {
-        self.ldf.borrow().clone().unnest(cols).into()
+    pub fn unnest(&self, columns: &RbSelector, separator: Option<String>) -> Self {
+        self.ldf
+            .read()
+            .clone()
+            .unnest(
+                columns.inner.clone(),
+                separator.as_deref().map(PlSmallStr::from_str),
+            )
+            .into()
     }
 
     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> {
+    pub fn merge_sorted(&self, other: &Self, key: String, maintain_order: bool) -> RbResult<Self> {
         let out = self
             .ldf
-            .borrow()
+            .read()
             .clone()
-            .merge_sorted(other.ldf.borrow().clone(), &key)
+            .merge_sorted(other.ldf.read().clone(), &key, maintain_order)
             .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: Some(false),
+                nulls_last: Some(false),
+            })
+            .collect::<Vec<_>>();
+
+        if !columns.is_empty() {
+            if descending.len() != 1 {
+                sorted
+                    .iter_mut()
+                    .zip(descending)
+                    .for_each(|(s, d)| s.descending = Some(d));
+            } else if descending[0] {
+                sorted.iter_mut().for_each(|s| s.descending = Some(true));
+            }
+
+            if nulls_last.len() != 1 {
+                sorted
+                    .iter_mut()
+                    .zip(nulls_last)
+                    .for_each(|(s, d)| s.nulls_last = Some(d));
+            } else if nulls_last[0] {
+                sorted.iter_mut().for_each(|s| s.nulls_last = Some(true));
+            }
+        }
+
+        let out = self
+            .ldf
+            .read()
+            .clone()
+            .hint(HintIR::Sorted(sorted.into()))
+            .map_err(RbPolarsErr::from)?;
+        Ok(out.into())
+    }
+}
+
+#[magnus::wrap(class = "Polars::RbCollectBatches")]
+pub struct RbCollectBatches {
+    inner: Arc<Mutex<CollectBatches>>,
+    ldf: LazyFrame,
+}
+
+impl RbCollectBatches {
+    pub fn next(rb: &Ruby, slf: &Self) -> RbResult<Option<RbDataFrame>> {
+        let inner = Arc::clone(&slf.inner);
+        rb.enter_polars(|| PolarsResult::Ok(inner.lock().next().transpose()?.map(RbDataFrame::new)))
+    }
+
+    pub fn __arrow_c_stream__(rb: &Ruby, self_: &Self) -> RbResult<Value> {
+        let mut ldf = self_.ldf.clone();
+        let schema = ldf
+            .collect_schema()
+            .map_err(RbPolarsErr::from)?
+            .to_arrow(CompatLevel::newest());
+
+        let dtype = ArrowDataType::Struct(schema.into_iter_values().collect());
+
+        let iter = Box::new(ArrowStreamIterator::new(self_.inner.clone(), dtype.clone()));
+        let field = ArrowField::new(PlSmallStr::EMPTY, dtype, false);
+        let stream = export_iterator(iter, field);
+        Ok(RbArrowArrayStream { stream }.into_value_with(rb))
+    }
+}
+
+pub struct ArrowStreamIterator {
+    inner: Arc<Mutex<CollectBatches>>,
+    dtype: ArrowDataType,
+}
+
+impl ArrowStreamIterator {
+    fn new(inner: Arc<Mutex<CollectBatches>>, schema: ArrowDataType) -> Self {
+        Self {
+            inner,
+            dtype: schema,
+        }
+    }
+}
+
+impl Iterator for ArrowStreamIterator {
+    type Item = PolarsResult<ArrayRef>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let next = self.inner.lock().next();
+        match next {
+            None => None,
+            Some(Err(err)) => Some(Err(err)),
+            Some(Ok(df)) => {
+                let height = df.height();
+                let arrays = df.rechunk_into_arrow(CompatLevel::newest());
+                Some(Ok(Box::new(arrow::array::StructArray::new(
+                    self.dtype.clone(),
+                    height,
+                    arrays,
+                    None,
+                ))))
+            }
+        }
+    }
 }