polars-df 0.10.0 → 0.11.0

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

@@ -7,187 +7,281 @@ use crate::RbExpr;
 impl RbExpr {
     pub fn rolling_sum(
         &self,
-        window_size: String,
+        window_size: usize,
         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),
+        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,
+        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_sum_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_min(
+        &self,
+        window_size: usize,
         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),
+        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,
+        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_min_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_max(
+        &self,
+        window_size: usize,
         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),
+        let options = RollingOptionsFixedWindow {
+            window_size: 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,
+        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_max_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_mean(
+        &self,
+        window_size: usize,
         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),
+        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,
+        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: 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 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,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+        ddof: u8,
+    ) -> Self {
+        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: 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 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,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+        ddof: u8,
+    ) -> Self {
+        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_var_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_median(
+        &self,
+        window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: usize,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-        warn_if_unsorted: bool,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
+        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,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-        warn_if_unsorted: bool,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&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 +290,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

@@ -259,39 +259,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/lazy.rs

@@ -130,6 +130,7 @@ pub fn concat_lf(
             rechunk,
             parallel,
             to_supertypes,
+            ..Default::default()
         },
     )
     .map_err(RbPolarsErr::from)?;
@@ -198,6 +199,7 @@ pub fn concat_lf_diagonal(
             rechunk,
             parallel,
             to_supertypes,
+            ..Default::default()
         },
     )
     .map_err(RbPolarsErr::from)?;

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

@@ -39,3 +39,77 @@ pub fn date_range(
     let time_unit = time_unit.map(|x| x.0);
     dsl::date_range(start, end, every, closed, time_unit, time_zone).into()
 }
+
+pub fn date_ranges(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    time_unit: Option<Wrap<TimeUnit>>,
+    time_zone: Option<TimeZone>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    let time_unit = time_unit.map(|x| x.0);
+    dsl::date_ranges(start, end, every, closed, time_unit, time_zone).into()
+}
+
+pub fn datetime_range(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    time_unit: Option<Wrap<TimeUnit>>,
+    time_zone: Option<TimeZone>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    let time_unit = time_unit.map(|x| x.0);
+    dsl::datetime_range(start, end, every, closed, time_unit, time_zone).into()
+}
+
+pub fn datetime_ranges(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    time_unit: Option<Wrap<TimeUnit>>,
+    time_zone: Option<TimeZone>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    let time_unit = time_unit.map(|x| x.0);
+    dsl::datetime_ranges(start, end, every, closed, time_unit, time_zone).into()
+}
+
+pub fn time_range(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    dsl::time_range(start, end, every, closed).into()
+}
+
+pub fn time_ranges(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    dsl::time_ranges(start, end, every, closed).into()
+}

data/ext/polars/src/interop/mod.rs

@@ -0,0 +1 @@
+pub mod numo;

data/ext/polars/src/interop/numo/mod.rs

@@ -0,0 +1,2 @@
+pub mod to_numo_df;
+pub mod to_numo_series;

data/ext/polars/src/interop/numo/to_numo_df.rs

@@ -0,0 +1,23 @@
+use magnus::Value;
+use polars_core::utils::try_get_supertype;
+
+use crate::dataframe::RbDataFrame;
+
+impl RbDataFrame {
+    pub fn to_numo(&self) -> Option<Value> {
+        let mut st = None;
+        for s in self.df.borrow().iter() {
+            let dt_i = s.dtype();
+            match st {
+                None => st = Some(dt_i.clone()),
+                Some(ref mut st) => {
+                    *st = try_get_supertype(st, dt_i).ok()?;
+                }
+            }
+        }
+        let _st = st?;
+
+        // TODO
+        None
+    }
+}

data/ext/polars/src/interop/numo/to_numo_series.rs

@@ -0,0 +1,60 @@
+use magnus::{class, prelude::*, Module, RArray, RClass, RModule, Value};
+use polars_core::prelude::*;
+
+use crate::error::RbPolarsErr;
+use crate::raise_err;
+use crate::series::RbSeries;
+use crate::RbResult;
+
+impl RbSeries {
+    /// For numeric types, this should only be called for Series with null types.
+    /// This will cast to floats so that `nil = NAN`
+    pub fn to_numo(&self) -> RbResult<Value> {
+        let s = &self.series.borrow();
+        match s.dtype() {
+            DataType::String => {
+                let ca = s.str().unwrap();
+
+                // TODO make more efficient
+                let np_arr = RArray::from_iter(ca);
+                class::object()
+                    .const_get::<_, RModule>("Numo")?
+                    .const_get::<_, RClass>("RObject")?
+                    .funcall("cast", (np_arr,))
+            }
+            dt if dt.is_numeric() => {
+                if s.bit_repr_is_large() {
+                    let s = s.cast(&DataType::Float64).unwrap();
+                    let ca = s.f64().unwrap();
+                    // TODO make more efficient
+                    let np_arr = RArray::from_iter(ca.into_iter().map(|opt_v| match opt_v {
+                        Some(v) => v,
+                        None => f64::NAN,
+                    }));
+                    class::object()
+                        .const_get::<_, RModule>("Numo")?
+                        .const_get::<_, RClass>("DFloat")?
+                        .funcall("cast", (np_arr,))
+                } else {
+                    let s = s.cast(&DataType::Float32).unwrap();
+                    let ca = s.f32().unwrap();
+                    // TODO make more efficient
+                    let np_arr = RArray::from_iter(ca.into_iter().map(|opt_v| match opt_v {
+                        Some(v) => v,
+                        None => f32::NAN,
+                    }));
+                    class::object()
+                        .const_get::<_, RModule>("Numo")?
+                        .const_get::<_, RClass>("SFloat")?
+                        .funcall("cast", (np_arr,))
+                }
+            }
+            dt => {
+                raise_err!(
+                    format!("'to_numo' not supported for dtype: {dt:?}"),
+                    ComputeError
+                );
+            }
+        }
+    }
+}