polars-df 0.21.1 → 0.23.0

data/ext/polars/src/series/aggregation.rs

@@ -1,10 +1,12 @@
 use crate::error::RbPolarsErr;
 use crate::prelude::*;
+use crate::utils::to_rb_err;
 use crate::{RbResult, RbSeries};
-use magnus::{IntoValue, Value};
+use magnus::{IntoValue, Ruby, Value};
 
 fn scalar_to_rb(scalar: RbResult<Scalar>) -> RbResult<Value> {
-    Ok(Wrap(scalar?.as_any_value()).into_value())
+    let ruby = Ruby::get().unwrap();
+    Ok(Wrap(scalar?.as_any_value()).into_value_with(&ruby))
 }
 
 impl RbSeries {
@@ -36,40 +38,50 @@ impl RbSeries {
         self.series.borrow().arg_min()
     }
 
-    pub fn max(&self) -> RbResult<Value> {
+    pub fn max(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
         Ok(Wrap(
-            self.series
+            rb_self
+                .series
                 .borrow()
                 .max_reduce()
                 .map_err(RbPolarsErr::from)?
                 .as_any_value(),
         )
-        .into_value())
+        .into_value_with(ruby))
     }
 
-    pub fn mean(&self) -> RbResult<Value> {
-        match self.series.borrow().dtype() {
+    pub fn mean(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
+        match rb_self.series.borrow().dtype() {
             DataType::Boolean => Ok(Wrap(
-                self.series
+                rb_self
+                    .series
                     .borrow()
                     .cast(&DataType::UInt8)
                     .unwrap()
                     .mean_reduce()
+                    .map_err(to_rb_err)?
                     .as_any_value(),
             )
-            .into_value()),
+            .into_value_with(ruby)),
             // For non-numeric output types we require mean_reduce.
-            dt if dt.is_temporal() => {
-                Ok(Wrap(self.series.borrow().mean_reduce().as_any_value()).into_value())
-            }
-            _ => Ok(self.series.borrow().mean().into_value()),
+            dt if dt.is_temporal() => Ok(Wrap(
+                rb_self
+                    .series
+                    .borrow()
+                    .mean_reduce()
+                    .map_err(to_rb_err)?
+                    .as_any_value(),
+            )
+            .into_value_with(ruby)),
+            _ => Ok(rb_self.series.borrow().mean().into_value_with(ruby)),
         }
     }
 
-    pub fn median(&self) -> RbResult<Value> {
-        match self.series.borrow().dtype() {
+    pub fn median(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
+        match rb_self.series.borrow().dtype() {
             DataType::Boolean => Ok(Wrap(
-                self.series
+                rb_self
+                    .series
                     .borrow()
                     .cast(&DataType::UInt8)
                     .unwrap()
@@ -77,50 +89,58 @@ impl RbSeries {
                     .map_err(RbPolarsErr::from)?
                     .as_any_value(),
             )
-            .into_value()),
+            .into_value_with(ruby)),
             // For non-numeric output types we require median_reduce.
             dt if dt.is_temporal() => Ok(Wrap(
-                self.series
+                rb_self
+                    .series
                     .borrow()
                     .median_reduce()
                     .map_err(RbPolarsErr::from)?
                     .as_any_value(),
             )
-            .into_value()),
-            _ => Ok(self.series.borrow().median().into_value()),
+            .into_value_with(ruby)),
+            _ => Ok(rb_self.series.borrow().median().into_value_with(ruby)),
         }
     }
 
-    pub fn min(&self) -> RbResult<Value> {
+    pub fn min(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
         Ok(Wrap(
-            self.series
+            rb_self
+                .series
                 .borrow()
                 .min_reduce()
                 .map_err(RbPolarsErr::from)?
                 .as_any_value(),
         )
-        .into_value())
+        .into_value_with(ruby))
     }
 
-    pub fn quantile(&self, quantile: f64, interpolation: Wrap<QuantileMethod>) -> RbResult<Value> {
-        let bind = self
+    pub fn quantile(
+        ruby: &Ruby,
+        rb_self: &Self,
+        quantile: f64,
+        interpolation: Wrap<QuantileMethod>,
+    ) -> RbResult<Value> {
+        let bind = rb_self
             .series
             .borrow()
             .quantile_reduce(quantile, interpolation.0);
         let sc = bind.map_err(RbPolarsErr::from)?;
 
-        Ok(Wrap(sc.as_any_value()).into_value())
+        Ok(Wrap(sc.as_any_value()).into_value_with(ruby))
     }
 
-    pub fn sum(&self) -> RbResult<Value> {
+    pub fn sum(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
         Ok(Wrap(
-            self.series
+            rb_self
+                .series
                 .borrow()
                 .sum_reduce()
                 .map_err(RbPolarsErr::from)?
                 .as_any_value(),
         )
-        .into_value())
+        .into_value_with(ruby))
     }
 
     pub fn first(&self) -> RbResult<Value> {

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

@@ -75,10 +75,12 @@ init_method_opt!(new_opt_u8, UInt8Type, u8);
 init_method_opt!(new_opt_u16, UInt16Type, u16);
 init_method_opt!(new_opt_u32, UInt32Type, u32);
 init_method_opt!(new_opt_u64, UInt64Type, u64);
+init_method_opt!(new_opt_u128, UInt128Type, u128);
 init_method_opt!(new_opt_i8, Int8Type, i8);
 init_method_opt!(new_opt_i16, Int16Type, i16);
 init_method_opt!(new_opt_i32, Int32Type, i32);
 init_method_opt!(new_opt_i64, Int64Type, i64);
+init_method_opt!(new_opt_i128, Int128Type, i128);
 init_method_opt!(new_opt_f32, Float32Type, f32);
 init_method_opt!(new_opt_f64, Float64Type, f64);
 

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

@@ -1,4 +1,4 @@
-use magnus::{IntoValue, RArray, Value, value::qnil};
+use magnus::{IntoValue, Ruby, Value, value::ReprValue};
 use polars_core::prelude::*;
 
 use crate::RbSeries;
@@ -11,42 +11,44 @@ impl RbSeries {
         let series = &self.series.borrow();
 
         fn to_a_recursive(series: &Series) -> Value {
+            let ruby = Ruby::get().unwrap();
             let rblist = match series.dtype() {
-                DataType::Boolean => RArray::from_iter(series.bool().unwrap()).into_value(),
-                DataType::UInt8 => RArray::from_iter(series.u8().unwrap()).into_value(),
-                DataType::UInt16 => RArray::from_iter(series.u16().unwrap()).into_value(),
-                DataType::UInt32 => RArray::from_iter(series.u32().unwrap()).into_value(),
-                DataType::UInt64 => RArray::from_iter(series.u64().unwrap()).into_value(),
-                DataType::Int8 => RArray::from_iter(series.i8().unwrap()).into_value(),
-                DataType::Int16 => RArray::from_iter(series.i16().unwrap()).into_value(),
-                DataType::Int32 => RArray::from_iter(series.i32().unwrap()).into_value(),
-                DataType::Int64 => RArray::from_iter(series.i64().unwrap()).into_value(),
-                DataType::Int128 => todo!(),
-                DataType::Float32 => RArray::from_iter(series.f32().unwrap()).into_value(),
-                DataType::Float64 => RArray::from_iter(series.f64().unwrap()).into_value(),
+                DataType::Boolean => ruby.ary_from_iter(series.bool().unwrap()).as_value(),
+                DataType::UInt8 => ruby.ary_from_iter(series.u8().unwrap()).as_value(),
+                DataType::UInt16 => ruby.ary_from_iter(series.u16().unwrap()).as_value(),
+                DataType::UInt32 => ruby.ary_from_iter(series.u32().unwrap()).as_value(),
+                DataType::UInt64 => ruby.ary_from_iter(series.u64().unwrap()).as_value(),
+                DataType::UInt128 => ruby.ary_from_iter(series.u128().unwrap()).as_value(),
+                DataType::Int8 => ruby.ary_from_iter(series.i8().unwrap()).as_value(),
+                DataType::Int16 => ruby.ary_from_iter(series.i16().unwrap()).as_value(),
+                DataType::Int32 => ruby.ary_from_iter(series.i32().unwrap()).as_value(),
+                DataType::Int64 => ruby.ary_from_iter(series.i64().unwrap()).as_value(),
+                DataType::Int128 => ruby.ary_from_iter(series.i128().unwrap()).as_value(),
+                DataType::Float32 => ruby.ary_from_iter(series.f32().unwrap()).as_value(),
+                DataType::Float64 => ruby.ary_from_iter(series.f64().unwrap()).as_value(),
                 DataType::Categorical(_, _) | DataType::Enum(_, _) => {
                     with_match_categorical_physical_type!(series.dtype().cat_physical().unwrap(), |$C| {
-                        RArray::from_iter(series.cat::<$C>().unwrap().iter_str()).into_value()
+                        ruby.ary_from_iter(series.cat::<$C>().unwrap().iter_str()).as_value()
                     })
                 }
                 DataType::Object(_) => {
-                    let v = RArray::with_capacity(series.len());
+                    let v = ruby.ary_new_capa(series.len());
                     for i in 0..series.len() {
                         let obj: Option<&ObjectValue> = series.get_object(i).map(|any| any.into());
                         match obj {
                             Some(val) => v.push(val.to_value()).unwrap(),
-                            None => v.push(qnil()).unwrap(),
+                            None => v.push(ruby.qnil()).unwrap(),
                         };
                     }
-                    v.into_value()
+                    v.as_value()
                 }
                 DataType::List(_) => {
-                    let v = RArray::new();
+                    let v = ruby.ary_new();
                     let ca = series.list().unwrap();
                     for opt_s in ca.amortized_iter() {
                         match opt_s {
                             None => {
-                                v.push(qnil()).unwrap();
+                                v.push(ruby.qnil()).unwrap();
                             }
                             Some(s) => {
                                 let rblst = to_a_recursive(s.as_ref());
@@ -54,15 +56,15 @@ impl RbSeries {
                             }
                         }
                     }
-                    v.into_value()
+                    v.as_value()
                 }
                 DataType::Array(_, _) => {
-                    let v = RArray::new();
+                    let v = ruby.ary_new();
                     let ca = series.array().unwrap();
                     for opt_s in ca.amortized_iter() {
                         match opt_s {
                             None => {
-                                v.push(qnil()).unwrap();
+                                v.push(ruby.qnil()).unwrap();
                             }
                             Some(s) => {
                                 let rblst = to_a_recursive(s.as_ref());
@@ -70,39 +72,39 @@ impl RbSeries {
                             }
                         }
                     }
-                    v.into_value()
+                    v.as_value()
                 }
                 DataType::Date => {
                     let ca = series.date().unwrap();
-                    return Wrap(ca).into_value();
+                    return Wrap(ca).into_value_with(&ruby);
                 }
                 DataType::Time => {
                     let ca = series.time().unwrap();
-                    return Wrap(ca).into_value();
+                    return Wrap(ca).into_value_with(&ruby);
                 }
                 DataType::Datetime(_, _) => {
                     let ca = series.datetime().unwrap();
-                    return Wrap(ca).into_value();
+                    return Wrap(ca).into_value_with(&ruby);
                 }
                 DataType::Decimal(_, _) => {
                     let ca = series.decimal().unwrap();
-                    return Wrap(ca).into_value();
+                    return Wrap(ca).into_value_with(&ruby);
                 }
                 DataType::String => {
                     let ca = series.str().unwrap();
-                    return Wrap(ca).into_value();
+                    return Wrap(ca).into_value_with(&ruby);
                 }
                 DataType::Struct(_) => {
                     let ca = series.struct_().unwrap();
-                    return Wrap(ca).into_value();
+                    return Wrap(ca).into_value_with(&ruby);
                 }
                 DataType::Duration(_) => {
                     let ca = series.duration().unwrap();
-                    return Wrap(ca).into_value();
+                    return Wrap(ca).into_value_with(&ruby);
                 }
                 DataType::Binary => {
                     let ca = series.binary().unwrap();
-                    return Wrap(ca).into_value();
+                    return Wrap(ca).into_value_with(&ruby);
                 }
                 DataType::Null => {
                     let null: Option<u8> = None;
@@ -125,7 +127,10 @@ impl RbSeries {
                     }
                     impl ExactSizeIterator for NullIter {}
 
-                    RArray::from_iter(NullIter { iter, n }).into_value()
+                    Ruby::get()
+                        .unwrap()
+                        .ary_from_iter(NullIter { iter, n })
+                        .as_value()
                 }
                 DataType::Unknown(_) => {
                     panic!("to_a not implemented for unknown")
@@ -134,7 +139,7 @@ impl RbSeries {
                     unreachable!()
                 }
             };
-            rblist.into_value()
+            rblist.as_value()
         }
 
         to_a_recursive(series)

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

@@ -1,4 +1,4 @@
-use magnus::{Error, IntoValue, RArray, Value, exception, value::ReprValue};
+use magnus::{Error, IntoValue, RArray, Ruby, Value, value::ReprValue};
 use polars::prelude::*;
 use polars::series::IsSorted;
 use polars_core::utils::flatten::flatten_series;
@@ -6,7 +6,7 @@ use polars_core::utils::flatten::flatten_series;
 use crate::conversion::*;
 use crate::exceptions::RbIndexError;
 use crate::rb_modules;
-use crate::{RbDataFrame, RbPolarsErr, RbResult, RbSeries};
+use crate::{RbDataFrame, RbErr, RbPolarsErr, RbResult, RbSeries};
 
 impl RbSeries {
     pub fn struct_unnest(&self) -> RbResult<RbDataFrame> {
@@ -16,7 +16,6 @@ impl RbSeries {
         Ok(df.into())
     }
 
-    // TODO add to Ruby
     pub fn struct_fields(&self) -> RbResult<Vec<String>> {
         let binding = self.series.borrow();
         let ca = binding.struct_().map_err(RbPolarsErr::from)?;
@@ -87,8 +86,8 @@ impl RbSeries {
         }
     }
 
-    pub fn get_index(&self, index: usize) -> RbResult<Value> {
-        let binding = self.series.borrow();
+    pub fn get_index(ruby: &Ruby, rb_self: &Self, index: usize) -> RbResult<Value> {
+        let binding = rb_self.series.borrow();
         let av = match binding.get(index) {
             Ok(v) => v,
             Err(PolarsError::OutOfBounds(err)) => {
@@ -100,26 +99,26 @@ impl RbSeries {
         match av {
             AnyValue::List(s) | AnyValue::Array(s, _) => {
                 let rbseries = RbSeries::new(s);
-                rb_modules::utils().funcall("wrap_s", (rbseries,))
+                rb_modules::pl_utils().funcall("wrap_s", (rbseries,))
             }
-            _ => Ok(Wrap(av).into_value()),
+            _ => Ok(Wrap(av).into_value_with(ruby)),
         }
     }
 
-    pub fn get_index_signed(&self, index: isize) -> RbResult<Value> {
+    pub fn get_index_signed(ruby: &Ruby, rb_self: &Self, index: isize) -> RbResult<Value> {
         let index = if index < 0 {
-            match self.len().checked_sub(index.unsigned_abs()) {
+            match rb_self.len().checked_sub(index.unsigned_abs()) {
                 Some(v) => v,
                 None => {
                     return Err(RbIndexError::new_err(
-                        polars_err!(oob = index, self.len()).to_string(),
+                        polars_err!(oob = index, rb_self.len()).to_string(),
                     ));
                 }
             }
         } else {
             usize::try_from(index).unwrap()
         };
-        self.get_index(index)
+        Self::get_index(ruby, rb_self, index)
     }
 
     pub fn bitand(&self, other: &RbSeries) -> RbResult<Self> {
@@ -149,16 +148,17 @@ impl RbSeries {
         self.series.borrow_mut().rename(name.into());
     }
 
-    pub fn dtype(&self) -> Value {
-        Wrap(self.series.borrow().dtype().clone()).into_value()
+    pub fn dtype(ruby: &Ruby, rb_self: &Self) -> Value {
+        Wrap(rb_self.series.borrow().dtype().clone()).into_value_with(ruby)
     }
 
-    pub fn inner_dtype(&self) -> Option<Value> {
-        self.series
+    pub fn inner_dtype(ruby: &Ruby, rb_self: &Self) -> Option<Value> {
+        rb_self
+            .series
             .borrow()
             .dtype()
             .inner_dtype()
-            .map(|dt| Wrap(dt.clone()).into_value())
+            .map(|dt| Wrap(dt.clone()).into_value_with(ruby))
     }
 
     pub fn set_sorted_flag(&self, descending: bool) -> Self {
@@ -175,11 +175,11 @@ impl RbSeries {
         self.series.borrow().n_chunks()
     }
 
-    pub fn append(&self, other: &RbSeries) -> RbResult<()> {
-        let mut binding = self.series.borrow_mut();
+    pub fn append(ruby: &Ruby, rb_self: &Self, other: &RbSeries) -> RbResult<()> {
+        let mut binding = rb_self.series.borrow_mut();
         let res = binding.append(&other.series.borrow());
         if let Err(e) = res {
-            Err(Error::new(exception::runtime_error(), e.to_string()))
+            Err(Error::new(ruby.exception_runtime_error(), e.to_string()))
         } else {
             Ok(())
         }
@@ -193,22 +193,30 @@ impl RbSeries {
         Ok(())
     }
 
-    pub fn new_from_index(&self, index: usize, length: usize) -> RbResult<Self> {
-        if index >= self.series.borrow().len() {
-            Err(Error::new(exception::arg_error(), "index is out of bounds"))
+    pub fn new_from_index(
+        ruby: &Ruby,
+        rb_self: &Self,
+        index: usize,
+        length: usize,
+    ) -> RbResult<Self> {
+        if index >= rb_self.series.borrow().len() {
+            Err(Error::new(
+                ruby.exception_arg_error(),
+                "index is out of bounds",
+            ))
         } else {
-            Ok(self.series.borrow().new_from_index(index, length).into())
+            Ok(rb_self.series.borrow().new_from_index(index, length).into())
         }
     }
 
-    pub fn filter(&self, filter: &RbSeries) -> RbResult<Self> {
+    pub fn filter(ruby: &Ruby, rb_self: &Self, filter: &RbSeries) -> RbResult<Self> {
         let filter_series = &filter.series.borrow();
         if let Ok(ca) = filter_series.bool() {
-            let series = self.series.borrow().filter(ca).unwrap();
+            let series = rb_self.series.borrow().filter(ca).unwrap();
             Ok(series.into())
         } else {
             Err(Error::new(
-                exception::runtime_error(),
+                ruby.exception_runtime_error(),
                 "Expected a boolean mask".to_string(),
             ))
         }
@@ -319,6 +327,35 @@ impl RbSeries {
         Ok((!bool).into_series().into())
     }
 
+    pub fn shrink_dtype(&self) -> RbResult<Self> {
+        self.series
+            .borrow()
+            .shrink_type()
+            .map(Into::into)
+            .map_err(RbPolarsErr::from)
+            .map_err(RbErr::from)
+    }
+
+    pub fn str_to_decimal_infer(&self, inference_length: usize) -> RbResult<Self> {
+        let s = self.series.borrow();
+        let ca = s.str().map_err(RbPolarsErr::from)?;
+        ca.to_decimal_infer(inference_length)
+            .map(Into::into)
+            .map_err(RbPolarsErr::from)
+            .map_err(RbErr::from)
+    }
+
+    pub fn str_json_decode(&self, infer_schema_length: Option<usize>) -> RbResult<Self> {
+        let lock = self.series.borrow();
+        lock.str()
+            .map_err(RbPolarsErr::from)?
+            .json_decode(None, infer_schema_length)
+            .map(|s| s.with_name(lock.name().clone()))
+            .map(Into::into)
+            .map_err(RbPolarsErr::from)
+            .map_err(RbErr::from)
+    }
+
     pub fn to_s(&self) -> String {
         format!("{}", self.series.borrow())
     }
@@ -404,11 +441,12 @@ impl RbSeries {
         Ok(out.into())
     }
 
-    pub fn get_chunks(&self) -> RbResult<RArray> {
-        flatten_series(&self.series.borrow())
-            .into_iter()
-            .map(|s| rb_modules::utils().funcall::<_, _, Value>("wrap_s", (Self::new(s),)))
-            .collect()
+    pub fn get_chunks(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        ruby.ary_try_from_iter(
+            flatten_series(&rb_self.series.borrow())
+                .into_iter()
+                .map(|s| rb_modules::pl_utils().funcall::<_, _, Value>("wrap_s", (Self::new(s),))),
+        )
     }
 
     pub fn is_sorted(&self, descending: bool, nulls_last: bool) -> RbResult<bool> {

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

@@ -8,13 +8,14 @@ mod import;
 mod map;
 mod scatter;
 
-use magnus::{RArray, prelude::*};
+use magnus::{DataTypeFunctions, RArray, Ruby, TypedData, gc, prelude::*};
 use polars::prelude::*;
 use std::cell::RefCell;
 
-use crate::RbResult;
+use crate::{ObjectValue, RbResult};
 
-#[magnus::wrap(class = "Polars::RbSeries")]
+#[derive(TypedData)]
+#[magnus(class = "Polars::RbSeries", mark)]
 pub struct RbSeries {
     pub series: RefCell<Series>,
 }
@@ -42,9 +43,33 @@ pub fn to_series(rs: RArray) -> RbResult<Vec<Series>> {
 }
 
 pub fn to_rbseries(s: Vec<Column>) -> RArray {
-    RArray::from_iter(
+    Ruby::get().unwrap().ary_from_iter(
         s.into_iter()
             .map(|c| c.take_materialized_series())
             .map(RbSeries::new),
     )
 }
+
+pub fn mark_series(marker: &gc::Marker, series: &Series) {
+    if let DataType::Object(_) = series.dtype() {
+        for i in 0..series.len() {
+            let obj: Option<&ObjectValue> = series.get_object(i).map(|any| any.into());
+            if let Some(o) = obj {
+                marker.mark(o.inner);
+            }
+        }
+    }
+}
+
+impl DataTypeFunctions for RbSeries {
+    fn mark(&self, marker: &gc::Marker) {
+        // this is not ideal, as objects will not be marked if unable to borrow
+        // this should never happen, but log for now to avoid panic,
+        // as most series will not use Object datatype
+        if let Ok(s) = &self.series.try_borrow() {
+            mark_series(marker, s);
+        } else {
+            eprintln!("[polars] Could not borrow!");
+        }
+    }
+}

data/lib/polars/array_expr.rb

@@ -38,7 +38,7 @@ module Polars
     # @param offset [Integer]
     #   Start index. Negative indexing is supported.
     # @param length [Integer]
-    #   Length of the slice. If set to `None` (default), the slice is taken to the
+    #   Length of the slice. If set to `nil` (default), the slice is taken to the
     #   end of the list.
     # @param as_array [Boolean]
     #   Return result as a fixed-length `Array`, otherwise as a `List`.