polars-df 0.2.5 → 0.3.1

data/ext/polars/src/error.rs

@@ -1,4 +1,4 @@
-use magnus::exception::arg_error;
+use magnus::exception;
 use magnus::Error;
 use polars::error::ArrowError;
 use polars::prelude::PolarsError;
@@ -8,23 +8,23 @@ pub struct RbPolarsErr {}
 impl RbPolarsErr {
     // convert to Error instead of Self
     pub fn from(e: PolarsError) -> Error {
-        Error::runtime_error(e.to_string())
+        Error::new(exception::runtime_error(), e.to_string())
     }
 
     pub fn arrow(e: ArrowError) -> Error {
-        Error::runtime_error(e.to_string())
+        Error::new(exception::runtime_error(), e.to_string())
     }
 
     pub fn io(e: std::io::Error) -> Error {
-        Error::runtime_error(e.to_string())
+        Error::new(exception::runtime_error(), e.to_string())
     }
 
     pub fn other(message: String) -> Error {
-        Error::runtime_error(message)
+        Error::new(exception::runtime_error(), message)
     }
 
     pub fn todo() -> Error {
-        Error::runtime_error("not implemented yet")
+        Error::new(exception::runtime_error(), "not implemented yet")
     }
 }
 
@@ -32,7 +32,7 @@ pub struct RbValueError {}
 
 impl RbValueError {
     pub fn new_err(message: String) -> Error {
-        Error::new(arg_error(), message)
+        Error::new(exception::arg_error(), message)
     }
 }
 
@@ -40,6 +40,14 @@ pub struct ComputeError {}
 
 impl ComputeError {
     pub fn new_err(message: String) -> Error {
-        Error::runtime_error(message)
+        Error::new(exception::runtime_error(), message)
     }
 }
+
+#[macro_export]
+macro_rules! raise_err(
+    ($msg:expr, $err:ident) => {{
+        Err(PolarsError::$err($msg.into())).map_err(RbPolarsErr::from)?;
+        unreachable!()
+    }}
+);

data/ext/polars/src/file.rs

@@ -1,4 +1,4 @@
-use magnus::{Error, RString, Value};
+use magnus::{exception, Error, RString, Value};
 use polars::io::mmap::MmapBytesReader;
 use std::fs::File;
 use std::io::Cursor;
@@ -9,9 +9,10 @@ use crate::RbResult;
 pub fn get_file_like(f: Value, truncate: bool) -> RbResult<File> {
     let str_slice = f.try_convert::<PathBuf>()?;
     let f = if truncate {
-        File::create(str_slice).map_err(|e| Error::runtime_error(e.to_string()))?
+        File::create(str_slice)
+            .map_err(|e| Error::new(exception::runtime_error(), e.to_string()))?
     } else {
-        File::open(str_slice).map_err(|e| Error::runtime_error(e.to_string()))?
+        File::open(str_slice).map_err(|e| Error::new(exception::runtime_error(), e.to_string()))?
     };
     Ok(f)
 }
@@ -23,7 +24,7 @@ pub fn get_mmap_bytes_reader(rb_f: Value) -> RbResult<Box<dyn MmapBytesReader>>
         Ok(Box::new(Cursor::new(bytes.to_vec())))
     } else {
         let p = rb_f.try_convert::<PathBuf>()?;
-        let f = File::open(p).map_err(|e| Error::runtime_error(e.to_string()))?;
+        let f = File::open(p).map_err(|e| Error::new(exception::runtime_error(), e.to_string()))?;
         Ok(Box::new(f))
     }
 }

data/ext/polars/src/lazy/apply.rs

@@ -4,7 +4,7 @@ use polars::prelude::*;
 use crate::lazy::dsl::RbExpr;
 use crate::Wrap;
 
-pub fn binary_lambda(_lambda: Value, _a: Series, _b: Series) -> PolarsResult<Series> {
+pub fn binary_lambda(_lambda: Value, _a: Series, _b: Series) -> PolarsResult<Option<Series>> {
     todo!();
 }
 

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

@@ -1,4 +1,4 @@
-use magnus::{RArray, RHash, Value};
+use magnus::{IntoValue, RArray, RHash, Value};
 use polars::io::RowCount;
 use polars::lazy::frame::{LazyFrame, LazyGroupBy};
 use polars::prelude::*;
@@ -140,7 +140,7 @@ impl RbLazyFrame {
             .with_infer_schema_length(infer_schema_length)
             .with_delimiter(delimiter)
             .has_header(has_header)
-            .with_ignore_parser_errors(ignore_errors)
+            .with_ignore_errors(ignore_errors)
             .with_skip_rows(skip_rows)
             .with_n_rows(n_rows)
             .with_cache(cache)
@@ -180,6 +180,8 @@ impl RbLazyFrame {
             rechunk,
             row_count,
             low_memory,
+            // TODO support cloud options
+            cloud_options: None,
         };
         let lf = LazyFrame::scan_parquet(path, args).map_err(RbPolarsErr::from)?;
         Ok(lf.into())
@@ -254,6 +256,7 @@ impl RbLazyFrame {
             SortOptions {
                 descending: reverse,
                 nulls_last,
+                multithreaded: true,
             },
         )
         .into()
@@ -597,10 +600,10 @@ impl RbLazyFrame {
         Ok(self.get_schema()?.iter_names().cloned().collect())
     }
 
-    pub fn dtypes(&self) -> RbResult<Vec<Value>> {
+    pub fn dtypes(&self) -> RbResult<RArray> {
         let schema = self.get_schema()?;
-        let iter = schema.iter_dtypes().map(|dt| Wrap(dt.clone()).into());
-        Ok(iter.collect())
+        let iter = schema.iter_dtypes().map(|dt| Wrap(dt.clone()).into_value());
+        Ok(RArray::from_iter(iter))
     }
 
     pub fn schema(&self) -> RbResult<RHash> {
@@ -610,7 +613,10 @@ impl RbLazyFrame {
         schema.iter_fields().for_each(|fld| {
             // TODO remove unwrap
             schema_dict
-                .aset::<String, Value>(fld.name().clone(), Wrap(fld.data_type().clone()).into())
+                .aset::<String, Value>(
+                    fld.name().clone(),
+                    Wrap(fld.data_type().clone()).into_value(),
+                )
                 .unwrap();
         });
         Ok(schema_dict)

data/ext/polars/src/lazy/dsl.rs

@@ -1,5 +1,5 @@
 use magnus::block::Proc;
-use magnus::{class, RArray, RString, Value};
+use magnus::{class, IntoValue, RArray, RString, Value};
 use polars::chunked_array::ops::SortOptions;
 use polars::lazy::dsl;
 use polars::lazy::dsl::Operator;
@@ -214,6 +214,7 @@ impl RbExpr {
             .sort_with(SortOptions {
                 descending,
                 nulls_last,
+                multithreaded: true,
             })
             .into()
     }
@@ -224,6 +225,7 @@ impl RbExpr {
             .arg_sort(SortOptions {
                 descending: reverse,
                 nulls_last,
+                multithreaded: true,
             })
             .into()
     }
@@ -240,10 +242,10 @@ impl RbExpr {
         self.clone().inner.arg_min().into()
     }
 
-    pub fn search_sorted(&self, element: &RbExpr) -> Self {
+    pub fn search_sorted(&self, element: &RbExpr, side: Wrap<SearchSortedSide>) -> Self {
         self.inner
             .clone()
-            .search_sorted(element.inner.clone())
+            .search_sorted(element.inner.clone(), side.0)
             .into()
     }
 
@@ -287,7 +289,10 @@ impl RbExpr {
         Ok(self
             .inner
             .clone()
-            .apply(move |s| s.fill_null(strat), GetOutput::same_type())
+            .apply(
+                move |s| s.fill_null(strat).map(Some),
+                GetOutput::same_type(),
+            )
             .with_fmt("fill_null_with_strategy")
             .into())
     }
@@ -335,7 +340,10 @@ impl RbExpr {
     pub fn take_every(&self, n: usize) -> Self {
         self.clone()
             .inner
-            .map(move |s: Series| Ok(s.take_every(n)), GetOutput::same_type())
+            .map(
+                move |s: Series| Ok(Some(s.take_every(n))),
+                GetOutput::same_type(),
+            )
             .with_fmt("take_every")
             .into()
     }
@@ -365,7 +373,7 @@ impl RbExpr {
     pub fn rechunk(&self) -> Self {
         self.inner
             .clone()
-            .map(|s| Ok(s.rechunk()), GetOutput::same_type())
+            .map(|s| Ok(Some(s.rechunk())), GetOutput::same_type())
             .into()
     }
 
@@ -527,6 +535,7 @@ impl RbExpr {
                 exact,
                 cache,
                 tz_aware: false,
+                utc: false,
             })
             .into()
     }
@@ -538,6 +547,7 @@ impl RbExpr {
         exact: bool,
         cache: bool,
         tz_aware: bool,
+        utc: bool,
     ) -> Self {
         let tu = match fmt {
             Some(ref fmt) => {
@@ -565,6 +575,7 @@ impl RbExpr {
                 exact,
                 cache,
                 tz_aware,
+                utc,
             })
             .into()
     }
@@ -586,26 +597,27 @@ impl RbExpr {
                 exact,
                 cache,
                 tz_aware: false,
+                utc: false,
             })
             .into()
     }
 
-    pub fn str_strip(&self, matches: Option<char>) -> Self {
+    pub fn str_strip(&self, matches: Option<String>) -> Self {
         self.inner.clone().str().strip(matches).into()
     }
 
-    pub fn str_rstrip(&self, matches: Option<char>) -> Self {
+    pub fn str_rstrip(&self, matches: Option<String>) -> Self {
         self.inner.clone().str().rstrip(matches).into()
     }
 
-    pub fn str_lstrip(&self, matches: Option<char>) -> Self {
+    pub fn str_lstrip(&self, matches: Option<String>) -> Self {
         self.inner.clone().str().lstrip(matches).into()
     }
 
     pub fn str_slice(&self, start: i64, length: Option<u64>) -> Self {
         let function = move |s: Series| {
             let ca = s.utf8()?;
-            Ok(ca.str_slice(start, length)?.into_series())
+            Ok(Some(ca.str_slice(start, length)?.into_series()))
         };
         self.clone()
             .inner
@@ -625,7 +637,7 @@ impl RbExpr {
     pub fn str_lengths(&self) -> Self {
         let function = |s: Series| {
             let ca = s.utf8()?;
-            Ok(ca.str_lengths().into_series())
+            Ok(Some(ca.str_lengths().into_series()))
         };
         self.clone()
             .inner
@@ -637,7 +649,7 @@ impl RbExpr {
     pub fn str_n_chars(&self) -> Self {
         let function = |s: Series| {
             let ca = s.utf8()?;
-            Ok(ca.str_n_chars().into_series())
+            Ok(Some(ca.str_n_chars().into_series()))
         };
         self.clone()
             .inner
@@ -674,37 +686,51 @@ impl RbExpr {
         self.clone().inner.str().rjust(width, fillchar).into()
     }
 
-    pub fn str_contains(&self, pat: String, literal: Option<bool>) -> Self {
+    pub fn str_contains(&self, pat: &RbExpr, literal: Option<bool>, strict: bool) -> Self {
         match literal {
-            Some(true) => self.inner.clone().str().contains_literal(pat).into(),
-            _ => self.inner.clone().str().contains(pat).into(),
+            Some(true) => self
+                .inner
+                .clone()
+                .str()
+                .contains_literal(pat.inner.clone())
+                .into(),
+            _ => self
+                .inner
+                .clone()
+                .str()
+                .contains(pat.inner.clone(), strict)
+                .into(),
         }
     }
 
-    pub fn str_ends_with(&self, sub: String) -> Self {
-        self.inner.clone().str().ends_with(sub).into()
+    pub fn str_ends_with(&self, sub: &RbExpr) -> Self {
+        self.inner.clone().str().ends_with(sub.inner.clone()).into()
     }
 
-    pub fn str_starts_with(&self, sub: String) -> Self {
-        self.inner.clone().str().starts_with(sub).into()
+    pub fn str_starts_with(&self, sub: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .starts_with(sub.inner.clone())
+            .into()
     }
 
     pub fn str_hex_encode(&self) -> Self {
         self.clone()
             .inner
             .map(
-                move |s| s.utf8().map(|s| s.hex_encode().into_series()),
+                move |s| s.utf8().map(|s| Some(s.hex_encode().into_series())),
                 GetOutput::same_type(),
             )
             .with_fmt("str.hex_encode")
             .into()
     }
 
-    pub fn str_hex_decode(&self, strict: Option<bool>) -> Self {
+    pub fn str_hex_decode(&self, strict: bool) -> Self {
         self.clone()
             .inner
             .map(
-                move |s| s.utf8()?.hex_decode(strict).map(|s| s.into_series()),
+                move |s| s.utf8()?.hex_decode(strict).map(|s| Some(s.into_series())),
                 GetOutput::same_type(),
             )
             .with_fmt("str.hex_decode")
@@ -715,18 +741,22 @@ impl RbExpr {
         self.clone()
             .inner
             .map(
-                move |s| s.utf8().map(|s| s.base64_encode().into_series()),
+                move |s| s.utf8().map(|s| Some(s.base64_encode().into_series())),
                 GetOutput::same_type(),
             )
             .with_fmt("str.base64_encode")
             .into()
     }
 
-    pub fn str_base64_decode(&self, strict: Option<bool>) -> Self {
+    pub fn str_base64_decode(&self, strict: bool) -> Self {
         self.clone()
             .inner
             .map(
-                move |s| s.utf8()?.base64_decode(strict).map(|s| s.into_series()),
+                move |s| {
+                    s.utf8()?
+                        .base64_decode(strict)
+                        .map(|s| Some(s.into_series()))
+                },
                 GetOutput::same_type(),
             )
             .with_fmt("str.base64_decode")
@@ -737,7 +767,7 @@ impl RbExpr {
         let function = move |s: Series| {
             let ca = s.utf8()?;
             match ca.json_path_match(&pat) {
-                Ok(ca) => Ok(ca.into_series()),
+                Ok(ca) => Ok(Some(ca.into_series())),
                 Err(e) => Err(PolarsError::ComputeError(format!("{:?}", e).into())),
             }
         };
@@ -864,7 +894,7 @@ impl RbExpr {
         self.inner
             .clone()
             .map(
-                |s| Ok(s.duration()?.days().into_series()),
+                |s| Ok(Some(s.duration()?.days().into_series())),
                 GetOutput::from_type(DataType::Int64),
             )
             .into()
@@ -874,7 +904,7 @@ impl RbExpr {
         self.inner
             .clone()
             .map(
-                |s| Ok(s.duration()?.hours().into_series()),
+                |s| Ok(Some(s.duration()?.hours().into_series())),
                 GetOutput::from_type(DataType::Int64),
             )
             .into()
@@ -884,7 +914,7 @@ impl RbExpr {
         self.inner
             .clone()
             .map(
-                |s| Ok(s.duration()?.minutes().into_series()),
+                |s| Ok(Some(s.duration()?.minutes().into_series())),
                 GetOutput::from_type(DataType::Int64),
             )
             .into()
@@ -894,7 +924,7 @@ impl RbExpr {
         self.inner
             .clone()
             .map(
-                |s| Ok(s.duration()?.seconds().into_series()),
+                |s| Ok(Some(s.duration()?.seconds().into_series())),
                 GetOutput::from_type(DataType::Int64),
             )
             .into()
@@ -904,7 +934,7 @@ impl RbExpr {
         self.inner
             .clone()
             .map(
-                |s| Ok(s.duration()?.nanoseconds().into_series()),
+                |s| Ok(Some(s.duration()?.nanoseconds().into_series())),
                 GetOutput::from_type(DataType::Int64),
             )
             .into()
@@ -914,7 +944,7 @@ impl RbExpr {
         self.inner
             .clone()
             .map(
-                |s| Ok(s.duration()?.microseconds().into_series()),
+                |s| Ok(Some(s.duration()?.microseconds().into_series())),
                 GetOutput::from_type(DataType::Int64),
             )
             .into()
@@ -924,7 +954,7 @@ impl RbExpr {
         self.inner
             .clone()
             .map(
-                |s| Ok(s.duration()?.milliseconds().into_series()),
+                |s| Ok(Some(s.duration()?.milliseconds().into_series())),
                 GetOutput::from_type(DataType::Int64),
             )
             .into()
@@ -945,7 +975,7 @@ impl RbExpr {
             .map(
                 |s| {
                     s.timestamp(TimeUnit::Milliseconds)
-                        .map(|ca| (ca / 1000).into_series())
+                        .map(|ca| Some((ca / 1000).into_series()))
                 },
                 GetOutput::from_type(DataType::Int64),
             )
@@ -956,18 +986,19 @@ impl RbExpr {
         self.inner.clone().dt().with_time_unit(tu.0).into()
     }
 
-    pub fn dt_with_time_zone(&self, tz: Option<TimeZone>) -> Self {
-        self.inner.clone().dt().with_time_zone(tz).into()
+    pub fn dt_convert_time_zone(&self, tz: TimeZone) -> Self {
+        self.inner.clone().dt().convert_time_zone(tz).into()
     }
 
     pub fn dt_cast_time_unit(&self, tu: Wrap<TimeUnit>) -> Self {
         self.inner.clone().dt().cast_time_unit(tu.0).into()
     }
 
-    pub fn dt_cast_time_zone(&self, tz: String) -> Self {
-        self.inner.clone().dt().cast_time_zone(tz).into()
+    pub fn dt_replace_time_zone(&self, tz: Option<String>) -> Self {
+        self.inner.clone().dt().replace_time_zone(tz).into()
     }
 
+    #[allow(deprecated)]
     pub fn dt_tz_localize(&self, tz: String) -> Self {
         self.inner.clone().dt().tz_localize(tz).into()
     }
@@ -989,7 +1020,7 @@ impl RbExpr {
     }
 
     pub fn reinterpret(&self, signed: bool) -> Self {
-        let function = move |s: Series| reinterpret(&s, signed);
+        let function = move |s: Series| reinterpret(&s, signed).map(Some);
         let dt = if signed {
             DataType::Int64
         } else {
@@ -1391,7 +1422,7 @@ impl RbExpr {
         self.inner
             .clone()
             .map(
-                |s| Ok(s.to_physical_repr().into_owned()),
+                |s| Ok(Some(s.to_physical_repr().into_owned())),
                 GetOutput::map_dtype(|dt| dt.to_physical()),
             )
             .with_fmt("to_physical")
@@ -1428,44 +1459,67 @@ impl RbExpr {
             .into()
     }
 
-    pub fn ewm_mean(&self, alpha: f64, adjust: bool, min_periods: usize) -> Self {
+    pub fn ewm_mean(
+        &self,
+        alpha: f64,
+        adjust: bool,
+        min_periods: usize,
+        ignore_nulls: bool,
+    ) -> Self {
         let options = EWMOptions {
             alpha,
             adjust,
             bias: false,
             min_periods,
+            ignore_nulls,
         };
         self.inner.clone().ewm_mean(options).into()
     }
 
-    pub fn ewm_std(&self, alpha: f64, adjust: bool, bias: bool, min_periods: usize) -> Self {
+    pub fn ewm_std(
+        &self,
+        alpha: f64,
+        adjust: bool,
+        bias: bool,
+        min_periods: usize,
+        ignore_nulls: bool,
+    ) -> Self {
         let options = EWMOptions {
             alpha,
             adjust,
             bias,
             min_periods,
+            ignore_nulls,
         };
         self.inner.clone().ewm_std(options).into()
     }
 
-    pub fn ewm_var(&self, alpha: f64, adjust: bool, bias: bool, min_periods: usize) -> Self {
+    pub fn ewm_var(
+        &self,
+        alpha: f64,
+        adjust: bool,
+        bias: bool,
+        min_periods: usize,
+        ignore_nulls: bool,
+    ) -> Self {
         let options = EWMOptions {
             alpha,
             adjust,
             bias,
             min_periods,
+            ignore_nulls,
         };
         self.inner.clone().ewm_var(options).into()
     }
 
     pub fn extend_constant(&self, value: Wrap<AnyValue>, n: usize) -> Self {
-        let value = Value::from(value);
+        let value = value.into_value();
         self.inner
             .clone()
             .apply(
                 move |s| {
                     let value = value.try_convert::<Wrap<AnyValue>>().unwrap().0;
-                    s.extend_constant(value, n)
+                    s.extend_constant(value, n).map(Some)
                 },
                 GetOutput::same_type(),
             )

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

@@ -1,14 +1,15 @@
-use crate::{RbExpr, RbPolarsErr, RbResult};
+use crate::{RArray, RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
-    pub fn meta_pop(&self) -> Vec<RbExpr> {
-        self.inner
-            .clone()
-            .meta()
-            .pop()
-            .into_iter()
-            .map(RbExpr::from)
-            .collect()
+    pub fn meta_pop(&self) -> RArray {
+        RArray::from_iter(
+            self.inner
+                .clone()
+                .meta()
+                .pop()
+                .into_iter()
+                .map(RbExpr::from),
+        )
     }
 
     pub fn meta_eq(&self, other: &RbExpr) -> bool {