polars-df 0.6.0 → 0.7.0

data/ext/polars/src/expr/rolling.rs

@@ -0,0 +1,196 @@
+use polars::prelude::*;
+use std::any::Any;
+
+use crate::conversion::Wrap;
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn rolling_sum(
+        &self,
+        window_size: String,
+        weights: Option<Vec<f64>>,
+        min_periods: usize,
+        center: bool,
+        by: Option<String>,
+        closed: Option<Wrap<ClosedWindow>>,
+    ) -> Self {
+        let options = RollingOptions {
+            window_size: Duration::parse(&window_size),
+            weights,
+            min_periods,
+            center,
+            by,
+            closed_window: closed.map(|c| c.0),
+            ..Default::default()
+        };
+        self.inner.clone().rolling_sum(options).into()
+    }
+
+    pub fn rolling_min(
+        &self,
+        window_size: String,
+        weights: Option<Vec<f64>>,
+        min_periods: usize,
+        center: bool,
+        by: Option<String>,
+        closed: Option<Wrap<ClosedWindow>>,
+    ) -> Self {
+        let options = RollingOptions {
+            window_size: Duration::parse(&window_size),
+            weights,
+            min_periods,
+            center,
+            by,
+            closed_window: closed.map(|c| c.0),
+            ..Default::default()
+        };
+        self.inner.clone().rolling_min(options).into()
+    }
+
+    pub fn rolling_max(
+        &self,
+        window_size: String,
+        weights: Option<Vec<f64>>,
+        min_periods: usize,
+        center: bool,
+        by: Option<String>,
+        closed: Option<Wrap<ClosedWindow>>,
+    ) -> Self {
+        let options = RollingOptions {
+            window_size: Duration::parse(&window_size),
+            weights,
+            min_periods,
+            center,
+            by,
+            closed_window: closed.map(|c| c.0),
+            ..Default::default()
+        };
+        self.inner.clone().rolling_max(options).into()
+    }
+
+    pub fn rolling_mean(
+        &self,
+        window_size: String,
+        weights: Option<Vec<f64>>,
+        min_periods: usize,
+        center: bool,
+        by: Option<String>,
+        closed: Option<Wrap<ClosedWindow>>,
+    ) -> Self {
+        let options = RollingOptions {
+            window_size: Duration::parse(&window_size),
+            weights,
+            min_periods,
+            center,
+            by,
+            closed_window: closed.map(|c| c.0),
+            ..Default::default()
+        };
+
+        self.inner.clone().rolling_mean(options).into()
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn rolling_std(
+        &self,
+        window_size: String,
+        weights: Option<Vec<f64>>,
+        min_periods: usize,
+        center: bool,
+        by: Option<String>,
+        closed: Option<Wrap<ClosedWindow>>,
+        ddof: u8,
+    ) -> Self {
+        let options = RollingOptions {
+            window_size: Duration::parse(&window_size),
+            weights,
+            min_periods,
+            center,
+            by,
+            closed_window: closed.map(|c| c.0),
+            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+        };
+
+        self.inner.clone().rolling_std(options).into()
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn rolling_var(
+        &self,
+        window_size: String,
+        weights: Option<Vec<f64>>,
+        min_periods: usize,
+        center: bool,
+        by: Option<String>,
+        closed: Option<Wrap<ClosedWindow>>,
+        ddof: u8,
+    ) -> Self {
+        let options = RollingOptions {
+            window_size: Duration::parse(&window_size),
+            weights,
+            min_periods,
+            center,
+            by,
+            closed_window: closed.map(|c| c.0),
+            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+        };
+
+        self.inner.clone().rolling_var(options).into()
+    }
+
+    pub fn rolling_median(
+        &self,
+        window_size: String,
+        weights: Option<Vec<f64>>,
+        min_periods: usize,
+        center: bool,
+        by: Option<String>,
+        closed: Option<Wrap<ClosedWindow>>,
+    ) -> Self {
+        let options = RollingOptions {
+            window_size: Duration::parse(&window_size),
+            weights,
+            min_periods,
+            center,
+            by,
+            closed_window: closed.map(|c| c.0),
+            fn_params: Some(Arc::new(RollingQuantileParams {
+                prob: 0.5,
+                interpol: QuantileInterpolOptions::Linear,
+            }) as Arc<dyn Any + Send + Sync>),
+        };
+        self.inner.clone().rolling_quantile(options).into()
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn rolling_quantile(
+        &self,
+        quantile: f64,
+        interpolation: Wrap<QuantileInterpolOptions>,
+        window_size: String,
+        weights: Option<Vec<f64>>,
+        min_periods: usize,
+        center: bool,
+        by: Option<String>,
+        closed: Option<Wrap<ClosedWindow>>,
+    ) -> Self {
+        let options = RollingOptions {
+            window_size: Duration::parse(&window_size),
+            weights,
+            min_periods,
+            center,
+            by,
+            closed_window: closed.map(|c| c.0),
+            fn_params: Some(Arc::new(RollingQuantileParams {
+                prob: quantile,
+                interpol: interpolation.0,
+            }) as Arc<dyn Any + Send + Sync>),
+        };
+
+        self.inner.clone().rolling_quantile(options).into()
+    }
+
+    pub fn rolling_skew(&self, window_size: usize, bias: bool) -> Self {
+        self.inner.clone().rolling_skew(window_size, bias).into()
+    }
+}

data/ext/polars/src/expr/string.rs

@@ -4,8 +4,12 @@ use crate::conversion::Wrap;
 use crate::RbExpr;
 
 impl RbExpr {
-    pub fn str_concat(&self, delimiter: String) -> Self {
-        self.inner.clone().str().concat(&delimiter).into()
+    pub fn str_concat(&self, delimiter: String, ignore_nulls: bool) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .concat(&delimiter, ignore_nulls)
+            .into()
     }
 
     pub fn str_to_date(
@@ -24,6 +28,7 @@ impl RbExpr {
         self.inner.clone().str().to_date(options).into()
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn str_to_datetime(
         &self,
         format: Option<String>,
@@ -32,6 +37,7 @@ impl RbExpr {
         strict: bool,
         exact: bool,
         cache: bool,
+        ambiguous: &Self,
     ) -> Self {
         let options = StrptimeOptions {
             format,
@@ -42,7 +48,12 @@ impl RbExpr {
         self.inner
             .clone()
             .str()
-            .to_datetime(time_unit.map(|tu| tu.0), time_zone, options)
+            .to_datetime(
+                time_unit.map(|tu| tu.0),
+                time_zone,
+                options,
+                ambiguous.inner.clone(),
+            )
             .into()
     }
 
@@ -56,30 +67,50 @@ impl RbExpr {
         self.inner.clone().str().to_time(options).into()
     }
 
-    pub fn str_strip(&self, matches: Option<String>) -> Self {
-        self.inner.clone().str().strip(matches).into()
+    pub fn str_strip_chars(&self, matches: &Self) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .strip_chars(matches.inner.clone())
+            .into()
     }
 
-    pub fn str_rstrip(&self, matches: Option<String>) -> Self {
-        self.inner.clone().str().rstrip(matches).into()
+    pub fn str_strip_chars_start(&self, matches: &Self) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .strip_chars_start(matches.inner.clone())
+            .into()
     }
 
-    pub fn str_lstrip(&self, matches: Option<String>) -> Self {
-        self.inner.clone().str().lstrip(matches).into()
+    pub fn str_strip_chars_end(&self, matches: &Self) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .strip_chars_end(matches.inner.clone())
+            .into()
     }
 
-    pub fn str_slice(&self, start: i64, length: Option<u64>) -> Self {
-        let function = move |s: Series| {
-            let ca = s.utf8()?;
-            Ok(Some(ca.str_slice(start, length)?.into_series()))
-        };
-        self.clone()
-            .inner
-            .map(function, GetOutput::from_type(DataType::Utf8))
-            .with_fmt("str.slice")
+    pub fn str_strip_prefix(&self, prefix: &Self) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .strip_prefix(prefix.inner.clone())
+            .into()
+    }
+
+    pub fn str_strip_suffix(&self, suffix: &Self) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .strip_suffix(suffix.inner.clone())
             .into()
     }
 
+    pub fn str_slice(&self, start: i64, length: Option<u64>) -> Self {
+        self.inner.clone().str().slice(start, length).into()
+    }
+
     pub fn str_explode(&self) -> Self {
         self.inner.clone().str().explode().into()
     }
@@ -92,28 +123,12 @@ impl RbExpr {
         self.inner.clone().str().to_lowercase().into()
     }
 
-    pub fn str_lengths(&self) -> Self {
-        let function = |s: Series| {
-            let ca = s.utf8()?;
-            Ok(Some(ca.str_lengths().into_series()))
-        };
-        self.clone()
-            .inner
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("str.lengths")
-            .into()
+    pub fn str_len_bytes(&self) -> Self {
+        self.inner.clone().str().len_bytes().into()
     }
 
-    pub fn str_n_chars(&self) -> Self {
-        let function = |s: Series| {
-            let ca = s.utf8()?;
-            Ok(Some(ca.str_n_chars().into_series()))
-        };
-        self.clone()
-            .inner
-            .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("str.n_chars")
-            .into()
+    pub fn str_len_chars(&self) -> Self {
+        self.inner.clone().str().len_chars().into()
     }
 
     pub fn str_replace_n(&self, pat: &RbExpr, val: &RbExpr, literal: bool, n: i64) -> Self {
@@ -132,16 +147,16 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_zfill(&self, alignment: usize) -> Self {
-        self.clone().inner.str().zfill(alignment).into()
+    pub fn str_pad_start(&self, length: usize, fillchar: char) -> Self {
+        self.clone().inner.str().pad_start(length, fillchar).into()
     }
 
-    pub fn str_ljust(&self, width: usize, fillchar: char) -> Self {
-        self.clone().inner.str().ljust(width, fillchar).into()
+    pub fn str_pad_end(&self, length: usize, fillchar: char) -> Self {
+        self.clone().inner.str().pad_end(length, fillchar).into()
     }
 
-    pub fn str_rjust(&self, width: usize, fillchar: char) -> Self {
-        self.clone().inner.str().rjust(width, fillchar).into()
+    pub fn str_zfill(&self, alignment: usize) -> Self {
+        self.clone().inner.str().zfill(alignment).into()
     }
 
     pub fn str_contains(&self, pat: &RbExpr, literal: Option<bool>, strict: bool) -> Self {
@@ -221,11 +236,11 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_parse_int(&self, radix: u32, strict: bool) -> Self {
+    pub fn str_to_integer(&self, base: u32, strict: bool) -> Self {
         self.inner
             .clone()
             .str()
-            .from_radix(radix, strict)
+            .to_integer(base, strict)
             .with_fmt("str.parse_int")
             .into()
     }
@@ -284,31 +299,43 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_count_match(&self, pat: String) -> Self {
-        self.inner.clone().str().count_match(&pat).into()
+    pub fn str_count_matches(&self, pat: &Self, literal: bool) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .count_matches(pat.inner.clone(), literal)
+            .into()
     }
 
-    pub fn str_split(&self, by: String) -> Self {
-        self.inner.clone().str().split(&by).into()
+    pub fn str_split(&self, by: &Self) -> Self {
+        self.inner.clone().str().split(by.inner.clone()).into()
     }
 
-    pub fn str_split_inclusive(&self, by: String) -> Self {
-        self.inner.clone().str().split_inclusive(&by).into()
+    pub fn str_split_inclusive(&self, by: &Self) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .split_inclusive(by.inner.clone())
+            .into()
     }
 
-    pub fn str_split_exact(&self, by: String, n: usize) -> Self {
-        self.inner.clone().str().split_exact(&by, n).into()
+    pub fn str_split_exact(&self, by: &Self, n: usize) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .split_exact(by.inner.clone(), n)
+            .into()
     }
 
-    pub fn str_split_exact_inclusive(&self, by: String, n: usize) -> Self {
+    pub fn str_split_exact_inclusive(&self, by: &Self, n: usize) -> Self {
         self.inner
             .clone()
             .str()
-            .split_exact_inclusive(&by, n)
+            .split_exact_inclusive(by.inner.clone(), n)
             .into()
     }
 
-    pub fn str_splitn(&self, by: String, n: usize) -> Self {
-        self.inner.clone().str().splitn(&by, n).into()
+    pub fn str_splitn(&self, by: &Self, n: usize) -> Self {
+        self.inner.clone().str().splitn(by.inner.clone(), n).into()
     }
 }

data/ext/polars/src/file.rs

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

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

@@ -0,0 +1,35 @@
+use magnus::RArray;
+use polars::lazy::dsl;
+
+use crate::rb_exprs_to_exprs;
+use crate::{RbExpr, RbPolarsErr, RbResult};
+
+pub fn all_horizontal(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let e = dsl::all_horizontal(exprs).map_err(RbPolarsErr::from)?;
+    Ok(e.into())
+}
+
+pub fn any_horizontal(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let e = dsl::any_horizontal(exprs).map_err(RbPolarsErr::from)?;
+    Ok(e.into())
+}
+
+pub fn max_horizontal(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let e = dsl::max_horizontal(exprs).map_err(RbPolarsErr::from)?;
+    Ok(e.into())
+}
+
+pub fn min_horizontal(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let e = dsl::min_horizontal(exprs).map_err(RbPolarsErr::from)?;
+    Ok(e.into())
+}
+
+pub fn sum_horizontal(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let e = dsl::sum_horizontal(exprs).map_err(RbPolarsErr::from)?;
+    Ok(e.into())
+}

data/ext/polars/src/functions/eager.rs

@@ -1,11 +1,9 @@
 use magnus::RArray;
-use polars::{functions, time};
-use polars_core::datatypes::{TimeUnit, TimeZone};
-use polars_core::prelude::{DataFrame, IntoSeries};
+use polars::functions;
+use polars_core::prelude::DataFrame;
 
-use crate::conversion::{get_df, get_series, Wrap};
+use crate::conversion::{get_df, get_series};
 use crate::error::RbPolarsErr;
-use crate::prelude::{ClosedWindow, Duration};
 use crate::{RbDataFrame, RbResult, RbSeries};
 
 pub fn concat_df(seq: RArray) -> RbResult<RbDataFrame> {
@@ -52,42 +50,20 @@ pub fn concat_series(seq: RArray) -> RbResult<RbSeries> {
     Ok(s.into())
 }
 
-pub fn date_range(
-    start: i64,
-    stop: i64,
-    every: String,
-    closed: Wrap<ClosedWindow>,
-    name: String,
-    tu: Wrap<TimeUnit>,
-    tz: Option<TimeZone>,
-) -> RbResult<RbSeries> {
-    let date_range = time::date_range_impl(
-        &name,
-        start,
-        stop,
-        Duration::parse(&every),
-        closed.0,
-        tu.0,
-        tz.as_ref(),
-    )
-    .map_err(RbPolarsErr::from)?;
-    Ok(date_range.into_series().into())
-}
-
-pub fn diag_concat_df(seq: RArray) -> RbResult<RbDataFrame> {
+pub fn concat_df_diagonal(seq: RArray) -> RbResult<RbDataFrame> {
     let mut dfs = Vec::new();
     for item in seq.each() {
         dfs.push(get_df(item?)?);
     }
-    let df = functions::diag_concat_df(&dfs).map_err(RbPolarsErr::from)?;
+    let df = functions::concat_df_diagonal(&dfs).map_err(RbPolarsErr::from)?;
     Ok(df.into())
 }
 
-pub fn hor_concat_df(seq: RArray) -> RbResult<RbDataFrame> {
+pub fn concat_df_horizontal(seq: RArray) -> RbResult<RbDataFrame> {
     let mut dfs = Vec::new();
     for item in seq.each() {
         dfs.push(get_df(item?)?);
     }
-    let df = functions::hor_concat_df(&dfs).map_err(RbPolarsErr::from)?;
+    let df = functions::concat_df_horizontal(&dfs).map_err(RbPolarsErr::from)?;
     Ok(df.into())
 }

data/ext/polars/src/functions/io.rs

@@ -5,30 +5,30 @@ use crate::file::get_file_like;
 use crate::prelude::DataType;
 use crate::{RbPolarsErr, RbResult};
 
-pub fn read_ipc_schema(rb_f: Value) -> RbResult<Value> {
+pub fn read_ipc_schema(rb_f: Value) -> RbResult<RHash> {
     use polars_core::export::arrow::io::ipc::read::read_file_metadata;
     let mut r = get_file_like(rb_f, false)?;
-    let metadata = read_file_metadata(&mut r).map_err(RbPolarsErr::arrow)?;
+    let metadata = read_file_metadata(&mut r).map_err(RbPolarsErr::from)?;
 
     let dict = RHash::new();
-    for field in metadata.schema.fields {
+    for field in &metadata.schema.fields {
         let dt: Wrap<DataType> = Wrap((&field.data_type).into());
-        dict.aset(field.name, dt)?;
+        dict.aset(field.name.clone(), dt)?;
     }
-    Ok(dict.into())
+    Ok(dict)
 }
 
-pub fn read_parquet_schema(rb_f: Value) -> RbResult<Value> {
-    use polars_core::export::arrow::io::parquet::read::{infer_schema, read_metadata};
+pub fn read_parquet_schema(rb_f: Value) -> RbResult<RHash> {
+    use polars_parquet::read::{infer_schema, read_metadata};
 
     let mut r = get_file_like(rb_f, false)?;
-    let metadata = read_metadata(&mut r).map_err(RbPolarsErr::arrow)?;
-    let arrow_schema = infer_schema(&metadata).map_err(RbPolarsErr::arrow)?;
+    let metadata = read_metadata(&mut r).map_err(RbPolarsErr::from)?;
+    let arrow_schema = infer_schema(&metadata).map_err(RbPolarsErr::from)?;
 
     let dict = RHash::new();
     for field in arrow_schema.fields {
         let dt: Wrap<DataType> = Wrap((&field.data_type).into());
         dict.aset(field.name, dt)?;
     }
-    Ok(dict.into())
+    Ok(dict)
 }