polars-df 0.6.0 → 0.7.0

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

@@ -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,25 +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 {
-    *memoize!(RClass: class::object().const_get("BigDecimal").unwrap())
+    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,4 +1,4 @@
-use magnus::RArray;
+use magnus::{prelude::*, RArray};
 use polars_core::prelude::*;
 
 use crate::conversion::{slice_extract_wrapped, vec_extract_wrapped, Wrap};
@@ -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)
 }

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()?;

data/ext/polars/src/series.rs

@@ -5,15 +5,14 @@ mod construction;
 mod export;
 mod set_at_idx;
 
-use magnus::{exception, Error, IntoValue, RArray, Value, QNIL};
+use magnus::{exception, prelude::*, value::qnil, Error, IntoValue, RArray, Value};
 use polars::prelude::*;
 use polars::series::IsSorted;
 use std::cell::RefCell;
 
-use crate::apply::series::{call_lambda_and_extract, ApplyLambda};
 use crate::apply_method_all_arrow_series2;
 use crate::conversion::*;
-use crate::series::set_at_idx::set_at_idx;
+use crate::map::series::{call_lambda_and_extract, ApplyLambda};
 use crate::{RbDataFrame, RbPolarsErr, RbResult};
 
 #[magnus::wrap(class = "Polars::RbSeries")]
@@ -38,7 +37,7 @@ impl RbSeries {
 pub fn to_series_collection(rs: RArray) -> RbResult<Vec<Series>> {
     let mut series = Vec::new();
     for item in rs.each() {
-        series.push(item?.try_convert::<&RbSeries>()?.series.borrow().clone());
+        series.push(<&RbSeries>::try_convert(item?)?.series.borrow().clone());
     }
     Ok(series)
 }
@@ -325,7 +324,7 @@ impl RbSeries {
                         let obj: Option<&ObjectValue> = series.get_object(i).map(|any| any.into());
                         match obj {
                             Some(val) => v.push(val.to_object()).unwrap(),
-                            None => v.push(QNIL).unwrap(),
+                            None => v.push(qnil()).unwrap(),
                         };
                     }
                     v.into_value()
@@ -333,10 +332,10 @@ impl RbSeries {
                 DataType::List(_) => {
                     let v = RArray::new();
                     let ca = series.list().unwrap();
-                    for opt_s in ca.amortized_iter() {
+                    for opt_s in unsafe { ca.amortized_iter() } {
                         match opt_s {
                             None => {
-                                v.push(QNIL).unwrap();
+                                v.push(qnil()).unwrap();
                             }
                             Some(s) => {
                                 let rblst = to_a_recursive(s.as_ref());
@@ -352,7 +351,7 @@ impl RbSeries {
                     for opt_s in ca.amortized_iter() {
                         match opt_s {
                             None => {
-                                v.push(QNIL).unwrap();
+                                v.push(qnil()).unwrap();
                             }
                             Some(s) => {
                                 let rblst = to_a_recursive(s.as_ref());
@@ -643,14 +642,6 @@ impl RbSeries {
         Ok(df.into())
     }
 
-    pub fn peak_max(&self) -> Self {
-        self.series.borrow().peak_max().into_series().into()
-    }
-
-    pub fn peak_min(&self) -> Self {
-        self.series.borrow().peak_min().into_series().into()
-    }
-
     pub fn n_unique(&self) -> RbResult<usize> {
         let n = self.series.borrow().n_unique().map_err(RbPolarsErr::from)?;
         Ok(n)
@@ -708,17 +699,6 @@ impl RbSeries {
             None
         }
     }
-
-    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)),
-        }
-    }
 }
 
 macro_rules! impl_set_with_mask {

data/ext/polars/src/sql.rs

@@ -0,0 +1,46 @@
+use polars::sql::SQLContext;
+use std::cell::RefCell;
+
+use crate::{RbLazyFrame, RbPolarsErr, RbResult};
+
+#[magnus::wrap(class = "Polars::RbSQLContext")]
+#[repr(transparent)]
+#[derive(Clone)]
+pub struct RbSQLContext {
+    pub context: RefCell<SQLContext>,
+}
+
+#[allow(
+    clippy::wrong_self_convention,
+    clippy::should_implement_trait,
+    clippy::len_without_is_empty
+)]
+impl RbSQLContext {
+    #[allow(clippy::new_without_default)]
+    pub fn new() -> RbSQLContext {
+        RbSQLContext {
+            context: SQLContext::new().into(),
+        }
+    }
+
+    pub fn execute(&self, query: String) -> RbResult<RbLazyFrame> {
+        Ok(self
+            .context
+            .borrow_mut()
+            .execute(&query)
+            .map_err(RbPolarsErr::from)?
+            .into())
+    }
+
+    pub fn get_tables(&self) -> RbResult<Vec<String>> {
+        Ok(self.context.borrow().get_tables())
+    }
+
+    pub fn register(&self, name: String, lf: &RbLazyFrame) {
+        self.context.borrow_mut().register(&name, lf.ldf.clone())
+    }
+
+    pub fn unregister(&self, name: String) {
+        self.context.borrow_mut().unregister(&name)
+    }
+}