polars-df 0.5.0 → 0.6.0

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

@@ -2,6 +2,8 @@ use magnus::{block::Proc, IntoValue, RArray, Value};
 use polars::lazy::dsl;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
+use polars_core::series::IsSorted;
+use std::any::Any;
 
 use crate::apply::lazy::map_single;
 use crate::conversion::{parse_fill_null_strategy, Wrap};
@@ -201,6 +203,7 @@ impl RbExpr {
                 descending,
                 nulls_last,
                 multithreaded: true,
+                maintain_order: false,
             })
             .into()
     }
@@ -212,6 +215,7 @@ impl RbExpr {
                 descending: reverse,
                 nulls_last,
                 multithreaded: true,
+                maintain_order: false,
             })
             .into()
     }
@@ -588,6 +592,7 @@ impl RbExpr {
             center,
             by,
             closed_window: closed.map(|c| c.0),
+            ..Default::default()
         };
         self.inner.clone().rolling_sum(options).into()
     }
@@ -608,6 +613,7 @@ impl RbExpr {
             center,
             by,
             closed_window: closed.map(|c| c.0),
+            ..Default::default()
         };
         self.inner.clone().rolling_min(options).into()
     }
@@ -628,6 +634,7 @@ impl RbExpr {
             center,
             by,
             closed_window: closed.map(|c| c.0),
+            ..Default::default()
         };
         self.inner.clone().rolling_max(options).into()
     }
@@ -648,11 +655,13 @@ impl RbExpr {
             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,
@@ -661,6 +670,7 @@ impl RbExpr {
         center: bool,
         by: Option<String>,
         closed: Option<Wrap<ClosedWindow>>,
+        ddof: u8,
     ) -> Self {
         let options = RollingOptions {
             window_size: Duration::parse(&window_size),
@@ -669,11 +679,13 @@ impl RbExpr {
             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,
@@ -682,6 +694,7 @@ impl RbExpr {
         center: bool,
         by: Option<String>,
         closed: Option<Wrap<ClosedWindow>>,
+        ddof: u8,
     ) -> Self {
         let options = RollingOptions {
             window_size: Duration::parse(&window_size),
@@ -690,6 +703,7 @@ impl RbExpr {
             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()
@@ -711,8 +725,12 @@ impl RbExpr {
             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_median(options).into()
+        self.inner.clone().rolling_quantile(options).into()
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -734,12 +752,13 @@ impl RbExpr {
             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(quantile, interpolation.0, options)
-            .into()
+        self.inner.clone().rolling_quantile(options).into()
     }
 
     pub fn rolling_skew(&self, window_size: usize, bias: bool) -> Self {
@@ -809,8 +828,8 @@ impl RbExpr {
             .into()
     }
 
-    pub fn shuffle(&self, seed: Option<u64>) -> Self {
-        self.inner.clone().shuffle(seed).into()
+    pub fn shuffle(&self, seed: Option<u64>, fixed_seed: bool) -> Self {
+        self.inner.clone().shuffle(seed, fixed_seed).into()
     }
 
     pub fn sample_n(
@@ -819,10 +838,11 @@ impl RbExpr {
         with_replacement: bool,
         shuffle: bool,
         seed: Option<u64>,
+        fixed_seed: bool,
     ) -> Self {
         self.inner
             .clone()
-            .sample_n(n, with_replacement, shuffle, seed)
+            .sample_n(n, with_replacement, shuffle, seed, fixed_seed)
             .into()
     }
 
@@ -832,10 +852,11 @@ impl RbExpr {
         with_replacement: bool,
         shuffle: bool,
         seed: Option<u64>,
+        fixed_seed: bool,
     ) -> Self {
         self.inner
             .clone()
-            .sample_frac(frac, with_replacement, shuffle, seed)
+            .sample_frac(frac, with_replacement, shuffle, seed, fixed_seed)
             .into()
     }
 
@@ -930,4 +951,13 @@ impl RbExpr {
     pub fn hash(&self, seed: u64, seed_1: u64, seed_2: u64, seed_3: u64) -> Self {
         self.inner.clone().hash(seed, seed_1, seed_2, seed_3).into()
     }
+
+    pub fn set_sorted_flag(&self, descending: bool) -> Self {
+        let is_sorted = if descending {
+            IsSorted::Descending
+        } else {
+            IsSorted::Ascending
+        };
+        self.inner.clone().set_sorted_flag(is_sorted).into()
+    }
 }

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

@@ -8,17 +8,17 @@ use crate::{RbExpr, RbResult};
 
 impl RbExpr {
     pub fn list_arg_max(&self) -> Self {
-        self.inner.clone().arr().arg_max().into()
+        self.inner.clone().list().arg_max().into()
     }
 
     pub fn list_arg_min(&self) -> Self {
-        self.inner.clone().arr().arg_min().into()
+        self.inner.clone().list().arg_min().into()
     }
 
     pub fn list_contains(&self, other: &RbExpr) -> Self {
         self.inner
             .clone()
-            .arr()
+            .list()
             .contains(other.inner.clone())
             .into()
     }
@@ -26,53 +26,58 @@ impl RbExpr {
     pub fn list_count_match(&self, expr: &RbExpr) -> Self {
         self.inner
             .clone()
-            .arr()
+            .list()
             .count_match(expr.inner.clone())
             .into()
     }
 
     pub fn list_diff(&self, n: i64, null_behavior: Wrap<NullBehavior>) -> RbResult<Self> {
-        Ok(self.inner.clone().arr().diff(n, null_behavior.0).into())
+        Ok(self.inner.clone().list().diff(n, null_behavior.0).into())
     }
 
     pub fn list_eval(&self, expr: &RbExpr, parallel: bool) -> Self {
         self.inner
             .clone()
-            .arr()
+            .list()
             .eval(expr.inner.clone(), parallel)
             .into()
     }
 
     pub fn list_get(&self, index: &RbExpr) -> Self {
-        self.inner.clone().arr().get(index.inner.clone()).into()
+        self.inner.clone().list().get(index.inner.clone()).into()
     }
 
     pub fn list_join(&self, separator: String) -> Self {
-        self.inner.clone().arr().join(&separator).into()
+        self.inner.clone().list().join(&separator).into()
     }
 
     pub fn list_lengths(&self) -> Self {
-        self.inner.clone().arr().lengths().into()
+        self.inner.clone().list().lengths().into()
     }
 
     pub fn list_max(&self) -> Self {
-        self.inner.clone().arr().max().into()
+        self.inner.clone().list().max().into()
     }
 
     pub fn list_mean(&self) -> Self {
-        self.inner.clone().arr().mean().with_fmt("arr.mean").into()
+        self.inner
+            .clone()
+            .list()
+            .mean()
+            .with_fmt("list.mean")
+            .into()
     }
 
     pub fn list_min(&self) -> Self {
-        self.inner.clone().arr().min().into()
+        self.inner.clone().list().min().into()
     }
 
     pub fn list_reverse(&self) -> Self {
-        self.inner.clone().arr().reverse().into()
+        self.inner.clone().list().reverse().into()
     }
 
     pub fn list_shift(&self, periods: i64) -> Self {
-        self.inner.clone().arr().shift(periods).into()
+        self.inner.clone().list().shift(periods).into()
     }
 
     pub fn list_slice(&self, offset: &RbExpr, length: Option<&RbExpr>) -> Self {
@@ -82,7 +87,7 @@ impl RbExpr {
         };
         self.inner
             .clone()
-            .arr()
+            .list()
             .slice(offset.inner.clone(), length)
             .into()
     }
@@ -90,23 +95,23 @@ impl RbExpr {
     pub fn list_sort(&self, reverse: bool) -> Self {
         self.inner
             .clone()
-            .arr()
+            .list()
             .sort(SortOptions {
                 descending: reverse,
                 ..Default::default()
             })
-            .with_fmt("arr.sort")
+            .with_fmt("list.sort")
             .into()
     }
 
     pub fn list_sum(&self) -> Self {
-        self.inner.clone().arr().sum().with_fmt("arr.sum").into()
+        self.inner.clone().list().sum().with_fmt("list.sum").into()
     }
 
     pub fn list_take(&self, index: &RbExpr, null_on_oob: bool) -> Self {
         self.inner
             .clone()
-            .arr()
+            .list()
             .take(index.inner.clone(), null_on_oob)
             .into()
     }
@@ -129,7 +134,7 @@ impl RbExpr {
         Ok(self
             .inner
             .clone()
-            .arr()
+            .list()
             .to_struct(width_strat.0, name_gen, upper_bound)
             .into())
     }
@@ -138,9 +143,9 @@ impl RbExpr {
         let e = self.inner.clone();
 
         if maintain_order {
-            e.arr().unique_stable().into()
+            e.list().unique_stable().into()
         } else {
-            e.arr().unique().into()
+            e.list().unique().into()
         }
     }
 }

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

@@ -20,12 +20,10 @@ impl RbExpr {
             strict,
             exact,
             cache,
-            ..Default::default()
         };
         self.inner.clone().str().to_date(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn str_to_datetime(
         &self,
         format: Option<String>,
@@ -34,16 +32,12 @@ impl RbExpr {
         strict: bool,
         exact: bool,
         cache: bool,
-        utc: bool,
-        tz_aware: bool,
     ) -> Self {
         let options = StrptimeOptions {
             format,
             strict,
             exact,
             cache,
-            tz_aware,
-            utc,
         };
         self.inner
             .clone()
@@ -58,7 +52,6 @@ impl RbExpr {
             strict,
             cache,
             exact: true,
-            ..Default::default()
         };
         self.inner.clone().str().to_time(options).into()
     }
@@ -87,6 +80,10 @@ 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()
     }
@@ -233,7 +230,11 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_json_extract(&self, dtype: Option<Wrap<DataType>>) -> Self {
+    pub fn str_json_extract(
+        &self,
+        dtype: Option<Wrap<DataType>>,
+        infer_schema_len: Option<usize>,
+    ) -> Self {
         let dtype = dtype.map(|wrap| wrap.0);
 
         let output_type = match dtype.clone() {
@@ -243,7 +244,7 @@ impl RbExpr {
 
         let function = move |s: Series| {
             let ca = s.utf8()?;
-            match ca.json_extract(dtype.clone()) {
+            match ca.json_extract(dtype.clone(), infer_schema_len) {
                 Ok(ca) => Ok(Some(ca.into_series())),
                 Err(e) => Err(PolarsError::ComputeError(format!("{e:?}").into())),
             }

data/ext/polars/src/expr.rs

@@ -1,3 +1,4 @@
+mod array;
 mod binary;
 mod categorical;
 mod datetime;

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

@@ -1,4 +1,5 @@
-use magnus::{class, RArray, RString, Value};
+use magnus::encoding::{self, EncodingCapable};
+use magnus::{class, Float, Integer, RArray, RString, Value};
 use polars::lazy::dsl;
 use polars::prelude::*;
 
@@ -6,7 +7,7 @@ use crate::apply::lazy::binary_lambda;
 use crate::conversion::{get_lf, get_rbseq, Wrap};
 use crate::prelude::vec_extract_wrapped;
 use crate::rb_exprs_to_exprs;
-use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
+use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
 macro_rules! set_unwrapped_or_0 {
     ($($var:ident),+ $(,)?) => {
@@ -57,7 +58,12 @@ pub fn cols(names: Vec<String>) -> RbExpr {
     dsl::cols(names).into()
 }
 
-pub fn concat_lf(lfs: Value, rechunk: bool, parallel: bool) -> RbResult<RbLazyFrame> {
+pub fn concat_lf(
+    lfs: Value,
+    rechunk: bool,
+    parallel: bool,
+    to_supertypes: bool,
+) -> RbResult<RbLazyFrame> {
     let (seq, len) = get_rbseq(lfs)?;
     let mut lfs = Vec::with_capacity(len);
 
@@ -67,7 +73,15 @@ pub fn concat_lf(lfs: Value, rechunk: bool, parallel: bool) -> RbResult<RbLazyFr
         lfs.push(lf);
     }
 
-    let lf = polars::lazy::dsl::concat(lfs, rechunk, parallel).map_err(RbPolarsErr::from)?;
+    let lf = dsl::concat(
+        lfs,
+        UnionArgs {
+            rechunk,
+            parallel,
+            to_supertypes,
+        },
+    )
+    .map_err(RbPolarsErr::from)?;
     Ok(lf.into())
 }
 
@@ -135,16 +149,11 @@ pub fn cumfold(acc: &RbExpr, lambda: Value, exprs: RArray, include_init: bool) -
     Ok(polars::lazy::dsl::cumfold_exprs(acc.inner.clone(), func, exprs, include_init).into())
 }
 
-// TODO improve
-pub fn lit(value: Value) -> RbResult<RbExpr> {
-    if value.is_nil() {
-        Ok(dsl::lit(Null {}).into())
-    } else if let Ok(series) = value.try_convert::<&RbSeries>() {
-        Ok(dsl::lit(series.series.borrow().clone()).into())
-    } else if let Some(v) = RString::from_value(value) {
-        Ok(dsl::lit(v.try_convert::<String>()?).into())
-    } else if value.is_kind_of(class::integer()) {
-        match value.try_convert::<i64>() {
+pub fn lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
+    if value.is_kind_of(class::true_class()) || value.is_kind_of(class::false_class()) {
+        Ok(dsl::lit(value.try_convert::<bool>()?).into())
+    } else if let Some(v) = Integer::from_value(value) {
+        match v.try_convert::<i64>() {
             Ok(val) => {
                 if val > 0 && val < i32::MAX as i64 || val < 0 && val > i32::MIN as i64 {
                     Ok(dsl::lit(val as i32).into())
@@ -157,17 +166,47 @@ pub fn lit(value: Value) -> RbResult<RbExpr> {
                 Ok(dsl::lit(val).into())
             }
         }
+    } else if let Some(v) = Float::from_value(value) {
+        Ok(dsl::lit(v.try_convert::<f64>()?).into())
+    } else if let Some(v) = RString::from_value(value) {
+        if v.enc_get() == encoding::Index::utf8() {
+            Ok(dsl::lit(v.try_convert::<String>()?).into())
+        } else {
+            Ok(dsl::lit(unsafe { v.as_slice() }).into())
+        }
+    } else if let Ok(series) = value.try_convert::<&RbSeries>() {
+        Ok(dsl::lit(series.series.borrow().clone()).into())
+    } else if value.is_nil() {
+        Ok(dsl::lit(Null {}).into())
+    } else if allow_object {
+        todo!()
     } else {
-        Ok(dsl::lit(value.try_convert::<f64>()?).into())
+        Err(RbValueError::new_err(format!(
+            "could not convert value {:?} as a Literal",
+            value.to_string()
+        )))
     }
 }
 
-pub fn repeat(value: Value, n_times: &RbExpr) -> RbResult<RbExpr> {
-    if value.is_nil() {
-        Ok(polars::lazy::dsl::repeat(Null {}, n_times.inner.clone()).into())
-    } else {
-        todo!();
+pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbResult<RbExpr> {
+    let mut value = value.inner.clone();
+    let n = n.inner.clone();
+
+    if let Some(dtype) = dtype {
+        value = value.cast(dtype.0);
+    }
+
+    if let Expr::Literal(lv) = &value {
+        let av = lv.to_anyvalue().unwrap();
+        // Integer inputs that fit in Int32 are parsed as such
+        if let DataType::Int64 = av.dtype() {
+            let int_value = av.try_extract::<i64>().unwrap();
+            if int_value >= i32::MIN as i64 && int_value <= i32::MAX as i64 {
+                value = value.cast(DataType::Int32);
+            }
+        }
     }
+    Ok(dsl::repeat(value, n).into())
 }
 
 pub fn pearson_corr(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
@@ -203,7 +242,8 @@ pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
     Ok(crate::functions::lazy::dtype_cols(dtypes))
 }
 
+// TODO rename to sum_horizontal
 pub fn sum_exprs(exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
-    Ok(polars::lazy::dsl::sum_exprs(exprs).into())
+    Ok(polars::lazy::dsl::sum_horizontal(exprs).into())
 }

data/ext/polars/src/lazyframe.rs

@@ -230,7 +230,13 @@ impl RbLazyFrame {
         ldf.into()
     }
 
-    pub fn sort(&self, by_column: String, reverse: bool, nulls_last: bool) -> Self {
+    pub fn sort(
+        &self,
+        by_column: String,
+        reverse: bool,
+        nulls_last: bool,
+        maintain_order: bool,
+    ) -> Self {
         let ldf = self.ldf.clone();
         ldf.sort(
             &by_column,
@@ -238,6 +244,7 @@ impl RbLazyFrame {
                 descending: reverse,
                 nulls_last,
                 multithreaded: true,
+                maintain_order,
             },
         )
         .into()
@@ -248,10 +255,13 @@ impl RbLazyFrame {
         by_column: RArray,
         reverse: Vec<bool>,
         nulls_last: bool,
+        maintain_order: bool,
     ) -> RbResult<Self> {
         let ldf = self.ldf.clone();
         let exprs = rb_exprs_to_exprs(by_column)?;
-        Ok(ldf.sort_by_exprs(exprs, reverse, nulls_last).into())
+        Ok(ldf
+            .sort_by_exprs(exprs, reverse, nulls_last, maintain_order)
+            .into())
     }
 
     pub fn cache(&self) -> Self {
@@ -328,6 +338,7 @@ impl RbLazyFrame {
         offset: String,
         closed: Wrap<ClosedWindow>,
         by: RArray,
+        check_sorted: bool,
     ) -> RbResult<RbLazyGroupBy> {
         let closed_window = closed.0;
         let ldf = self.ldf.clone();
@@ -340,6 +351,7 @@ impl RbLazyFrame {
                 period: Duration::parse(&period),
                 offset: Duration::parse(&offset),
                 closed_window,
+                check_sorted,
             },
         );