polars-df 0.5.0 → 0.7.0

data/ext/polars/src/{apply → map}/dataframe.rs

@@ -1,4 +1,4 @@
-use magnus::{class, IntoValue, RArray, TryConvert, Value};
+use magnus::{class, prelude::*, typed_data::Obj, IntoValue, RArray, TryConvert, Value};
 use polars::prelude::*;
 use polars_core::frame::row::{rows_to_schema_first_non_null, Row};
 use polars_core::series::SeriesIter;
@@ -34,20 +34,20 @@ pub fn apply_lambda_unknown<'a>(
             null_count += 1;
             continue;
         } else if out.is_kind_of(class::true_class()) || out.is_kind_of(class::false_class()) {
-            let first_value = out.try_convert::<bool>().ok();
+            let first_value = bool::try_convert(out).ok();
             return Ok((
-                RbSeries::new(
+                Obj::wrap(RbSeries::new(
                     apply_lambda_with_bool_out_type(df, lambda, null_count, first_value)
                         .into_series(),
-                )
-                .into(),
+                ))
+                .as_value(),
                 false,
             ));
         } else if out.is_kind_of(class::float()) {
-            let first_value = out.try_convert::<f64>().ok();
+            let first_value = f64::try_convert(out).ok();
 
             return Ok((
-                RbSeries::new(
+                Obj::wrap(RbSeries::new(
                     apply_lambda_with_primitive_out_type::<Float64Type>(
                         df,
                         lambda,
@@ -55,14 +55,14 @@ pub fn apply_lambda_unknown<'a>(
                         first_value,
                     )
                     .into_series(),
-                )
-                .into(),
+                ))
+                .as_value(),
                 false,
             ));
         } else if out.is_kind_of(class::integer()) {
-            let first_value = out.try_convert::<i64>().ok();
+            let first_value = i64::try_convert(out).ok();
             return Ok((
-                RbSeries::new(
+                Obj::wrap(RbSeries::new(
                     apply_lambda_with_primitive_out_type::<Int64Type>(
                         df,
                         lambda,
@@ -70,12 +70,12 @@ pub fn apply_lambda_unknown<'a>(
                         first_value,
                     )
                     .into_series(),
-                )
-                .into(),
+                ))
+                .as_value(),
                 false,
             ));
         // } else if out.is_kind_of(class::string()) {
-        //     let first_value = out.try_convert::<String>().ok();
+        //     let first_value = String::try_convert(out).ok();
         //     return Ok((
         //         RbSeries::new(
         //             apply_lambda_with_utf8_out_type(df, lambda, null_count, first_value)
@@ -85,25 +85,21 @@ pub fn apply_lambda_unknown<'a>(
         //         false,
         //     ));
         } else if out.respond_to("_s", true)? {
-            let rb_rbseries: Value = out.funcall("_s", ()).unwrap();
-            let series = rb_rbseries
-                .try_convert::<&RbSeries>()
-                .unwrap()
-                .series
-                .borrow();
+            let rb_rbseries: Obj<RbSeries> = out.funcall("_s", ()).unwrap();
+            let series = rb_rbseries.series.borrow();
             let dt = series.dtype();
             return Ok((
-                RbSeries::new(
+                Obj::wrap(RbSeries::new(
                     apply_lambda_with_list_out_type(df, lambda, null_count, Some(&series), dt)?
                         .into_series(),
-                )
-                .into(),
+                ))
+                .as_value(),
                 false,
             ));
-        } else if out.try_convert::<Wrap<Row<'a>>>().is_ok() {
-            let first_value = out.try_convert::<Wrap<Row<'a>>>().unwrap().0;
+        } else if Wrap::<Row<'a>>::try_convert(out).is_ok() {
+            let first_value = Wrap::<Row<'a>>::try_convert(out).unwrap().0;
             return Ok((
-                RbDataFrame::from(
+                Obj::wrap(RbDataFrame::from(
                     apply_lambda_with_rows_output(
                         df,
                         lambda,
@@ -112,8 +108,8 @@ pub fn apply_lambda_unknown<'a>(
                         inference_size,
                     )
                     .map_err(RbPolarsErr::from)?,
-                )
-                .into(),
+                ))
+                .as_value(),
                 true,
             ));
         } else if out.is_kind_of(class::array()) {
@@ -143,7 +139,7 @@ where
         let iter = iters.iter_mut().map(|it| Wrap(it.next().unwrap()));
         let tpl = (RArray::from_iter(iter),);
         match lambda.funcall::<_, _, Value>("call", tpl) {
-            Ok(val) => val.try_convert::<T>().ok(),
+            Ok(val) => T::try_convert(val).ok(),
             Err(e) => panic!("ruby function failed {}", e),
         }
     })
@@ -219,8 +215,7 @@ pub fn apply_lambda_with_list_out_type(
             let tpl = (RArray::from_iter(iter),);
             match lambda.funcall::<_, _, Value>("call", tpl) {
                 Ok(val) => match val.funcall::<_, _, Value>("_s", ()) {
-                    Ok(val) => val
-                        .try_convert::<&RbSeries>()
+                    Ok(val) => Obj::<RbSeries>::try_convert(val)
                         .ok()
                         .map(|ps| ps.series.borrow().clone()),
                     Err(_) => {
@@ -257,11 +252,11 @@ pub fn apply_lambda_with_rows_output<'a>(
         let tpl = (RArray::from_iter(iter),);
         match lambda.funcall::<_, _, Value>("call", tpl) {
             Ok(val) => {
-                match val.try_convert::<RArray>().ok() {
+                match RArray::try_convert(val).ok() {
                     Some(tuple) => {
                         row_buf.0.clear();
                         for v in tuple.each() {
-                            let v = v.unwrap().try_convert::<Wrap<AnyValue>>().unwrap().0;
+                            let v = Wrap::<AnyValue>::try_convert(v.unwrap()).unwrap().0;
                             row_buf.0.push(v);
                         }
                         let ptr = &row_buf as *const Row;

data/ext/polars/src/{apply → map}/mod.rs

@@ -2,7 +2,7 @@ pub mod dataframe;
 pub mod lazy;
 pub mod series;
 
-use magnus::{RHash, Value};
+use magnus::{prelude::*, RHash, Value};
 use polars::chunked_array::builder::get_list_builder;
 use polars::prelude::*;
 use polars_core::export::rayon::prelude::*;
@@ -68,7 +68,7 @@ fn iterator_to_struct(
                 }
             }
             Some(dict) => {
-                let dict = dict.try_convert::<RHash>()?;
+                let dict = RHash::try_convert(dict)?;
                 if dict.len() != struct_width {
                     return Err(crate::error::ComputeError::new_err(
                         format!("Cannot create struct type.\n> The struct dtype expects {} fields, but it got a dict with {} fields.", struct_width, dict.len())
@@ -78,7 +78,7 @@ fn iterator_to_struct(
                 // the first item determines the output name
                 todo!()
                 // for ((_, val), field_items) in dict.iter().zip(&mut items) {
-                //     let item = val.try_convert::<Wrap<AnyValue>>()?;
+                //     let item = Wrap::<AnyValue>::try_convert(val)?;
                 //     field_items.push(item.0)
                 // }
             }
@@ -237,15 +237,19 @@ fn iterator_to_list(
     for _ in 0..init_null_count {
         builder.append_null()
     }
-    builder.append_opt_series(first_value);
+    builder
+        .append_opt_series(first_value)
+        .map_err(RbPolarsErr::from)?;
     for opt_val in it {
         match opt_val {
             None => builder.append_null(),
             Some(s) => {
                 if s.len() == 0 && s.dtype() != dt {
-                    builder.append_series(&Series::full_null("", 0, dt))
+                    builder
+                        .append_series(&Series::full_null("", 0, dt))
+                        .unwrap()
                 } else {
-                    builder.append_series(&s)
+                    builder.append_series(&s).map_err(RbPolarsErr::from)?
                 }
             }
         }

data/ext/polars/src/{apply → map}/series.rs

@@ -1,4 +1,4 @@
-use magnus::{class, IntoValue, RHash, TryConvert, Value};
+use magnus::{class, prelude::*, typed_data::Obj, IntoValue, RHash, TryConvert, Value};
 use polars::prelude::*;
 
 use super::*;
@@ -14,12 +14,12 @@ fn infer_and_finish<'a, A: ApplyLambda<'a>>(
     null_count: usize,
 ) -> RbResult<RbSeries> {
     if out.is_kind_of(class::true_class()) || out.is_kind_of(class::false_class()) {
-        let first_value = out.try_convert::<bool>().unwrap();
+        let first_value = bool::try_convert(out).unwrap();
         applyer
             .apply_lambda_with_bool_out_type(lambda, null_count, Some(first_value))
             .map(|ca| ca.into_series().into())
     } else if out.is_kind_of(class::float()) {
-        let first_value = out.try_convert::<f64>().unwrap();
+        let first_value = f64::try_convert(out).unwrap();
         applyer
             .apply_lambda_with_primitive_out_type::<Float64Type>(
                 lambda,
@@ -28,7 +28,7 @@ fn infer_and_finish<'a, A: ApplyLambda<'a>>(
             )
             .map(|ca| ca.into_series().into())
     } else if out.is_kind_of(class::string()) {
-        let first_value = out.try_convert::<String>().unwrap();
+        let first_value = String::try_convert(out).unwrap();
         applyer
             .apply_lambda_with_utf8_out_type(lambda, null_count, Some(first_value.as_str()))
             .map(|ca| ca.into_series().into())
@@ -37,13 +37,13 @@ fn infer_and_finish<'a, A: ApplyLambda<'a>>(
     } else if out.is_kind_of(class::array()) {
         todo!()
     } else if out.is_kind_of(class::hash()) {
-        let first = out.try_convert::<Wrap<AnyValue<'_>>>()?;
+        let first = Wrap::<AnyValue<'_>>::try_convert(out)?;
         applyer.apply_to_struct(lambda, null_count, first.0)
     }
     // this succeeds for numpy ints as well, where checking if it is pyint fails
     // we do this later in the chain so that we don't extract integers from string chars.
-    else if out.try_convert::<i64>().is_ok() {
-        let first_value = out.try_convert::<i64>().unwrap();
+    else if i64::try_convert(out).is_ok() {
+        let first_value = i64::try_convert(out).unwrap();
         applyer
             .apply_lambda_with_primitive_out_type::<Int64Type>(
                 lambda,
@@ -51,7 +51,7 @@ fn infer_and_finish<'a, A: ApplyLambda<'a>>(
                 Some(first_value),
             )
             .map(|ca| ca.into_series().into())
-    } else if let Ok(av) = out.try_convert::<Wrap<AnyValue>>() {
+    } else if let Ok(av) = Wrap::<AnyValue>::try_convert(out) {
         applyer
             .apply_extract_any_values(lambda, null_count, av.0)
             .map(|s| s.into())
@@ -141,7 +141,7 @@ where
     S: TryConvert,
 {
     match call_lambda(lambda, in_val) {
-        Ok(out) => out.try_convert::<S>(),
+        Ok(out) => S::try_convert(out),
         Err(e) => panic!("ruby function failed {}", e),
     }
 }
@@ -151,13 +151,9 @@ where
     T: IntoValue,
 {
     let out: Value = lambda.funcall("call", (in_val,))?;
-    let py_series: Value = out.funcall("_s", ())?;
-    Ok(py_series
-        .try_convert::<&RbSeries>()
-        .unwrap()
-        .series
-        .borrow()
-        .clone())
+    let py_series: Obj<RbSeries> = out.funcall("_s", ())?;
+    let tmp = py_series.series.borrow();
+    Ok(tmp.clone())
 }
 
 impl<'a> ApplyLambda<'a> for BooleanChunked {

data/ext/polars/src/object.rs

@@ -10,7 +10,7 @@ use polars_core::prelude::AnyValue;
 use crate::prelude::ObjectValue;
 use crate::Wrap;
 
-// pub(crate) const OBJECT_NAME: &str = "object";
+pub(crate) const OBJECT_NAME: &str = "object";
 
 pub(crate) fn register_object_builder() {
     if !registry::is_object_builder_registered() {
@@ -21,7 +21,7 @@ pub(crate) fn register_object_builder() {
 
         let object_converter = Arc::new(|av: AnyValue| {
             let object = ObjectValue {
-                inner: Wrap(av).into_value(),
+                inner: Wrap(av).into_value().into(),
             };
             Box::new(object) as Box<dyn Any>
         });

data/ext/polars/src/rb_modules.rs

@@ -1,21 +1,40 @@
-use magnus::{class, memoize, Module, RClass, RModule};
+use magnus::{value::Lazy, Module, RClass, RModule, Ruby};
+
+static POLARS: Lazy<RModule> = Lazy::new(|ruby| ruby.class_object().const_get("Polars").unwrap());
 
 pub(crate) fn polars() -> RModule {
-    *memoize!(RModule: class::object().const_get("Polars").unwrap())
+    Ruby::get().unwrap().get_inner(&POLARS)
 }
 
+static SERIES: Lazy<RClass> =
+    Lazy::new(|ruby| ruby.get_inner(&POLARS).const_get("Series").unwrap());
+
 pub(crate) fn series() -> RClass {
-    *memoize!(RClass: polars().const_get("Series").unwrap())
+    Ruby::get().unwrap().get_inner(&SERIES)
 }
 
+static UTILS: Lazy<RModule> = Lazy::new(|ruby| ruby.get_inner(&POLARS).const_get("Utils").unwrap());
+
 pub(crate) fn utils() -> RModule {
-    *memoize!(RModule: polars().const_get("Utils").unwrap())
+    Ruby::get().unwrap().get_inner(&UTILS)
+}
+
+static BIGDECIMAL: Lazy<RClass> =
+    Lazy::new(|ruby| ruby.class_object().const_get("BigDecimal").unwrap());
+
+pub(crate) fn bigdecimal() -> RClass {
+    Ruby::get().unwrap().get_inner(&BIGDECIMAL)
 }
 
+static DATE: Lazy<RClass> = Lazy::new(|ruby| ruby.class_object().const_get("Date").unwrap());
+
 pub(crate) fn date() -> RClass {
-    *memoize!(RClass: class::object().const_get("Date").unwrap())
+    Ruby::get().unwrap().get_inner(&DATE)
 }
 
+static DATETIME: Lazy<RClass> =
+    Lazy::new(|ruby| ruby.class_object().const_get("DateTime").unwrap());
+
 pub(crate) fn datetime() -> RClass {
-    *memoize!(RClass: class::object().const_get("DateTime").unwrap())
+    Ruby::get().unwrap().get_inner(&DATETIME)
 }

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

@@ -1,10 +1,10 @@
-use magnus::RArray;
+use magnus::{prelude::*, RArray};
 use polars_core::prelude::*;
 
-use crate::conversion::{slice_extract_wrapped, Wrap};
+use crate::conversion::{slice_extract_wrapped, vec_extract_wrapped, Wrap};
 use crate::prelude::ObjectValue;
 use crate::series::to_series_collection;
-use crate::{RbPolarsErr, RbResult, RbSeries};
+use crate::{RbPolarsErr, RbResult, RbSeries, RbValueError};
 
 impl RbSeries {
     pub fn new_opt_bool(name: String, obj: RArray, strict: bool) -> RbResult<RbSeries> {
@@ -16,7 +16,7 @@ impl RbSeries {
                 if item.is_nil() {
                     builder.append_null()
                 } else {
-                    match item.try_convert::<bool>() {
+                    match bool::try_convert(*item) {
                         Ok(val) => builder.append_value(val),
                         Err(e) => {
                             if strict {
@@ -49,7 +49,7 @@ where
             if item.is_nil() {
                 builder.append_null()
             } else {
-                match item.try_convert::<T::Native>() {
+                match T::Native::try_convert(*item) {
                     Ok(val) => builder.append_value(val),
                     Err(e) => {
                         if strict {
@@ -92,7 +92,7 @@ init_method_opt!(new_opt_f64, Float64Type, f64);
 fn vec_wrap_any_value<'s>(arr: RArray) -> RbResult<Vec<Wrap<AnyValue<'s>>>> {
     let mut val = Vec::with_capacity(arr.len());
     for v in arr.each() {
-        val.push(v?.try_convert()?);
+        val.push(Wrap::<AnyValue<'s>>::try_convert(v?)?);
     }
     Ok(val)
 }
@@ -137,6 +137,32 @@ impl RbSeries {
         Ok(Series::new(&name, &series_vec).into())
     }
 
+    pub fn new_array(
+        width: usize,
+        inner: Option<Wrap<DataType>>,
+        name: String,
+        val: RArray,
+        _strict: bool,
+    ) -> RbResult<Self> {
+        let val = vec_wrap_any_value(val)?;
+        let val = vec_extract_wrapped(val);
+        let out = Series::new(&name, &val);
+        match out.dtype() {
+            DataType::List(list_inner) => {
+                let out = out
+                    .cast(&DataType::Array(
+                        Box::new(inner.map(|dt| dt.0).unwrap_or(*list_inner.clone())),
+                        width,
+                    ))
+                    .map_err(RbPolarsErr::from)?;
+                Ok(out.into())
+            }
+            _ => Err(RbValueError::new_err(
+                "could not create Array from input".to_string(),
+            )),
+        }
+    }
+
     pub fn new_decimal(name: String, val: RArray, strict: bool) -> RbResult<Self> {
         let val = vec_wrap_any_value(val)?;
         // TODO: do we have to respect 'strict' here? it's possible if we want to

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

@@ -1,4 +1,4 @@
-use magnus::{class, Module, RArray, RClass, RModule, Value};
+use magnus::{class, prelude::*, Module, RArray, RClass, RModule, Value};
 use polars_core::prelude::*;
 
 use crate::{raise_err, RbPolarsErr, RbResult, RbSeries};
@@ -13,7 +13,7 @@ impl RbSeries {
                 let ca = s.utf8().unwrap();
 
                 // TODO make more efficient
-                let np_arr = RArray::from_iter(ca.into_iter());
+                let np_arr = RArray::from_iter(ca);
                 class::object()
                     .const_get::<_, RModule>("Numo")?
                     .const_get::<_, RClass>("RObject")?

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

@@ -1,18 +1,34 @@
-// use polars::export::arrow2::array::Array;
+use polars::export::arrow::array::Array;
 use polars::prelude::*;
 
-pub fn set_at_idx(mut s: Series, idx: &Series, values: &Series) -> PolarsResult<Series> {
+use crate::error::RbPolarsErr;
+use crate::{RbResult, RbSeries};
+
+impl RbSeries {
+    pub fn set_at_idx(&self, idx: &RbSeries, values: &RbSeries) -> RbResult<()> {
+        let mut s = self.series.borrow_mut();
+        match set_at_idx(s.clone(), &idx.series.borrow(), &values.series.borrow()) {
+            Ok(out) => {
+                *s = out;
+                Ok(())
+            }
+            Err(e) => Err(RbPolarsErr::from(e)),
+        }
+    }
+}
+
+fn set_at_idx(mut s: Series, idx: &Series, values: &Series) -> PolarsResult<Series> {
     let logical_dtype = s.dtype().clone();
     let idx = idx.cast(&IDX_DTYPE)?;
     let idx = idx.rechunk();
     let idx = idx.idx().unwrap();
     let idx = idx.downcast_iter().next().unwrap();
 
-    // if idx.null_count() > 0 {
-    //     return Err(PolarsError::ComputeError(
-    //         "index values should not be null".into(),
-    //     ));
-    // }
+    if idx.null_count() > 0 {
+        return Err(PolarsError::ComputeError(
+            "index values should not be null".into(),
+        ));
+    }
 
     let idx = idx.values().as_slice();
 
@@ -27,52 +43,52 @@ pub fn set_at_idx(mut s: Series, idx: &Series, values: &Series) -> PolarsResult<
         DataType::Int8 => {
             let ca: &mut ChunkedArray<Int8Type> = mutable_s.as_mut();
             let values = values.i8()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::Int16 => {
             let ca: &mut ChunkedArray<Int16Type> = mutable_s.as_mut();
             let values = values.i16()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::Int32 => {
             let ca: &mut ChunkedArray<Int32Type> = mutable_s.as_mut();
             let values = values.i32()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::Int64 => {
             let ca: &mut ChunkedArray<Int64Type> = mutable_s.as_mut();
             let values = values.i64()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::UInt8 => {
             let ca: &mut ChunkedArray<UInt8Type> = mutable_s.as_mut();
             let values = values.u8()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::UInt16 => {
             let ca: &mut ChunkedArray<UInt16Type> = mutable_s.as_mut();
             let values = values.u16()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::UInt32 => {
             let ca: &mut ChunkedArray<UInt32Type> = mutable_s.as_mut();
             let values = values.u32()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::UInt64 => {
             let ca: &mut ChunkedArray<UInt64Type> = mutable_s.as_mut();
             let values = values.u64()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::Float32 => {
             let ca: &mut ChunkedArray<Float32Type> = mutable_s.as_mut();
             let values = values.f32()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::Float64 => {
             let ca: &mut ChunkedArray<Float64Type> = mutable_s.as_mut();
             let values = values.f64()?;
-            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+            std::mem::take(ca).set_at_idx2(idx, values)
         }
         DataType::Boolean => {
             let ca = s.bool()?;