polars-df 0.6.0 → 0.8.0

data/ext/polars/src/{series.rs → series/mod.rs}

@@ -3,17 +3,16 @@ mod arithmetic;
 mod comparison;
 mod construction;
 mod export;
-mod set_at_idx;
+mod scatter;
 
-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)
 }
@@ -81,7 +80,7 @@ impl RbSeries {
 
     pub fn get_fmt(&self, index: usize, str_lengths: usize) -> String {
         let val = format!("{}", self.series.borrow().get(index).unwrap());
-        if let DataType::Utf8 | DataType::Categorical(_) = self.series.borrow().dtype() {
+        if let DataType::String | DataType::Categorical(_, _) = self.series.borrow().dtype() {
             let v_trunc = &val[..val
                 .char_indices()
                 .take(str_lengths)
@@ -91,7 +90,7 @@ impl RbSeries {
             if val == v_trunc {
                 val
             } else {
-                format!("{}...", v_trunc)
+                format!("{}…", v_trunc)
             }
         } else {
             val
@@ -274,15 +273,13 @@ impl RbSeries {
         Ok(s.into())
     }
 
-    pub fn series_equal(&self, other: &RbSeries, null_equal: bool, strict: bool) -> bool {
+    pub fn equals(&self, other: &RbSeries, null_equal: bool, strict: bool) -> bool {
         if strict {
             self.series.borrow().eq(&other.series.borrow())
         } else if null_equal {
-            self.series
-                .borrow()
-                .series_equal_missing(&other.series.borrow())
+            self.series.borrow().equals_missing(&other.series.borrow())
         } else {
-            self.series.borrow().series_equal(&other.series.borrow())
+            self.series.borrow().equals(&other.series.borrow())
         }
     }
 
@@ -316,16 +313,16 @@ impl RbSeries {
                 DataType::Int64 => RArray::from_iter(series.i64().unwrap()).into_value(),
                 DataType::Float32 => RArray::from_iter(series.f32().unwrap()).into_value(),
                 DataType::Float64 => RArray::from_iter(series.f64().unwrap()).into_value(),
-                DataType::Categorical(_) => {
+                DataType::Categorical(_, _) => {
                     RArray::from_iter(series.categorical().unwrap().iter_str()).into_value()
                 }
-                DataType::Object(_) => {
+                DataType::Object(_, _) => {
                     let v = RArray::with_capacity(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_object()).unwrap(),
-                            None => v.push(QNIL).unwrap(),
+                            None => v.push(qnil()).unwrap(),
                         };
                     }
                     v.into_value()
@@ -333,10 +330,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 +349,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());
@@ -378,8 +375,8 @@ impl RbSeries {
                     let ca = series.decimal().unwrap();
                     return Wrap(ca).into_value();
                 }
-                DataType::Utf8 => {
-                    let ca = series.utf8().unwrap();
+                DataType::String => {
+                    let ca = series.str().unwrap();
                     return Wrap(ca).into_value();
                 }
                 DataType::Struct(_) => {
@@ -443,7 +440,7 @@ impl RbSeries {
 
         macro_rules! dispatch_apply {
             ($self:expr, $method:ident, $($args:expr),*) => {
-                if matches!($self.dtype(), DataType::Object(_)) {
+                if matches!($self.dtype(), DataType::Object(_, _)) {
                     // let ca = $self.0.unpack::<ObjectType<ObjectValue>>().unwrap();
                     // ca.$method($($args),*)
                     todo!()
@@ -464,7 +461,7 @@ impl RbSeries {
             DataType::Datetime(_, _)
                 | DataType::Date
                 | DataType::Duration(_)
-                | DataType::Categorical(_)
+                | DataType::Categorical(_, _)
                 | DataType::Time
         ) || !skip_nulls
         {
@@ -605,12 +602,12 @@ impl RbSeries {
                 )?;
                 ca.into_datetime(tu, tz).into_series()
             }
-            Some(DataType::Utf8) => {
+            Some(DataType::String) => {
                 let ca = dispatch_apply!(series, apply_lambda_with_utf8_out_type, lambda, 0, None)?;
 
                 ca.into_series()
             }
-            Some(DataType::Object(_)) => {
+            Some(DataType::Object(_, _)) => {
                 let ca =
                     dispatch_apply!(series, apply_lambda_with_object_out_type, lambda, 0, None)?;
                 ca.into_series()
@@ -643,14 +640,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)
@@ -665,8 +654,13 @@ impl RbSeries {
         self.series.borrow_mut().shrink_to_fit();
     }
 
-    pub fn dot(&self, other: &RbSeries) -> Option<f64> {
-        self.series.borrow().dot(&other.series.borrow())
+    pub fn dot(&self, other: &RbSeries) -> RbResult<f64> {
+        let out = self
+            .series
+            .borrow()
+            .dot(&other.series.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(out)
     }
 
     pub fn skew(&self, bias: bool) -> RbResult<Option<f64>> {
@@ -708,17 +702,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/series/{set_at_idx.rs → scatter.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 scatter(&self, idx: &RbSeries, values: &RbSeries) -> RbResult<()> {
+        let mut s = self.series.borrow_mut();
+        match scatter(s.clone(), &idx.series.borrow(), &values.series.borrow()) {
+            Ok(out) => {
+                *s = out;
+                Ok(())
+            }
+            Err(e) => Err(RbPolarsErr::from(e)),
+        }
+    }
+}
+
+fn scatter(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,62 +43,62 @@ 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).scatter(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).scatter(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).scatter(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).scatter(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).scatter(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).scatter(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).scatter(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).scatter(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).scatter(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).scatter(idx, values)
         }
         DataType::Boolean => {
             let ca = s.bool()?;
             let values = values.bool()?;
-            ca.set_at_idx2(idx, values)
+            ca.scatter(idx, values)
         }
-        DataType::Utf8 => {
-            let ca = s.utf8()?;
-            let values = values.utf8()?;
-            ca.set_at_idx2(idx, values)
+        DataType::String => {
+            let ca = s.str()?;
+            let values = values.str()?;
+            ca.scatter(idx, values)
         }
         _ => panic!("not yet implemented for dtype: {}", logical_dtype),
     };

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)
+    }
+}

data/ext/polars/src/utils.rs

@@ -23,7 +23,7 @@ macro_rules! apply_method_all_arrow_series2 {
     ($self:expr, $method:ident, $($args:expr),*) => {
         match $self.dtype() {
             DataType::Boolean => $self.bool().unwrap().$method($($args),*),
-            DataType::Utf8 => $self.utf8().unwrap().$method($($args),*),
+            DataType::String => $self.str().unwrap().$method($($args),*),
             DataType::UInt8 => $self.u8().unwrap().$method($($args),*),
             DataType::UInt16 => $self.u16().unwrap().$method($($args),*),
             DataType::UInt32 => $self.u32().unwrap().$method($($args),*),