polars-df 0.10.0 → 0.12.0

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

@@ -7,187 +7,289 @@ use crate::RbExpr;
 impl RbExpr {
     pub fn rolling_sum(
         &self,
-        window_size: String,
+        window_size: usize,
         weights: Option<Vec<f64>>,
-        min_periods: usize,
+        min_periods: Option<usize>,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            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(
+    pub fn rolling_sum_by(
         &self,
+        by: &RbExpr,
         window_size: String,
-        weights: Option<Vec<f64>>,
         min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
+        closed: Wrap<ClosedWindow>,
     ) -> Self {
-        let options = RollingOptions {
+        let options = RollingOptionsDynamicWindow {
             window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+        self.inner
+            .clone()
+            .rolling_sum_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_min(
+        &self,
+        window_size: usize,
+        weights: Option<Vec<f64>>,
+        min_periods: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            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(
+    pub fn rolling_min_by(
         &self,
+        by: &RbExpr,
         window_size: String,
-        weights: Option<Vec<f64>>,
         min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
+        closed: Wrap<ClosedWindow>,
     ) -> Self {
-        let options = RollingOptions {
+        let options = RollingOptionsDynamicWindow {
             window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+        self.inner
+            .clone()
+            .rolling_min_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_max(
+        &self,
+        window_size: usize,
+        weights: Option<Vec<f64>>,
+        min_periods: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            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(
+    pub fn rolling_max_by(
         &self,
+        by: &RbExpr,
         window_size: String,
-        weights: Option<Vec<f64>>,
         min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
+        closed: Wrap<ClosedWindow>,
     ) -> Self {
-        let options = RollingOptions {
+        let options = RollingOptionsDynamicWindow {
             window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+        self.inner
+            .clone()
+            .rolling_max_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_mean(
+        &self,
+        window_size: usize,
+        weights: Option<Vec<f64>>,
+        min_periods: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            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(
+    pub fn rolling_mean_by(
         &self,
+        by: &RbExpr,
         window_size: String,
-        weights: Option<Vec<f64>>,
         min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+
+        self.inner
+            .clone()
+            .rolling_mean_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_std(
+        &self,
+        window_size: usize,
+        weights: Option<Vec<f64>>,
+        min_periods: Option<usize>,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
         ddof: u8,
-        warn_if_unsorted: bool,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            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>),
-            warn_if_unsorted,
         };
 
         self.inner.clone().rolling_std(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub fn rolling_var(
+    pub fn rolling_std_by(
         &self,
+        by: &RbExpr,
         window_size: String,
-        weights: Option<Vec<f64>>,
         min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
+        closed: Wrap<ClosedWindow>,
         ddof: u8,
-        warn_if_unsorted: bool,
     ) -> Self {
-        let options = RollingOptions {
+        let options = RollingOptionsDynamicWindow {
             window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+        };
+
+        self.inner
+            .clone()
+            .rolling_std_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_var(
+        &self,
+        window_size: usize,
+        weights: Option<Vec<f64>>,
+        min_periods: Option<usize>,
+        center: bool,
+        ddof: u8,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            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>),
-            warn_if_unsorted,
         };
 
         self.inner.clone().rolling_var(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub fn rolling_median(
+    pub fn rolling_var_by(
         &self,
+        by: &RbExpr,
         window_size: String,
-        weights: Option<Vec<f64>>,
         min_periods: usize,
-        center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-        warn_if_unsorted: bool,
+        closed: Wrap<ClosedWindow>,
+        ddof: u8,
     ) -> Self {
-        let options = RollingOptions {
+        let options = RollingOptionsDynamicWindow {
             window_size: Duration::parse(&window_size),
-            weights,
             min_periods,
+            closed_window: closed.0,
+            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+        };
+
+        self.inner
+            .clone()
+            .rolling_var_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_median(
+        &self,
+        window_size: usize,
+        weights: Option<Vec<f64>>,
+        min_periods: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            window_size,
+            min_periods,
+            weights,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             fn_params: None,
-            warn_if_unsorted,
         };
         self.inner.clone().rolling_median(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
+    pub fn rolling_median_by(
+        &self,
+        by: &RbExpr,
+        window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+        self.inner
+            .clone()
+            .rolling_median_by(by.inner.clone(), options)
+            .into()
+    }
+
     pub fn rolling_quantile(
         &self,
         quantile: f64,
         interpolation: Wrap<QuantileInterpolOptions>,
-        window_size: String,
+        window_size: usize,
         weights: Option<Vec<f64>>,
-        min_periods: usize,
+        min_periods: Option<usize>,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-        warn_if_unsorted: bool,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            window_size,
             weights,
             min_periods,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             fn_params: None,
-            warn_if_unsorted,
         };
 
         self.inner
@@ -196,6 +298,28 @@ impl RbExpr {
             .into()
     }
 
+    pub fn rolling_quantile_by(
+        &self,
+        by: &RbExpr,
+        quantile: f64,
+        interpolation: Wrap<QuantileInterpolOptions>,
+        window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+
+        self.inner
+            .clone()
+            .rolling_quantile_by(by.inner.clone(), interpolation.0, 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,11 +4,11 @@ use crate::conversion::Wrap;
 use crate::{RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
-    pub fn str_concat(&self, delimiter: String, ignore_nulls: bool) -> Self {
+    pub fn str_join(&self, delimiter: String, ignore_nulls: bool) -> Self {
         self.inner
             .clone()
             .str()
-            .concat(&delimiter, ignore_nulls)
+            .join(&delimiter, ignore_nulls)
             .into()
     }
 
@@ -115,10 +115,6 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_explode(&self) -> Self {
-        self.inner.clone().str().explode().into()
-    }
-
     pub fn str_to_uppercase(&self) -> Self {
         self.inner.clone().str().to_uppercase().into()
     }
@@ -259,39 +255,18 @@ impl RbExpr {
         infer_schema_len: Option<usize>,
     ) -> Self {
         let dtype = dtype.map(|wrap| wrap.0);
-
-        let output_type = match dtype.clone() {
-            Some(dtype) => GetOutput::from_type(dtype),
-            None => GetOutput::from_type(DataType::Unknown),
-        };
-
-        let function = move |s: Series| {
-            let ca = s.str()?;
-            match ca.json_decode(dtype.clone(), infer_schema_len) {
-                Ok(ca) => Ok(Some(ca.into_series())),
-                Err(e) => Err(PolarsError::ComputeError(format!("{e:?}").into())),
-            }
-        };
-
-        self.clone()
-            .inner
-            .map(function, output_type)
-            .with_fmt("str.json_decode")
+        self.inner
+            .clone()
+            .str()
+            .json_decode(dtype, infer_schema_len)
             .into()
     }
 
-    pub fn str_json_path_match(&self, pat: String) -> Self {
-        let function = move |s: Series| {
-            let ca = s.str()?;
-            match ca.json_path_match(&pat) {
-                Ok(ca) => Ok(Some(ca.into_series())),
-                Err(e) => Err(PolarsError::ComputeError(format!("{:?}", e).into())),
-            }
-        };
-        self.clone()
-            .inner
-            .map(function, GetOutput::from_type(DataType::String))
-            .with_fmt("str.json_path_match")
+    pub fn str_json_path_match(&self, pat: &Self) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .json_path_match(pat.inner.clone())
             .into()
     }
 

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

@@ -9,8 +9,8 @@ use crate::{RbDataFrame, RbResult, RbSeries};
 pub fn concat_df(seq: RArray) -> RbResult<RbDataFrame> {
     use polars_core::error::PolarsResult;
 
-    let mut iter = seq.each();
-    let first = iter.next().unwrap()?;
+    let mut iter = seq.into_iter();
+    let first = iter.next().unwrap();
 
     let first_rdf = get_df(first)?;
     let identity_df = first_rdf.slice(0, 0);
@@ -18,7 +18,7 @@ pub fn concat_df(seq: RArray) -> RbResult<RbDataFrame> {
     let mut rdfs: Vec<PolarsResult<DataFrame>> = vec![Ok(first_rdf)];
 
     for item in iter {
-        let rdf = get_df(item?)?;
+        let rdf = get_df(item)?;
         rdfs.push(Ok(rdf));
     }
 
@@ -37,13 +37,13 @@ pub fn concat_df(seq: RArray) -> RbResult<RbDataFrame> {
 }
 
 pub fn concat_series(seq: RArray) -> RbResult<RbSeries> {
-    let mut iter = seq.each();
-    let first = iter.next().unwrap()?;
+    let mut iter = seq.into_iter();
+    let first = iter.next().unwrap();
 
     let mut s = get_series(first)?;
 
     for res in iter {
-        let item = res?;
+        let item = res;
         let item = get_series(item)?;
         s.append(&item).map_err(RbPolarsErr::from)?;
     }
@@ -52,8 +52,8 @@ pub fn concat_series(seq: RArray) -> RbResult<RbSeries> {
 
 pub fn concat_df_diagonal(seq: RArray) -> RbResult<RbDataFrame> {
     let mut dfs = Vec::new();
-    for item in seq.each() {
-        dfs.push(get_df(item?)?);
+    for item in seq.into_iter() {
+        dfs.push(get_df(item)?);
     }
     let df = functions::concat_df_diagonal(&dfs).map_err(RbPolarsErr::from)?;
     Ok(df.into())
@@ -61,8 +61,8 @@ pub fn concat_df_diagonal(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?)?);
+    for item in seq.into_iter() {
+        dfs.push(get_df(item)?);
     }
     let df = functions::concat_df_horizontal(&dfs).map_err(RbPolarsErr::from)?;
     Ok(df.into())

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

@@ -58,7 +58,7 @@ pub fn rolling_cov(
 pub fn arg_sort_by(
     by: RArray,
     descending: Vec<bool>,
-    nulls_last: bool,
+    nulls_last: Vec<bool>,
     multithreaded: bool,
     maintain_order: bool,
 ) -> RbResult<RbExpr> {
@@ -95,12 +95,12 @@ pub fn col(name: String) -> RbExpr {
 
 pub fn collect_all(lfs: RArray) -> RbResult<RArray> {
     let lfs = lfs
-        .each()
-        .map(|v| <&RbLazyFrame>::try_convert(v?))
+        .into_iter()
+        .map(<&RbLazyFrame>::try_convert)
         .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
 
     Ok(RArray::from_iter(lfs.iter().map(|lf| {
-        let df = lf.ldf.clone().collect().unwrap();
+        let df = lf.ldf.borrow().clone().collect().unwrap();
         RbDataFrame::new(df)
     })))
 }
@@ -118,8 +118,8 @@ pub fn concat_lf(
     let (seq, len) = get_rbseq(lfs)?;
     let mut lfs = Vec::with_capacity(len);
 
-    for res in seq.each() {
-        let item = res?;
+    for res in seq.into_iter() {
+        let item = res;
         let lf = get_lf(item)?;
         lfs.push(lf);
     }
@@ -130,6 +130,7 @@ pub fn concat_lf(
             rechunk,
             parallel,
             to_supertypes,
+            ..Default::default()
         },
     )
     .map_err(RbPolarsErr::from)?;
@@ -183,14 +184,9 @@ pub fn concat_lf_diagonal(
     parallel: bool,
     to_supertypes: bool,
 ) -> RbResult<RbLazyFrame> {
-    let iter = lfs.each();
+    let iter = lfs.into_iter();
 
-    let lfs = iter
-        .map(|item| {
-            let item = item?;
-            get_lf(item)
-        })
-        .collect::<RbResult<Vec<_>>>()?;
+    let lfs = iter.map(get_lf).collect::<RbResult<Vec<_>>>()?;
 
     let lf = dsl::functions::concat_lf_diagonal(
         lfs,
@@ -198,23 +194,29 @@ pub fn concat_lf_diagonal(
             rechunk,
             parallel,
             to_supertypes,
+            ..Default::default()
         },
     )
     .map_err(RbPolarsErr::from)?;
     Ok(lf.into())
 }
 
-pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
-    dsl::dtype_cols(dtypes).into()
-}
-
-pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
+pub fn dtype_cols(dtypes: RArray) -> RbResult<RbExpr> {
     let dtypes = dtypes
-        .each()
-        .map(|v| Wrap::<DataType>::try_convert(v?))
+        .into_iter()
+        .map(Wrap::<DataType>::try_convert)
         .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
     let dtypes = vec_extract_wrapped(dtypes);
-    Ok(crate::functions::lazy::dtype_cols(dtypes))
+    Ok(dsl::dtype_cols(dtypes).into())
+}
+
+pub fn index_cols(indices: Vec<i64>) -> RbExpr {
+    if indices.len() == 1 {
+        dsl::nth(indices[0])
+    } else {
+        dsl::index_cols(indices)
+    }
+    .into()
 }
 
 #[allow(clippy::too_many_arguments)]