polars-df 0.21.1 → 0.23.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) {
@@ -549,11 +552,6 @@ impl RbDataFrame {
         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())
-    }
-
     pub fn clear(&self) -> Self {
         self.df.borrow().clear().into()
     }

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)

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

@@ -4,10 +4,14 @@ mod general;
 mod io;
 mod serde;
 
+use magnus::{DataTypeFunctions, TypedData, gc};
 use polars::prelude::*;
 use std::cell::RefCell;
 
-#[magnus::wrap(class = "Polars::RbDataFrame")]
+use crate::series::mark_series;
+
+#[derive(TypedData)]
+#[magnus(class = "Polars::RbDataFrame", mark)]
 pub struct RbDataFrame {
     pub df: RefCell<DataFrame>,
 }
@@ -25,3 +29,23 @@ impl RbDataFrame {
         }
     }
 }
+
+impl DataTypeFunctions for RbDataFrame {
+    fn mark(&self, marker: &gc::Marker) {
+        // this is really, really not ideal, as objects will not be marked if unable to borrow
+        // currently, this should only happen for write_* methods,
+        // which should refuse to write Object datatype, and therefore be safe,
+        // since GC will not have a chance to run
+        if let Ok(df) = self.df.try_borrow() {
+            for column in df.get_columns() {
+                if let DataType::Object(_) = column.dtype() {
+                    match column {
+                        Column::Series(s) => mark_series(marker, s),
+                        Column::Partitioned(s) => mark_series(marker, s.partitions()),
+                        Column::Scalar(s) => mark_series(marker, &s.as_single_value_series()),
+                    }
+                }
+            }
+        }
+    }
+}

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

@@ -1,15 +1,30 @@
-use crate::exceptions::ComputeError;
+use std::io::{BufReader, BufWriter};
+
+use polars::prelude::*;
+
 use crate::file::get_file_like;
-use crate::{RbDataFrame, RbResult};
+use crate::utils::to_rb_err;
+use crate::{RbDataFrame, RbPolarsErr, RbResult};
 use magnus::Value;
-use std::io::BufWriter;
 
 impl RbDataFrame {
-    // TODO add to Ruby
-    pub fn serialize_json(&self, rb_f: Value) -> RbResult<()> {
+    pub fn serialize_binary(&self, rb_f: Value) -> RbResult<()> {
         let file = get_file_like(rb_f, true)?;
-        let writer = BufWriter::new(file);
-        serde_json::to_writer(writer, &self.df)
-            .map_err(|err| ComputeError::new_err(err.to_string()))
+        let mut writer = BufWriter::new(file);
+
+        Ok(self
+            .df
+            .borrow_mut()
+            .serialize_into_writer(&mut writer)
+            .map_err(RbPolarsErr::from)?)
+    }
+
+    pub fn deserialize_binary(rb_f: Value) -> RbResult<Self> {
+        let file = get_file_like(rb_f, false)?;
+        let mut file = BufReader::new(file);
+
+        DataFrame::deserialize_from_reader(&mut file)
+            .map(|v| v.into())
+            .map_err(to_rb_err)
     }
 }

data/ext/polars/src/exceptions.rs

@@ -1,5 +1,5 @@
 use crate::rb_modules;
-use magnus::{Error, exception};
+use magnus::{Error, Ruby};
 use std::borrow::Cow;
 
 macro_rules! create_exception {
@@ -17,9 +17,12 @@ macro_rules! create_exception {
     };
 }
 
-create_exception!(RbTypeError, exception::type_error());
-create_exception!(RbValueError, exception::arg_error());
-create_exception!(RbOverflowError, exception::range_error());
-create_exception!(RbIndexError, exception::index_error());
+create_exception!(RbTypeError, Ruby::get().unwrap().exception_type_error());
+create_exception!(RbValueError, Ruby::get().unwrap().exception_arg_error());
+create_exception!(
+    RbOverflowError,
+    Ruby::get().unwrap().exception_range_error()
+);
+create_exception!(RbIndexError, Ruby::get().unwrap().exception_index_error());
 create_exception!(ComputeError, rb_modules::compute_error());
 create_exception!(InvalidOperationError, rb_modules::invalid_operation_error());

data/ext/polars/src/expr/datatype.rs

@@ -1,4 +1,4 @@
-use magnus::{IntoValue, Value};
+use magnus::{IntoValue, Ruby, Value};
 use polars::prelude::{DataType, DataTypeExpr, Schema};
 
 use crate::prelude::Wrap;
@@ -26,12 +26,12 @@ impl RbDataTypeExpr {
         DataTypeExpr::OfExpr(Box::new(expr.inner.clone())).into()
     }
 
-    pub fn collect_dtype(&self, schema: Wrap<Schema>) -> RbResult<Value> {
-        let dtype = self
+    pub fn collect_dtype(ruby: &Ruby, rb_self: &Self, schema: Wrap<Schema>) -> RbResult<Value> {
+        let dtype = rb_self
             .clone()
             .inner
             .into_datatype(&schema.0)
             .map_err(RbPolarsErr::from)?;
-        Ok(Wrap(dtype).into_value())
+        Ok(Wrap(dtype).into_value_with(ruby))
     }
 }

data/ext/polars/src/expr/datetime.rs

@@ -26,20 +26,6 @@ impl RbExpr {
         self.inner.clone().dt().offset_by(by.inner.clone()).into()
     }
 
-    pub fn dt_epoch_seconds(&self) -> Self {
-        self.inner
-            .clone()
-            .map(
-                |s| {
-                    s.take_materialized_series()
-                        .timestamp(TimeUnit::Milliseconds)
-                        .map(|ca| Some((ca / 1000).into_column()))
-                },
-                GetOutput::from_type(DataType::Int64),
-            )
-            .into()
-    }
-
     pub fn dt_with_time_unit(&self, tu: Wrap<TimeUnit>) -> Self {
         self.inner.clone().dt().with_time_unit(tu.0).into()
     }
@@ -232,31 +218,39 @@ impl RbExpr {
         self.inner.clone().dt().timestamp(tu.0).into()
     }
 
-    pub fn dt_total_days(&self) -> Self {
-        self.inner.clone().dt().total_days().into()
+    pub fn dt_total_days(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_days(fractional).into()
     }
 
-    pub fn dt_total_hours(&self) -> Self {
-        self.inner.clone().dt().total_hours().into()
+    pub fn dt_total_hours(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_hours(fractional).into()
     }
 
-    pub fn dt_total_minutes(&self) -> Self {
-        self.inner.clone().dt().total_minutes().into()
+    pub fn dt_total_minutes(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_minutes(fractional).into()
     }
 
-    pub fn dt_total_seconds(&self) -> Self {
-        self.inner.clone().dt().total_seconds().into()
+    pub fn dt_total_seconds(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_seconds(fractional).into()
     }
 
-    pub fn dt_total_milliseconds(&self) -> Self {
-        self.inner.clone().dt().total_milliseconds().into()
+    pub fn dt_total_milliseconds(&self, fractional: bool) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .total_milliseconds(fractional)
+            .into()
     }
 
-    pub fn dt_total_microseconds(&self) -> Self {
-        self.inner.clone().dt().total_microseconds().into()
+    pub fn dt_total_microseconds(&self, fractional: bool) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .total_microseconds(fractional)
+            .into()
     }
 
-    pub fn dt_total_nanoseconds(&self) -> Self {
-        self.inner.clone().dt().total_nanoseconds().into()
+    pub fn dt_total_nanoseconds(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_nanoseconds(fractional).into()
     }
 }

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

@@ -483,10 +483,7 @@ impl RbExpr {
     }
 
     pub fn rechunk(&self) -> Self {
-        self.inner
-            .clone()
-            .map(|s| Ok(Some(s.rechunk())), GetOutput::same_type())
-            .into()
+        self.inner.clone().rechunk().into()
     }
 
     pub fn round(&self, decimals: u32, mode: Wrap<RoundMode>) -> Self {
@@ -665,10 +662,6 @@ impl RbExpr {
         self.inner.clone().product().into()
     }
 
-    pub fn shrink_dtype(&self) -> Self {
-        self.inner.clone().shrink_dtype().into()
-    }
-
     pub fn map_batches(
         &self,
         lambda: Value,
@@ -868,8 +861,8 @@ impl RbExpr {
         self.inner.clone().all(drop_nulls).into()
     }
 
-    pub fn log(&self, base: f64) -> Self {
-        self.inner.clone().log(base).into()
+    pub fn log(&self, base: &RbExpr) -> Self {
+        self.inner.clone().log(base.inner.clone()).into()
     }
 
     pub fn log1p(&self) -> Self {

data/ext/polars/src/expr/list.rs

@@ -1,4 +1,4 @@
-use magnus::{Ruby, Value, prelude::*, value::Opaque};
+use magnus::{RArray, prelude::*};
 use polars::lazy::dsl::lit;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
@@ -195,33 +195,17 @@ impl RbExpr {
         self.inner.clone().list().to_array(width).into()
     }
 
-    pub fn list_to_struct(
-        &self,
-        width_strat: Wrap<ListToStructWidthStrategy>,
-        name_gen: Option<Value>,
-        upper_bound: Option<usize>,
-    ) -> RbResult<Self> {
-        let name_gen = name_gen.map(|lambda| {
-            let lambda = Opaque::from(lambda);
-            Arc::new(move |idx: usize| {
-                let lambda = Ruby::get().unwrap().get_inner(lambda);
-                let out: String = lambda.funcall("call", (idx,)).unwrap();
-                PlSmallStr::from_string(out)
-            });
-
-            // non-Ruby thread
-            todo!();
-        });
-
+    pub fn list_to_struct(&self, names: RArray) -> RbResult<Self> {
         Ok(self
             .inner
             .clone()
             .list()
-            .to_struct(ListToStruct::InferWidth {
-                infer_field_strategy: width_strat.0,
-                get_index_name: name_gen,
-                max_fields: upper_bound,
-            })
+            .to_struct(
+                names
+                    .into_iter()
+                    .map(|x| Ok(Wrap::<PlSmallStr>::try_convert(x)?.0))
+                    .collect::<RbResult<Arc<[_]>>>()?,
+            )
             .into())
     }
 

data/ext/polars/src/expr/meta.rs

@@ -1,4 +1,4 @@
-use magnus::RArray;
+use magnus::{RArray, Ruby};
 use polars::prelude::Schema;
 
 use crate::{RbExpr, RbPolarsErr, RbResult, Wrap};
@@ -8,17 +8,15 @@ impl RbExpr {
         self.inner == other.inner
     }
 
-    pub fn meta_pop(&self, schema: Option<Wrap<Schema>>) -> RbResult<RArray> {
+    pub fn meta_pop(ruby: &Ruby, rb_self: &Self, schema: Option<Wrap<Schema>>) -> RbResult<RArray> {
         let schema = schema.as_ref().map(|s| &s.0);
-        let exprs = self
+        let exprs = rb_self
             .inner
             .clone()
             .meta()
             .pop(schema)
             .map_err(RbPolarsErr::from)?;
-        Ok(RArray::from_iter(
-            exprs.iter().map(|e| RbExpr::from(e.clone())),
-        ))
+        Ok(ruby.ary_from_iter(exprs.iter().map(|e| RbExpr::from(e.clone()))))
     }
 
     pub fn meta_root_names(&self) -> Vec<String> {

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

@@ -10,6 +10,8 @@ mod meta;
 mod name;
 mod rolling;
 pub mod selector;
+#[cfg(feature = "serialize_binary")]
+mod serde;
 mod string;
 mod r#struct;
 

data/ext/polars/src/expr/name.rs

@@ -11,20 +11,17 @@ impl RbExpr {
 
     pub fn name_map(&self, lambda: Proc) -> Self {
         let lambda = Opaque::from(lambda);
-        self.inner
-            .clone()
-            .name()
-            .map(move |name| {
-                let lambda = Ruby::get().unwrap().get_inner(lambda);
-                let out = lambda.call::<_, String>((name.as_str(),));
-                match out {
-                    Ok(out) => Ok(format_pl_smallstr!("{}", out)),
-                    Err(e) => Err(PolarsError::ComputeError(
-                        format!("Ruby function in 'name.map' produced an error: {e}.").into(),
-                    )),
-                }
-            })
-            .into()
+        let func = PlanCallback::new(move |name: PlSmallStr| {
+            let lambda = Ruby::get().unwrap().get_inner(lambda);
+            let out = lambda.call::<_, String>((name.as_str(),));
+            match out {
+                Ok(out) => Ok(format_pl_smallstr!("{}", out)),
+                Err(e) => Err(PolarsError::ComputeError(
+                    format!("Ruby function in 'name.map' produced an error: {e}.").into(),
+                )),
+            }
+        });
+        self.inner.clone().name().map(func).into()
     }
 
     pub fn name_prefix(&self, prefix: String) -> Self {

data/ext/polars/src/expr/serde.rs

@@ -0,0 +1,28 @@
+use std::io::{BufReader, BufWriter};
+
+use magnus::Value;
+use polars::lazy::prelude::Expr;
+use polars_utils::pl_serialize;
+
+use crate::exceptions::ComputeError;
+use crate::file::get_file_like;
+use crate::{RbExpr, RbResult};
+
+impl RbExpr {
+    pub fn serialize_binary(&self, rb_f: Value) -> RbResult<()> {
+        let file = get_file_like(rb_f, true)?;
+        let writer = BufWriter::new(file);
+        pl_serialize::SerializeOptions::default()
+            .serialize_into_writer::<_, _, true>(writer, &self.inner)
+            .map_err(|err| ComputeError::new_err(err.to_string()))
+    }
+
+    pub fn deserialize_binary(rb_f: Value) -> RbResult<RbExpr> {
+        let file = get_file_like(rb_f, false)?;
+        let reader = BufReader::new(file);
+        let expr: Expr = pl_serialize::SerializeOptions::default()
+            .deserialize_from_reader::<_, _, true>(reader)
+            .map_err(|err| ComputeError::new_err(err.to_string()))?;
+        Ok(expr.into())
+    }
+}

data/ext/polars/src/expr/string.rs

@@ -1,7 +1,7 @@
 use polars::prelude::*;
 
 use crate::conversion::Wrap;
-use crate::{RbExpr, RbPolarsErr, RbResult};
+use crate::{RbDataTypeExpr, RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
     pub fn str_join(&self, delimiter: String, ignore_nulls: bool) -> Self {
@@ -265,16 +265,11 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_json_decode(
-        &self,
-        dtype: Option<Wrap<DataType>>,
-        infer_schema_len: Option<usize>,
-    ) -> Self {
-        let dtype = dtype.map(|wrap| wrap.0);
+    pub fn str_json_decode(&self, dtype: &RbDataTypeExpr) -> Self {
         self.inner
             .clone()
             .str()
-            .json_decode(dtype, infer_schema_len)
+            .json_decode(dtype.inner.clone())
             .into()
     }
 
@@ -352,8 +347,8 @@ impl RbExpr {
         self.inner.clone().str().splitn(by.inner.clone(), n).into()
     }
 
-    pub fn str_to_decimal(&self, infer_len: usize) -> Self {
-        self.inner.clone().str().to_decimal(infer_len).into()
+    pub fn str_to_decimal(&self, scale: usize) -> Self {
+        self.inner.clone().str().to_decimal(scale).into()
     }
 
     pub fn str_contains_any(&self, patterns: &RbExpr, ascii_case_insensitive: bool) -> Self {