polars-df 0.21.0 → 0.22.0

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

@@ -1,16 +1,17 @@
 pub(crate) mod any_value;
 mod categorical;
 mod chunked_array;
+mod datetime;
 
+use std::collections::BTreeMap;
 use std::fmt::{Debug, Display, Formatter};
 use std::fs::File;
 use std::hash::{Hash, Hasher};
-use std::num::NonZeroUsize;
 
 pub use categorical::RbCategories;
 use magnus::{
-    IntoValue, Module, RArray, RHash, Ruby, Symbol, TryConvert, Value, class, exception,
-    prelude::*, r_hash::ForEach, try_convert::TryConvertOwned, value::Opaque,
+    IntoValue, Module, RArray, RHash, Ruby, Symbol, TryConvert, Value, prelude::*, r_hash::ForEach,
+    try_convert::TryConvertOwned, value::Opaque,
 };
 use polars::chunked_array::object::PolarsObjectSafe;
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
@@ -18,9 +19,13 @@ use polars::datatypes::AnyValue;
 use polars::frame::row::Row;
 use polars::io::avro::AvroCompression;
 use polars::io::cloud::CloudOptions;
+use polars::prelude::default_values::{
+    DefaultFieldValues, IcebergIdentityTransformedPartitionFields,
+};
 use polars::prelude::deletion::DeletionFilesList;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
+use polars_core::schema::iceberg::IcebergSchema;
 use polars_core::utils::arrow::array::Array;
 use polars_core::utils::materialize_dyn_int;
 use polars_plan::dsl::ScanSources;
@@ -29,7 +34,8 @@ use polars_utils::total_ord::{TotalEq, TotalHash};
 
 use crate::file::{RubyScanSourceInput, get_ruby_scan_source_input};
 use crate::object::OBJECT_NAME;
-use crate::rb_modules::series;
+use crate::rb_modules::pl_series;
+use crate::utils::to_rb_err;
 use crate::{RbDataFrame, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbTypeError, RbValueError};
 
 pub(crate) fn slice_extract_wrapped<T>(slice: &[Wrap<T>]) -> &[T] {
@@ -84,7 +90,7 @@ pub(crate) fn get_series(obj: Value) -> RbResult<Series> {
 }
 
 pub(crate) fn to_series(s: RbSeries) -> Value {
-    let series = series();
+    let series = pl_series();
     series
         .funcall::<_, _, Value>("_from_rbseries", (s,))
         .unwrap()
@@ -119,16 +125,16 @@ impl TryConvert for Wrap<NullValues> {
     }
 }
 
-fn struct_dict<'a>(vals: impl Iterator<Item = AnyValue<'a>>, flds: &[Field]) -> Value {
-    let dict = RHash::new();
+fn struct_dict<'a>(ruby: &Ruby, vals: impl Iterator<Item = AnyValue<'a>>, flds: &[Field]) -> Value {
+    let dict = ruby.hash_new();
     for (fld, val) in flds.iter().zip(vals) {
         dict.aset(fld.name().as_str(), Wrap(val)).unwrap()
     }
-    dict.into_value()
+    dict.as_value()
 }
 
 impl IntoValue for Wrap<DataType> {
-    fn into_value_with(self, _: &Ruby) -> Value {
+    fn into_value_with(self, ruby: &Ruby) -> Value {
         let pl = crate::rb_modules::polars();
 
         match self.0 {
@@ -234,8 +240,10 @@ impl IntoValue for Wrap<DataType> {
                 let categories: Value = categories_class
                     .funcall("_from_rb_categories", (RbCategories::from(cats.clone()),))
                     .unwrap();
-                let kwargs = RHash::new();
-                kwargs.aset(Symbol::new("categories"), categories).unwrap();
+                let kwargs = ruby.hash_new();
+                kwargs
+                    .aset(ruby.to_symbol("categories"), categories)
+                    .unwrap();
                 categorical_class.funcall("new", (kwargs,)).unwrap()
             }
             DataType::Enum(_, mapping) => {
@@ -262,7 +270,7 @@ impl IntoValue for Wrap<DataType> {
                         .funcall::<_, _, Value>("new", (name, dtype))
                         .unwrap()
                 });
-                let fields = RArray::from_iter(iter);
+                let fields = ruby.ary_from_iter(iter);
                 let struct_class = pl.const_get::<_, Value>("Struct").unwrap();
                 struct_class
                     .funcall::<_, _, Value>("new", (fields,))
@@ -273,7 +281,7 @@ impl IntoValue for Wrap<DataType> {
                 class.funcall("new", ()).unwrap()
             }
             DataType::Unknown(UnknownKind::Int(v)) => {
-                Wrap(materialize_dyn_int(v).dtype()).into_value()
+                Wrap(materialize_dyn_int(v).dtype()).into_value_with(ruby)
             }
             DataType::Unknown(_) => {
                 let class = pl.const_get::<_, Value>("Unknown").unwrap();
@@ -291,19 +299,19 @@ enum CategoricalOrdering {
 }
 
 impl IntoValue for Wrap<CategoricalOrdering> {
-    fn into_value_with(self, _: &Ruby) -> Value {
-        "lexical".into_value()
+    fn into_value_with(self, ruby: &Ruby) -> Value {
+        "lexical".into_value_with(ruby)
     }
 }
 
 impl IntoValue for Wrap<TimeUnit> {
-    fn into_value_with(self, _: &Ruby) -> Value {
+    fn into_value_with(self, ruby: &Ruby) -> Value {
         let tu = match self.0 {
             TimeUnit::Nanoseconds => "ns",
             TimeUnit::Microseconds => "us",
             TimeUnit::Milliseconds => "ms",
         };
-        tu.into_value()
+        tu.into_value_with(ruby)
     }
 }
 
@@ -317,7 +325,8 @@ impl TryConvert for Wrap<Field> {
 
 impl TryConvert for Wrap<DataType> {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        let dtype = if ob.is_kind_of(class::class()) {
+        let ruby = Ruby::get_with(ob);
+        let dtype = if ob.is_kind_of(ruby.class_class()) {
             let name = ob.funcall::<_, _, String>("name", ())?;
             match name.as_str() {
                 "Polars::Int8" => DataType::Int8,
@@ -524,6 +533,60 @@ impl TryConvert for Wrap<Schema> {
     }
 }
 
+impl TryConvert for Wrap<ArrowSchema> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let ruby = Ruby::get_with(ob);
+        // TODO improve
+        let ob = RHash::try_convert(ob)?;
+        let fields: RArray = ob.aref(ruby.to_symbol("fields"))?;
+        let mut arrow_schema = ArrowSchema::with_capacity(fields.len());
+        for f in fields {
+            let f = RHash::try_convert(f)?;
+            let name: String = f.aref(ruby.to_symbol("name"))?;
+            let rb_dtype: String = f.aref(ruby.to_symbol("type"))?;
+            let dtype = match rb_dtype.as_str() {
+                "null" => ArrowDataType::Null,
+                "boolean" => ArrowDataType::Boolean,
+                "int8" => ArrowDataType::Int8,
+                "int16" => ArrowDataType::Int16,
+                "int32" => ArrowDataType::Int32,
+                "int64" => ArrowDataType::Int64,
+                "uint8" => ArrowDataType::UInt8,
+                "uint16" => ArrowDataType::UInt16,
+                "uint32" => ArrowDataType::UInt32,
+                "uint64" => ArrowDataType::UInt64,
+                "float16" => ArrowDataType::Float16,
+                "float32" => ArrowDataType::Float32,
+                "float64" => ArrowDataType::Float64,
+                "date32" => ArrowDataType::Date32,
+                "date64" => ArrowDataType::Date64,
+                "binary" => ArrowDataType::Binary,
+                "large_binary" => ArrowDataType::LargeBinary,
+                "string" => ArrowDataType::Utf8,
+                "large_string" => ArrowDataType::LargeUtf8,
+                "binary_view" => ArrowDataType::BinaryView,
+                "string_view" => ArrowDataType::Utf8View,
+                "unknown" => ArrowDataType::Unknown,
+                _ => todo!(),
+            };
+            let is_nullable = f.aref(ruby.to_symbol("nullable"))?;
+            let rb_metadata: RHash = f.aref(ruby.to_symbol("metadata"))?;
+            let mut metadata = BTreeMap::new();
+            rb_metadata.foreach(|k: String, v: String| {
+                metadata.insert(k.into(), v.into());
+                Ok(ForEach::Continue)
+            })?;
+            arrow_schema
+                .try_insert(
+                    name.clone().into(),
+                    ArrowField::new(name.into(), dtype, is_nullable).with_metadata(metadata),
+                )
+                .map_err(to_rb_err)?;
+        }
+        Ok(Wrap(arrow_schema))
+    }
+}
+
 impl TryConvert for Wrap<ScanSources> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let list = RArray::try_convert(ob)?;
@@ -665,7 +728,7 @@ impl From<&dyn PolarsObjectSafe> for &ObjectValue {
 
 impl ObjectValue {
     pub fn to_value(&self) -> Value {
-        self.clone().into_value()
+        self.clone().into_value_with(&Ruby::get().unwrap())
     }
 }
 
@@ -979,6 +1042,22 @@ impl TryConvert for Wrap<RankMethod> {
     }
 }
 
+impl TryConvert for Wrap<Roll> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match String::try_convert(ob)?.as_str() {
+            "raise" => Roll::Raise,
+            "forward" => Roll::Forward,
+            "backward" => Roll::Backward,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "`roll` must be one of {{'raise', 'forward', 'backward'}}, got {v}",
+                )));
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<TimeUnit> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match String::try_convert(ob)?.as_str() {
@@ -1156,7 +1235,109 @@ impl TryConvert for Wrap<CastColumnsPolicy> {
             let out = Wrap(CastColumnsPolicy::ERROR_ON_MISMATCH);
             return Ok(out);
         }
-        todo!();
+
+        let integer_upcast = match &*ob.funcall::<_, _, String>("integer_cast", ())? {
+            "upcast" => true,
+            "forbid" => false,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "unknown option for integer_cast: {v}"
+                )));
+            }
+        };
+
+        let mut float_upcast = false;
+        let mut float_downcast = false;
+
+        let float_cast_object: Value = ob.funcall("float_cast", ())?;
+
+        parse_multiple_options("float_cast", float_cast_object, |v| {
+            match v {
+                "forbid" => {}
+                "upcast" => float_upcast = true,
+                "downcast" => float_downcast = true,
+                v => {
+                    return Err(RbValueError::new_err(format!(
+                        "unknown option for float_cast: {v}"
+                    )));
+                }
+            }
+
+            Ok(())
+        })?;
+
+        let mut datetime_nanoseconds_downcast = false;
+        let mut datetime_convert_timezone = false;
+
+        let datetime_cast_object: Value = ob.funcall("datetime_cast", ())?;
+
+        parse_multiple_options("datetime_cast", datetime_cast_object, |v| {
+            match v {
+                "forbid" => {}
+                "nanosecond-downcast" => datetime_nanoseconds_downcast = true,
+                "convert-timezone" => datetime_convert_timezone = true,
+                v => {
+                    return Err(RbValueError::new_err(format!(
+                        "unknown option for datetime_cast: {v}"
+                    )));
+                }
+            };
+
+            Ok(())
+        })?;
+
+        let missing_struct_fields =
+            match &*ob.funcall::<_, _, String>("missing_struct_fields", ())? {
+                "insert" => MissingColumnsPolicy::Insert,
+                "raise" => MissingColumnsPolicy::Raise,
+                v => {
+                    return Err(RbValueError::new_err(format!(
+                        "unknown option for missing_struct_fields: {v}"
+                    )));
+                }
+            };
+
+        let extra_struct_fields = match &*ob.funcall::<_, _, String>("extra_struct_fields", ())? {
+            "ignore" => ExtraColumnsPolicy::Ignore,
+            "raise" => ExtraColumnsPolicy::Raise,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "unknown option for extra_struct_fields: {v}"
+                )));
+            }
+        };
+
+        return Ok(Wrap(CastColumnsPolicy {
+            integer_upcast,
+            float_upcast,
+            float_downcast,
+            datetime_nanoseconds_downcast,
+            datetime_microseconds_downcast: false,
+            datetime_convert_timezone,
+            null_upcast: true,
+            missing_struct_fields,
+            extra_struct_fields,
+        }));
+
+        fn parse_multiple_options(
+            parameter_name: &'static str,
+            rb_object: Value,
+            mut parser_func: impl FnMut(&str) -> RbResult<()>,
+        ) -> RbResult<()> {
+            if let Ok(v) = String::try_convert(rb_object) {
+                parser_func(&v)?;
+            } else if let Ok(v) = RArray::try_convert(rb_object) {
+                for v in v {
+                    parser_func(&String::try_convert(v)?)?;
+                }
+            } else {
+                return Err(RbValueError::new_err(format!(
+                    "unknown type for {parameter_name}: {rb_object}"
+                )));
+            }
+
+            Ok(())
+        }
     }
 }
 
@@ -1174,7 +1355,7 @@ pub fn parse_fill_null_strategy(
         "one" => FillNullStrategy::One,
         e => {
             return Err(magnus::Error::new(
-                exception::runtime_error(),
+                Ruby::get().unwrap().exception_runtime_error(),
                 format!(
                     "strategy must be one of {{'forward', 'backward', 'min', 'max', 'mean', 'zero', 'one'}}, got {e}",
                 ),
@@ -1225,15 +1406,6 @@ pub fn parse_parquet_compression(
     Ok(parsed)
 }
 
-impl TryConvert for Wrap<NonZeroUsize> {
-    fn try_convert(ob: Value) -> RbResult<Self> {
-        let v = usize::try_convert(ob)?;
-        NonZeroUsize::new(v)
-            .map(Wrap)
-            .ok_or(RbValueError::new_err("must be non-zero"))
-    }
-}
-
 pub(crate) fn strings_to_pl_smallstr<I, S>(container: I) -> Vec<PlSmallStr>
 where
     I: IntoIterator<Item = S>,
@@ -1270,6 +1442,23 @@ impl TryConvert for RbCompatLevel {
     }
 }
 
+impl TryConvert for Wrap<UnicodeForm> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match String::try_convert(ob)?.as_str() {
+            "NFC" => UnicodeForm::NFC,
+            "NFKC" => UnicodeForm::NFKC,
+            "NFD" => UnicodeForm::NFD,
+            "NFKD" => UnicodeForm::NFKD,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "`form` must be one of {{'NFC', 'NFKC', 'NFD', 'NFKD'}}, got {v}",
+                )));
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<Option<TimeZone>> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let tz = Option::<Wrap<PlSmallStr>>::try_convert(ob)?;
@@ -1313,13 +1502,97 @@ impl TryConvert for Wrap<MissingColumnsPolicy> {
 }
 
 impl TryConvert for Wrap<ColumnMapping> {
-    fn try_convert(_ob: Value) -> RbResult<Self> {
-        todo!()
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let (column_mapping_type, ob) = <(String, Value)>::try_convert(ob)?;
+
+        Ok(Wrap(match column_mapping_type.as_str() {
+            "iceberg-column-mapping" => {
+                let arrow_schema = Wrap::<ArrowSchema>::try_convert(ob)?;
+                ColumnMapping::Iceberg(Arc::new(
+                    IcebergSchema::from_arrow_schema(&arrow_schema.0).map_err(to_rb_err)?,
+                ))
+            }
+
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "unknown column mapping type: {v}"
+                )));
+            }
+        }))
     }
 }
 
 impl TryConvert for Wrap<DeletionFilesList> {
-    fn try_convert(_ob: Value) -> RbResult<Self> {
-        todo!();
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let (deletion_file_type, ob) = <(String, Value)>::try_convert(ob)?;
+
+        Ok(Wrap(match deletion_file_type.as_str() {
+            "iceberg-position-delete" => {
+                let dict = RHash::try_convert(ob)?;
+
+                let mut out = PlIndexMap::new();
+
+                dict.foreach(|k: usize, v: RArray| {
+                    let files = v
+                        .into_iter()
+                        .map(|x| {
+                            let x = String::try_convert(x)?;
+                            Ok(x)
+                        })
+                        .collect::<RbResult<Arc<[String]>>>()?;
+
+                    if !files.is_empty() {
+                        out.insert(k, files);
+                    }
+
+                    Ok(ForEach::Continue)
+                })?;
+
+                DeletionFilesList::IcebergPositionDelete(Arc::new(out))
+            }
+
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "unknown deletion file type: {v}"
+                )));
+            }
+        }))
+    }
+}
+
+impl TryConvert for Wrap<DefaultFieldValues> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let (default_values_type, ob) = <(String, Value)>::try_convert(ob)?;
+
+        Ok(Wrap(match &*default_values_type {
+            "iceberg" => {
+                let dict = RHash::try_convert(ob)?;
+
+                let mut out = PlIndexMap::new();
+
+                dict.foreach(|k: u32, v: Value| {
+                    let v: Result<Column, String> = if let Ok(s) = get_series(v) {
+                        Ok(s.into_column())
+                    } else {
+                        let err_msg = String::try_convert(v)?;
+                        Err(err_msg)
+                    };
+
+                    out.insert(k, v);
+
+                    Ok(ForEach::Continue)
+                })?;
+
+                DefaultFieldValues::Iceberg(Arc::new(IcebergIdentityTransformedPartitionFields(
+                    out,
+                )))
+            }
+
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "unknown deletion file type: {v}"
+                )));
+            }
+        }))
     }
 }

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

@@ -1,4 +1,4 @@
-use magnus::{RArray, RHash, Symbol, Value, prelude::*, r_hash::ForEach};
+use magnus::{RArray, RHash, Ruby, Symbol, Value, prelude::*, r_hash::ForEach};
 use polars::frame::row::{Row, rows_to_schema_supertypes, rows_to_supertypes};
 use polars::prelude::*;
 
@@ -125,6 +125,7 @@ where
 }
 
 fn dicts_to_rows<'a>(data: &Value, names: &'a [String], _strict: bool) -> RbResult<Vec<Row<'a>>> {
+    let ruby = Ruby::get_with(*data);
     let (data, len) = get_rbseq(*data)?;
     let mut rows = Vec::with_capacity(len);
     for d in data.into_iter() {
@@ -132,8 +133,8 @@ fn dicts_to_rows<'a>(data: &Value, names: &'a [String], _strict: bool) -> RbResu
 
         let mut row = Vec::with_capacity(names.len());
         for k in names.iter() {
-            // TODO improve performance
-            let val = match d.get(k.clone()).or_else(|| d.get(Symbol::new(k))) {
+            // TODO improve performance (must work with GC)
+            let val = match d.get(k.clone()).or_else(|| d.get(ruby.to_symbol(k))) {
                 None => AnyValue::Null,
                 Some(val) => Wrap::<AnyValue>::try_convert(val)?.0,
             };

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

@@ -1,4 +1,4 @@
-use magnus::{IntoValue, RArray, Value, prelude::*};
+use magnus::{IntoValue, Ruby, Value, prelude::*};
 
 use super::*;
 use crate::RbResult;
@@ -6,14 +6,15 @@ use crate::conversion::{ObjectValue, Wrap};
 use crate::interop::arrow::to_ruby::dataframe_to_stream;
 
 impl RbDataFrame {
-    pub fn row_tuple(&self, idx: i64) -> Value {
+    pub fn row_tuple(ruby: &Ruby, rb_self: &Self, idx: i64) -> Value {
         let idx = if idx < 0 {
-            (self.df.borrow().height() as i64 + idx) as usize
+            (rb_self.df.borrow().height() as i64 + idx) as usize
         } else {
             idx as usize
         };
-        RArray::from_iter(
-            self.df
+        ruby.ary_from_iter(
+            rb_self
+                .df
                 .borrow()
                 .get_columns()
                 .iter()
@@ -22,17 +23,18 @@ impl RbDataFrame {
                         let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
                         obj.unwrap().to_value()
                     }
-                    _ => Wrap(s.get(idx).unwrap()).into_value(),
+                    _ => Wrap(s.get(idx).unwrap()).into_value_with(ruby),
                 }),
         )
         .as_value()
     }
 
-    pub fn row_tuples(&self) -> Value {
-        let df = &self.df;
-        RArray::from_iter((0..df.borrow().height()).map(|idx| {
-            RArray::from_iter(
-                self.df
+    pub fn row_tuples(ruby: &Ruby, rb_self: &Self) -> Value {
+        let df = &rb_self.df;
+        ruby.ary_from_iter((0..df.borrow().height()).map(|idx| {
+            ruby.ary_from_iter(
+                rb_self
+                    .df
                     .borrow()
                     .get_columns()
                     .iter()
@@ -41,15 +43,15 @@ impl RbDataFrame {
                             let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
                             obj.unwrap().to_value()
                         }
-                        _ => Wrap(s.get(idx).unwrap()).into_value(),
+                        _ => Wrap(s.get(idx).unwrap()).into_value_with(ruby),
                     }),
             )
         }))
         .as_value()
     }
 
-    pub fn __arrow_c_stream__(&self) -> RbResult<Value> {
-        self.df.borrow_mut().align_chunks();
-        dataframe_to_stream(&self.df.borrow())
+    pub fn __arrow_c_stream__(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
+        rb_self.df.borrow_mut().align_chunks();
+        dataframe_to_stream(&rb_self.df.borrow(), ruby)
     }
 }

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

@@ -1,7 +1,7 @@
 use std::hash::BuildHasher;
 
 use either::Either;
-use magnus::{IntoValue, RArray, Value, prelude::*, typed_data::Obj};
+use magnus::{IntoValue, RArray, Ruby, Value, prelude::*};
 use polars::prelude::pivot::{pivot, pivot_stable};
 use polars::prelude::*;
 
@@ -149,12 +149,13 @@ impl RbDataFrame {
         Ok(())
     }
 
-    pub fn dtypes(&self) -> RArray {
-        RArray::from_iter(
-            self.df
+    pub fn dtypes(ruby: &Ruby, rb_self: &Self) -> RArray {
+        ruby.ary_from_iter(
+            rb_self
+                .df
                 .borrow()
                 .iter()
-                .map(|s| Wrap(s.dtype().clone()).into_value()),
+                .map(|s| Wrap(s.dtype().clone()).into_value_with(ruby)),
         )
     }
 
@@ -393,18 +394,19 @@ impl RbDataFrame {
     }
 
     pub fn partition_by(
-        &self,
+        ruby: &Ruby,
+        rb_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)
+            rb_self.df.borrow().partition_by_stable(by, include_key)
         } else {
-            self.df.borrow().partition_by(by, include_key)
+            rb_self.df.borrow().partition_by(by, include_key)
         }
         .map_err(RbPolarsErr::from)?;
-        Ok(RArray::from_iter(out.into_iter().map(RbDataFrame::new)))
+        Ok(ruby.ary_from_iter(out.into_iter().map(RbDataFrame::new)))
     }
 
     pub fn lazy(&self) -> RbLazyFrame {
@@ -440,12 +442,13 @@ impl RbDataFrame {
     }
 
     pub fn map_rows(
-        &self,
+        ruby: &Ruby,
+        rb_self: &Self,
         lambda: Value,
         output_type: Option<Wrap<DataType>>,
         inference_size: usize,
     ) -> RbResult<(Value, bool)> {
-        let df = &self.df.borrow();
+        let df = &rb_self.df.borrow();
 
         let output_type = output_type.map(|dt| dt.0);
         let out = match output_type {
@@ -490,7 +493,7 @@ impl RbDataFrame {
             _ => return apply_lambda_unknown(df, lambda, inference_size),
         };
 
-        Ok((Obj::wrap(RbSeries::from(out)).as_value(), false))
+        Ok((ruby.obj_wrap(RbSeries::from(out)).as_value(), false))
     }
 
     pub fn shrink_to_fit(&self) {

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

@@ -249,14 +249,13 @@ impl RbDataFrame {
         include_header: bool,
         separator: u8,
         quote_char: u8,
-        batch_size: Wrap<NonZeroUsize>,
+        batch_size: NonZeroUsize,
         datetime_format: Option<String>,
         date_format: Option<String>,
         time_format: Option<String>,
         float_precision: Option<usize>,
         null_value: Option<String>,
     ) -> RbResult<()> {
-        let batch_size = batch_size.0;
         let null = null_value.unwrap_or_default();
         let mut buf = get_file_like(rb_f, true)?;
         CsvWriter::new(&mut buf)