polars-df 0.15.0 → 0.26.0

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

@@ -1,139 +1,136 @@
+use std::hash::BuildHasher;
+
+use arrow::bitmap::MutableBitmap;
 use either::Either;
-use magnus::{prelude::*, typed_data::Obj, IntoValue, RArray, Value};
-use polars::frame::NullStrategy;
-use polars::prelude::pivot::{pivot, pivot_stable};
+use magnus::{RArray, Ruby, Value, prelude::*, value::Opaque};
 use polars::prelude::*;
 
 use crate::conversion::*;
-use crate::map::dataframe::{
-    apply_lambda_unknown, apply_lambda_with_bool_out_type, apply_lambda_with_primitive_out_type,
-    apply_lambda_with_utf8_out_type,
-};
 use crate::prelude::strings_to_pl_smallstr;
-use crate::series::{to_rbseries, to_series};
-use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
+use crate::rb_modules::pl_utils;
+use crate::ruby::exceptions::RbIndexError;
+use crate::ruby::gvl::GvlExt;
+use crate::ruby::utils::TryIntoValue;
+use crate::series::ToRbSeries;
+use crate::series::to_series;
+use crate::utils::EnterPolarsExt;
+use crate::{RbDataFrame, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
 
 impl RbDataFrame {
     pub fn init(columns: RArray) -> RbResult<Self> {
         let mut cols = Vec::new();
         for i in columns.into_iter() {
-            cols.push(<&RbSeries>::try_convert(i)?.series.borrow().clone().into());
+            cols.push(<&RbSeries>::try_convert(i)?.series.read().clone().into());
         }
-        let df = DataFrame::new(cols).map_err(RbPolarsErr::from)?;
+        let df = DataFrame::new_infer_height(cols).map_err(RbPolarsErr::from)?;
         Ok(RbDataFrame::new(df))
     }
 
     pub fn estimated_size(&self) -> usize {
-        self.df.borrow().estimated_size()
+        self.df.read().estimated_size()
     }
 
     pub fn dtype_strings(&self) -> Vec<String> {
         self.df
-            .borrow()
-            .get_columns()
+            .read()
+            .columns()
             .iter()
             .map(|s| format!("{}", s.dtype()))
             .collect()
     }
 
-    pub fn add(&self, s: &RbSeries) -> RbResult<Self> {
-        let df = (&*self.df.borrow() + &*s.series.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn add(rb: &Ruby, self_: &Self, s: &RbSeries) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() + &*s.series.read())
     }
 
-    pub fn sub(&self, s: &RbSeries) -> RbResult<Self> {
-        let df = (&*self.df.borrow() - &*s.series.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn sub(rb: &Ruby, self_: &Self, s: &RbSeries) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() - &*s.series.read())
     }
 
-    pub fn div(&self, s: &RbSeries) -> RbResult<Self> {
-        let df = (&*self.df.borrow() / &*s.series.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn div(rb: &Ruby, self_: &Self, s: &RbSeries) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() / &*s.series.read())
     }
 
-    pub fn mul(&self, s: &RbSeries) -> RbResult<Self> {
-        let df = (&*self.df.borrow() * &*s.series.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn mul(rb: &Ruby, self_: &Self, s: &RbSeries) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() * &*s.series.read())
     }
 
-    pub fn rem(&self, s: &RbSeries) -> RbResult<Self> {
-        let df = (&*self.df.borrow() % &*s.series.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn rem(rb: &Ruby, self_: &Self, s: &RbSeries) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() % &*s.series.read())
     }
 
-    pub fn add_df(&self, s: &Self) -> RbResult<Self> {
-        let df = (&*self.df.borrow() + &*s.df.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn add_df(rb: &Ruby, self_: &Self, s: &Self) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() + &*s.df.read())
     }
 
-    pub fn sub_df(&self, s: &Self) -> RbResult<Self> {
-        let df = (&*self.df.borrow() - &*s.df.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn sub_df(rb: &Ruby, self_: &Self, s: &Self) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() - &*s.df.read())
     }
 
-    pub fn div_df(&self, s: &Self) -> RbResult<Self> {
-        let df = (&*self.df.borrow() / &*s.df.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn div_df(rb: &Ruby, self_: &Self, s: &Self) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() / &*s.df.read())
     }
 
-    pub fn mul_df(&self, s: &Self) -> RbResult<Self> {
-        let df = (&*self.df.borrow() * &*s.df.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn mul_df(rb: &Ruby, self_: &Self, s: &Self) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() * &*s.df.read())
     }
 
-    pub fn rem_df(&self, s: &Self) -> RbResult<Self> {
-        let df = (&*self.df.borrow() % &*s.df.borrow()).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn rem_df(rb: &Ruby, self_: &Self, s: &Self) -> RbResult<Self> {
+        rb.enter_polars_df(|| &*self_.df.read() % &*s.df.read())
     }
 
     pub fn sample_n(
-        &self,
+        rb: &Ruby,
+        self_: &Self,
         n: &RbSeries,
         with_replacement: bool,
         shuffle: bool,
         seed: Option<u64>,
     ) -> RbResult<Self> {
-        let df = self
-            .df
-            .borrow()
-            .sample_n(&n.series.borrow(), with_replacement, shuffle, seed)
-            .map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+        rb.enter_polars_df(|| {
+            self_
+                .df
+                .read()
+                .sample_n(&n.series.read(), with_replacement, shuffle, seed)
+        })
     }
 
     pub fn sample_frac(
-        &self,
+        rb: &Ruby,
+        self_: &Self,
         frac: &RbSeries,
         with_replacement: bool,
         shuffle: bool,
         seed: Option<u64>,
     ) -> RbResult<Self> {
-        let df = self
-            .df
-            .borrow()
-            .sample_frac(&frac.series.borrow(), with_replacement, shuffle, seed)
-            .map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+        rb.enter_polars_df(|| {
+            self_
+                .df
+                .read()
+                .sample_frac(&frac.series.read(), with_replacement, shuffle, seed)
+        })
     }
 
-    pub fn rechunk(&self) -> Self {
-        let mut df = self.df.borrow_mut().clone();
-        df.as_single_chunk_par();
-        df.into()
+    pub fn rechunk(rb: &Ruby, self_: &Self) -> RbResult<Self> {
+        rb.enter_polars_df(|| {
+            let mut df = self_.df.write().clone();
+            df.rechunk_mut_par();
+            Ok(df)
+        })
     }
 
     pub fn as_str(&self) -> String {
-        format!("{}", self.df.borrow())
+        format!("{}", self.df.read())
     }
 
-    pub fn get_columns(&self) -> RArray {
-        let cols = self.df.borrow().get_columns().to_vec();
-        to_rbseries(cols)
+    pub fn get_columns(rb: &Ruby, self_: &Self) -> RArray {
+        let cols = self_.df.read().columns().to_vec();
+        cols.to_rbseries(rb)
     }
 
     pub fn columns(&self) -> Vec<String> {
         self.df
-            .borrow()
+            .read()
             .get_column_names()
             .iter()
             .map(|v| v.to_string())
@@ -142,415 +139,334 @@ impl RbDataFrame {
 
     pub fn set_column_names(&self, names: Vec<String>) -> RbResult<()> {
         self.df
-            .borrow_mut()
+            .write()
             .set_column_names(&names)
             .map_err(RbPolarsErr::from)?;
         Ok(())
     }
 
-    pub fn dtypes(&self) -> RArray {
-        RArray::from_iter(
-            self.df
-                .borrow()
-                .iter()
-                .map(|s| Wrap(s.dtype().clone()).into_value()),
-        )
+    pub fn dtypes(ruby: &Ruby, self_: &Self) -> RbResult<RArray> {
+        let df = self_.df.read();
+        let iter = df
+            .columns()
+            .iter()
+            .map(|s| Wrap(s.dtype().clone()).try_into_value_with(ruby));
+        ruby.ary_try_from_iter(iter)
     }
 
     pub fn n_chunks(&self) -> usize {
-        self.df.borrow().n_chunks()
+        self.df.read().first_col_n_chunks()
     }
 
     pub fn shape(&self) -> (usize, usize) {
-        self.df.borrow().shape()
+        self.df.read().shape()
     }
 
     pub fn height(&self) -> usize {
-        self.df.borrow().height()
+        self.df.read().height()
     }
 
     pub fn width(&self) -> usize {
-        self.df.borrow().width()
+        self.df.read().width()
     }
 
-    pub fn hstack(&self, columns: RArray) -> RbResult<Self> {
+    pub fn hstack(rb: &Ruby, self_: &Self, columns: RArray) -> RbResult<Self> {
         let columns = to_series(columns)?;
         let columns = columns.into_iter().map(Into::into).collect::<Vec<_>>();
-        let df = self
-            .df
-            .borrow()
-            .hstack(&columns)
-            .map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+        rb.enter_polars_df(|| self_.df.read().hstack(&columns))
     }
 
-    pub fn hstack_mut(&self, columns: RArray) -> RbResult<()> {
+    pub fn hstack_mut(rb: &Ruby, self_: &Self, columns: RArray) -> RbResult<()> {
         let columns = to_series(columns)?;
         let columns = columns.into_iter().map(Into::into).collect::<Vec<_>>();
-        self.df
-            .borrow_mut()
-            .hstack_mut(&columns)
-            .map_err(RbPolarsErr::from)?;
+        rb.enter_polars(|| self_.df.write().hstack_mut(&columns).map(drop))?;
         Ok(())
     }
 
-    pub fn vstack(&self, df: &RbDataFrame) -> RbResult<Self> {
-        let df = self
-            .df
-            .borrow()
-            .vstack(&df.df.borrow())
-            .map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn vstack(rb: &Ruby, self_: &Self, other: &RbDataFrame) -> RbResult<Self> {
+        rb.enter_polars_df(|| self_.df.read().vstack(&other.df.read()))
     }
 
-    pub fn vstack_mut(&self, df: &RbDataFrame) -> RbResult<()> {
-        self.df
-            .borrow_mut()
-            .vstack_mut(&df.df.borrow())
-            .map_err(RbPolarsErr::from)?;
+    pub fn vstack_mut(rb: &Ruby, self_: &Self, other: &RbDataFrame) -> RbResult<()> {
+        rb.enter_polars(|| {
+            // Prevent self-vstack deadlocks.
+            let other = other.df.read().clone();
+            self_.df.write().vstack_mut(&other)?;
+            PolarsResult::Ok(())
+        })?;
         Ok(())
     }
 
-    pub fn extend(&self, df: &RbDataFrame) -> RbResult<()> {
-        self.df
-            .borrow_mut()
-            .extend(&df.df.borrow())
-            .map_err(RbPolarsErr::from)?;
+    pub fn extend(rb: &Ruby, self_: &Self, other: &RbDataFrame) -> RbResult<()> {
+        rb.enter_polars(|| {
+            // Prevent self-extend deadlocks.
+            let other = other.df.read().clone();
+            self_.df.write().extend(&other)
+        })?;
         Ok(())
     }
 
     pub fn drop_in_place(&self, name: String) -> RbResult<RbSeries> {
         let s = self
             .df
-            .borrow_mut()
+            .write()
             .drop_in_place(&name)
             .map_err(RbPolarsErr::from)?;
         let s = s.take_materialized_series();
         Ok(RbSeries::new(s))
     }
 
-    pub fn select_at_idx(&self, idx: usize) -> Option<RbSeries> {
-        self.df
-            .borrow()
-            .select_at_idx(idx)
-            .map(|s| RbSeries::new(s.as_materialized_series().clone()))
+    pub fn to_series(&self, index: isize) -> RbResult<RbSeries> {
+        let df = &self.df.read();
+
+        let index_adjusted = if index < 0 {
+            df.width().checked_sub(index.unsigned_abs())
+        } else {
+            Some(usize::try_from(index).unwrap())
+        };
+
+        let s = index_adjusted.and_then(|i| df.select_at_idx(i));
+        match s {
+            Some(s) => Ok(RbSeries::new(s.as_materialized_series().clone())),
+            None => Err(RbIndexError::new_err(
+                polars_err!(oob = index, df.width()).to_string(),
+            )),
+        }
     }
 
     pub fn get_column_index(&self, name: String) -> Option<usize> {
-        self.df.borrow().get_column_index(&name)
+        self.df.read().get_column_index(&name)
     }
 
     pub fn get_column(&self, name: String) -> RbResult<RbSeries> {
         let series = self
             .df
-            .borrow()
+            .read()
             .column(&name)
             .map(|s| RbSeries::new(s.as_materialized_series().clone()))
             .map_err(RbPolarsErr::from)?;
         Ok(series)
     }
 
-    pub fn select(&self, selection: Vec<String>) -> RbResult<Self> {
-        let df = self
-            .df
-            .borrow()
-            .select(selection)
-            .map_err(RbPolarsErr::from)?;
-        Ok(RbDataFrame::new(df))
+    pub fn select(rb: &Ruby, self_: &Self, columns: Vec<String>) -> RbResult<Self> {
+        rb.enter_polars_df(|| self_.df.read().select(columns.iter().map(|x| &**x)))
     }
 
-    pub fn gather(&self, indices: Vec<IdxSize>) -> RbResult<Self> {
+    pub fn gather(rb: &Ruby, self_: &Self, indices: Vec<IdxSize>) -> RbResult<Self> {
         let indices = IdxCa::from_vec("".into(), indices);
-        let df = self.df.borrow().take(&indices).map_err(RbPolarsErr::from)?;
-        Ok(RbDataFrame::new(df))
+        rb.enter_polars_df(|| self_.df.read().take(&indices))
     }
 
-    pub fn take_with_series(&self, indices: &RbSeries) -> RbResult<Self> {
-        let binding = indices.series.borrow();
-        let idx = binding.idx().map_err(RbPolarsErr::from)?;
-        let df = self.df.borrow().take(idx).map_err(RbPolarsErr::from)?;
-        Ok(RbDataFrame::new(df))
+    pub fn gather_with_series(rb: &Ruby, self_: &Self, indices: &RbSeries) -> RbResult<Self> {
+        let idx_s = indices.series.read();
+        let indices = idx_s.idx().map_err(RbPolarsErr::from)?;
+        rb.enter_polars_df(|| self_.df.read().take(indices))
     }
 
     pub fn replace(&self, column: String, new_col: &RbSeries) -> RbResult<()> {
         self.df
-            .borrow_mut()
-            .replace(&column, new_col.series.borrow().clone())
+            .write()
+            .replace(&column, new_col.clone().series.into_inner().into_column())
             .map_err(RbPolarsErr::from)?;
         Ok(())
     }
 
     pub fn replace_column(&self, index: usize, new_col: &RbSeries) -> RbResult<()> {
         self.df
-            .borrow_mut()
-            .replace_column(index, new_col.series.borrow().clone())
+            .write()
+            .replace_column(index, new_col.clone().series.into_inner().into_column())
             .map_err(RbPolarsErr::from)?;
         Ok(())
     }
 
     pub fn insert_column(&self, index: usize, new_col: &RbSeries) -> RbResult<()> {
         self.df
-            .borrow_mut()
-            .insert_column(index, new_col.series.borrow().clone())
+            .write()
+            .insert_column(index, new_col.clone().series.into_inner().into_column())
             .map_err(RbPolarsErr::from)?;
         Ok(())
     }
 
-    pub fn slice(&self, offset: usize, length: Option<usize>) -> Self {
-        let df = self.df.borrow().slice(
-            offset as i64,
-            length.unwrap_or_else(|| self.df.borrow().height()),
-        );
-        df.into()
+    pub fn slice(rb: &Ruby, self_: &Self, offset: i64, length: Option<usize>) -> RbResult<Self> {
+        rb.enter_polars_df(|| {
+            let df = self_.df.read();
+            Ok(df.slice(offset, length.unwrap_or_else(|| df.height())))
+        })
     }
 
-    pub fn head(&self, length: Option<usize>) -> Self {
-        self.df.borrow().head(length).into()
+    pub fn head(rb: &Ruby, self_: &Self, n: usize) -> RbResult<Self> {
+        rb.enter_polars_df(|| Ok(self_.df.read().head(Some(n))))
     }
 
-    pub fn tail(&self, length: Option<usize>) -> Self {
-        self.df.borrow().tail(length).into()
+    pub fn tail(rb: &Ruby, self_: &Self, n: usize) -> RbResult<Self> {
+        rb.enter_polars_df(|| Ok(self_.df.read().tail(Some(n))))
     }
 
-    pub fn is_unique(&self) -> RbResult<RbSeries> {
-        let mask = self.df.borrow().is_unique().map_err(RbPolarsErr::from)?;
-        Ok(mask.into_series().into())
+    pub fn is_unique(rb: &Ruby, self_: &Self) -> RbResult<RbSeries> {
+        rb.enter_polars_series(|| self_.df.read().is_unique())
     }
 
-    pub fn is_duplicated(&self) -> RbResult<RbSeries> {
-        let mask = self
-            .df
-            .borrow()
-            .is_duplicated()
-            .map_err(RbPolarsErr::from)?;
-        Ok(mask.into_series().into())
+    pub fn is_duplicated(rb: &Ruby, self_: &Self) -> RbResult<RbSeries> {
+        rb.enter_polars_series(|| self_.df.read().is_duplicated())
     }
 
-    pub fn equals(&self, other: &RbDataFrame, null_equal: bool) -> bool {
+    pub fn equals(
+        rb: &Ruby,
+        self_: &Self,
+        other: &RbDataFrame,
+        null_equal: bool,
+    ) -> RbResult<bool> {
         if null_equal {
-            self.df.borrow().equals_missing(&other.df.borrow())
+            rb.enter_polars_ok(|| self_.df.read().equals_missing(&other.df.read()))
         } else {
-            self.df.borrow().equals(&other.df.borrow())
+            rb.enter_polars_ok(|| self_.df.read().equals(&other.df.read()))
         }
     }
 
-    pub fn with_row_index(&self, name: String, offset: Option<IdxSize>) -> RbResult<Self> {
-        let df = self
-            .df
-            .borrow()
-            .with_row_index(name.into(), offset)
-            .map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+    pub fn with_row_index(
+        rb: &Ruby,
+        self_: &Self,
+        name: String,
+        offset: Option<IdxSize>,
+    ) -> RbResult<Self> {
+        rb.enter_polars_df(|| self_.df.read().with_row_index(name.into(), offset))
+    }
+
+    pub fn group_by_map_groups(
+        rb: &Ruby,
+        self_: &Self,
+        by: Vec<String>,
+        lambda: Value,
+        maintain_order: bool,
+    ) -> RbResult<Self> {
+        let lambda = Opaque::from(lambda);
+        rb.enter_polars_df(|| {
+            let df = self_.df.read().clone(); // Clone so we can't deadlock on re-entrance from lambda.
+            let gb = if maintain_order {
+                df.group_by_stable(by.iter().map(|x| &**x))
+            } else {
+                df.group_by(by.iter().map(|x| &**x))
+            }?;
+
+            let function = move |df: DataFrame| {
+                Ruby::attach(|rb| {
+                    let lambda = rb.get_inner(lambda);
+
+                    let rbpolars = pl_utils(rb);
+                    let rbdf = RbDataFrame::new(df);
+                    let ruby_df_wrapper: Value =
+                        rbpolars.funcall("wrap_df", (rbdf,)).unwrap();
+
+                    // Call the lambda and get a Ruby-side DataFrame wrapper.
+                    let result_df_wrapper: Value = match lambda.funcall("call", (ruby_df_wrapper,)) {
+                        Ok(rbobj) => rbobj,
+                        Err(e) => panic!("UDF failed: {}", e),
+                    };
+                    let rbdf: &RbDataFrame = result_df_wrapper.funcall("_df", ()).expect(
+                        "Could not get DataFrame attribute '_df'. Make sure that you return a DataFrame object.",
+                    );
+
+                    Ok(rbdf.clone().df.into_inner())
+                })
+            };
+
+            gb.apply(function)
+        })
     }
 
     pub fn clone(&self) -> Self {
-        RbDataFrame::new(self.df.borrow().clone())
+        Clone::clone(self)
     }
 
     pub fn unpivot(
-        &self,
-        on: Vec<String>,
+        rb: &Ruby,
+        self_: &Self,
+        on: Option<Vec<String>>,
         index: Vec<String>,
         value_name: Option<String>,
         variable_name: Option<String>,
     ) -> RbResult<Self> {
-        let args = UnpivotArgsIR {
-            on: strings_to_pl_smallstr(on),
-            index: strings_to_pl_smallstr(index),
-            value_name: value_name.map(|s| s.into()),
-            variable_name: variable_name.map(|s| s.into()),
-        };
+        use polars_ops::unpivot::UnpivotDF;
+        let args = UnpivotArgsIR::new(
+            self_.df.read().get_column_names_owned(),
+            on.map(strings_to_pl_smallstr),
+            strings_to_pl_smallstr(index),
+            value_name.map(|s| s.into()),
+            variable_name.map(|s| s.into()),
+        );
 
-        let df = self.df.borrow().unpivot2(args).map_err(RbPolarsErr::from)?;
-        Ok(RbDataFrame::new(df))
-    }
-
-    #[allow(clippy::too_many_arguments)]
-    pub fn pivot_expr(
-        &self,
-        on: Vec<String>,
-        index: Option<Vec<String>>,
-        values: Option<Vec<String>>,
-        maintain_order: bool,
-        sort_columns: bool,
-        aggregate_expr: Option<&RbExpr>,
-        separator: Option<String>,
-    ) -> RbResult<Self> {
-        let fun = if maintain_order { pivot_stable } else { pivot };
-        let agg_expr = aggregate_expr.map(|aggregate_expr| aggregate_expr.inner.clone());
-        let df = fun(
-            &self.df.borrow(),
-            on,
-            index,
-            values,
-            sort_columns,
-            agg_expr,
-            separator.as_deref(),
-        )
-        .map_err(RbPolarsErr::from)?;
-        Ok(RbDataFrame::new(df))
+        rb.enter_polars_df(|| self_.df.read().unpivot2(args))
     }
 
     pub fn partition_by(
-        &self,
+        rb: &Ruby,
+        self_: &Self,
         by: Vec<String>,
         maintain_order: bool,
         include_key: bool,
     ) -> RbResult<RArray> {
         let out = if maintain_order {
-            self.df.borrow().partition_by_stable(by, include_key)
+            self_.df.read().partition_by_stable(by, include_key)
         } else {
-            self.df.borrow().partition_by(by, include_key)
+            self_.df.read().partition_by(by, include_key)
         }
         .map_err(RbPolarsErr::from)?;
-        Ok(RArray::from_iter(out.into_iter().map(RbDataFrame::new)))
+        Ok(rb.ary_from_iter(out.into_iter().map(RbDataFrame::new)))
     }
 
     pub fn lazy(&self) -> RbLazyFrame {
-        self.df.borrow().clone().lazy().into()
-    }
-
-    pub fn max_horizontal(&self) -> RbResult<Option<RbSeries>> {
-        let s = self
-            .df
-            .borrow()
-            .max_horizontal()
-            .map_err(RbPolarsErr::from)?;
-        Ok(s.map(|s| s.take_materialized_series().into()))
-    }
-
-    pub fn min_horizontal(&self) -> RbResult<Option<RbSeries>> {
-        let s = self
-            .df
-            .borrow()
-            .min_horizontal()
-            .map_err(RbPolarsErr::from)?;
-        Ok(s.map(|s| s.take_materialized_series().into()))
-    }
-
-    pub fn sum_horizontal(&self, ignore_nulls: bool) -> RbResult<Option<RbSeries>> {
-        let null_strategy = if ignore_nulls {
-            NullStrategy::Ignore
-        } else {
-            NullStrategy::Propagate
-        };
-        let s = self
-            .df
-            .borrow()
-            .sum_horizontal(null_strategy)
-            .map_err(RbPolarsErr::from)?;
-        Ok(s.map(|s| s.into()))
-    }
-
-    pub fn mean_horizontal(&self, ignore_nulls: bool) -> RbResult<Option<RbSeries>> {
-        let null_strategy = if ignore_nulls {
-            NullStrategy::Ignore
-        } else {
-            NullStrategy::Propagate
-        };
-        let s = self
-            .df
-            .borrow()
-            .mean_horizontal(null_strategy)
-            .map_err(RbPolarsErr::from)?;
-        Ok(s.map(|s| s.into()))
+        self.df.read().clone().lazy().into()
     }
 
     pub fn to_dummies(
-        &self,
+        rb: &Ruby,
+        self_: &Self,
         columns: Option<Vec<String>>,
         separator: Option<String>,
         drop_first: bool,
+        drop_nulls: bool,
     ) -> RbResult<Self> {
-        let df = match columns {
-            Some(cols) => self.df.borrow().columns_to_dummies(
+        rb.enter_polars_df(|| match columns {
+            Some(cols) => self_.df.read().columns_to_dummies(
                 cols.iter().map(|x| x as &str).collect(),
                 separator.as_deref(),
                 drop_first,
+                drop_nulls,
             ),
-            None => self
+            None => self_
                 .df
-                .borrow()
-                .to_dummies(separator.as_deref(), drop_first),
-        }
-        .map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+                .read()
+                .to_dummies(separator.as_deref(), drop_first, drop_nulls),
+        })
     }
 
-    pub fn null_count(&self) -> Self {
-        let df = self.df.borrow().null_count();
-        df.into()
+    pub fn null_count(rb: &Ruby, self_: &Self) -> RbResult<Self> {
+        rb.enter_polars_df(|| Ok(self_.df.read().null_count()))
     }
 
-    pub fn map_rows(
-        &self,
-        lambda: Value,
-        output_type: Option<Wrap<DataType>>,
-        inference_size: usize,
-    ) -> RbResult<(Value, bool)> {
-        let df = &self.df.borrow();
-
-        let output_type = output_type.map(|dt| dt.0);
-        let out = match output_type {
-            Some(DataType::Int32) => {
-                apply_lambda_with_primitive_out_type::<Int32Type>(df, lambda, 0, None).into_series()
-            }
-            Some(DataType::Int64) => {
-                apply_lambda_with_primitive_out_type::<Int64Type>(df, lambda, 0, None).into_series()
-            }
-            Some(DataType::UInt32) => {
-                apply_lambda_with_primitive_out_type::<UInt32Type>(df, lambda, 0, None)
-                    .into_series()
-            }
-            Some(DataType::UInt64) => {
-                apply_lambda_with_primitive_out_type::<UInt64Type>(df, lambda, 0, None)
-                    .into_series()
-            }
-            Some(DataType::Float32) => {
-                apply_lambda_with_primitive_out_type::<Float32Type>(df, lambda, 0, None)
-                    .into_series()
-            }
-            Some(DataType::Float64) => {
-                apply_lambda_with_primitive_out_type::<Float64Type>(df, lambda, 0, None)
-                    .into_series()
-            }
-            Some(DataType::Boolean) => {
-                apply_lambda_with_bool_out_type(df, lambda, 0, None).into_series()
-            }
-            Some(DataType::Date) => {
-                apply_lambda_with_primitive_out_type::<Int32Type>(df, lambda, 0, None)
-                    .into_date()
-                    .into_series()
-            }
-            Some(DataType::Datetime(tu, tz)) => {
-                apply_lambda_with_primitive_out_type::<Int64Type>(df, lambda, 0, None)
-                    .into_datetime(tu, tz)
-                    .into_series()
-            }
-            Some(DataType::String) => {
-                apply_lambda_with_utf8_out_type(df, lambda, 0, None).into_series()
-            }
-            _ => return apply_lambda_unknown(df, lambda, inference_size),
-        };
-
-        Ok((Obj::wrap(RbSeries::from(out)).as_value(), false))
-    }
-
-    pub fn shrink_to_fit(&self) {
-        self.df.borrow_mut().shrink_to_fit();
+    pub fn shrink_to_fit(rb: &Ruby, self_: &Self) -> RbResult<()> {
+        rb.enter_polars_ok(|| self_.df.write().shrink_to_fit())
     }
 
-    pub fn hash_rows(&self, k0: u64, k1: u64, k2: u64, k3: u64) -> RbResult<RbSeries> {
-        let hb = ahash::RandomState::with_seeds(k0, k1, k2, k3);
-        let hash = self
-            .df
-            .borrow_mut()
-            .hash_rows(Some(hb))
-            .map_err(RbPolarsErr::from)?;
-        Ok(hash.into_series().into())
+    pub fn hash_rows(
+        rb: &Ruby,
+        self_: &Self,
+        k0: u64,
+        k1: u64,
+        k2: u64,
+        k3: u64,
+    ) -> RbResult<RbSeries> {
+        let seed = PlFixedStateQuality::default().hash_one((k0, k1, k2, k3));
+        let hb = PlSeedableRandomStateQuality::seed_from_u64(seed);
+        rb.enter_polars_series(|| self_.df.write().hash_rows(Some(hb)))
     }
 
-    pub fn transpose(&self, keep_names_as: Option<String>, column_names: Value) -> RbResult<Self> {
+    pub fn transpose(
+        rb: &Ruby,
+        self_: &Self,
+        keep_names_as: Option<String>,
+        column_names: Value,
+    ) -> RbResult<Self> {
         let new_col_names = if let Ok(name) = <Vec<String>>::try_convert(column_names) {
             Some(Either::Right(name))
         } else if let Ok(name) = String::try_convert(column_names) {
@@ -558,45 +474,61 @@ impl RbDataFrame {
         } else {
             None
         };
-        Ok(self
-            .df
-            .borrow_mut()
-            .transpose(keep_names_as.as_deref(), new_col_names)
-            .map_err(RbPolarsErr::from)?
-            .into())
+        rb.enter_polars_df(|| {
+            self_
+                .df
+                .write()
+                .transpose(keep_names_as.as_deref(), new_col_names)
+        })
     }
 
     pub fn upsample(
-        &self,
+        rb: &Ruby,
+        self_: &Self,
         by: Vec<String>,
         index_column: String,
         every: String,
         stable: bool,
     ) -> RbResult<Self> {
-        let out = if stable {
-            self.df
-                .borrow()
-                .upsample_stable(by, &index_column, Duration::parse(&every))
-        } else {
-            self.df
-                .borrow()
-                .upsample(by, &index_column, Duration::parse(&every))
-        };
-        let out = out.map_err(RbPolarsErr::from)?;
-        Ok(out.into())
-    }
-
-    pub fn to_struct(&self, name: String) -> RbSeries {
-        let s = self.df.borrow().clone().into_struct(name.into());
-        s.into_series().into()
-    }
-
-    pub fn unnest(&self, names: Vec<String>) -> RbResult<Self> {
-        let df = self.df.borrow().unnest(names).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+        rb.enter_polars_df(|| {
+            if stable {
+                self_
+                    .df
+                    .read()
+                    .upsample_stable(by, &index_column, Duration::parse(&every))
+            } else {
+                self_
+                    .df
+                    .read()
+                    .upsample(by, &index_column, Duration::parse(&every))
+            }
+        })
+    }
+
+    pub fn to_struct(
+        rb: &Ruby,
+        self_: &Self,
+        name: String,
+        invalid_indices: Vec<usize>,
+    ) -> RbResult<RbSeries> {
+        rb.enter_polars_series(|| {
+            let mut ca = self_.df.read().clone().into_struct(name.into());
+
+            if !invalid_indices.is_empty() {
+                let mut validity = MutableBitmap::with_capacity(ca.len());
+                validity.extend_constant(ca.len(), true);
+                for i in invalid_indices {
+                    validity.set(i, false);
+                }
+                ca.rechunk_mut();
+                Ok(ca.with_outer_validity(Some(validity.freeze())))
+            } else {
+                Ok(ca)
+            }
+        })
     }
 
-    pub fn clear(&self) -> Self {
-        self.df.borrow().clear().into()
+    pub fn clear(rb: &Ruby, self_: &Self) -> RbResult<Self> {
+        rb.enter_polars_df(|| Ok(self_.df.read().clear()))
     }
 }