polars-df 0.4.0 → 0.6.0

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

@@ -0,0 +1,249 @@
+use magnus::encoding::{self, EncodingCapable};
+use magnus::{class, Float, Integer, RArray, RString, Value};
+use polars::lazy::dsl;
+use polars::prelude::*;
+
+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, RbValueError};
+
+macro_rules! set_unwrapped_or_0 {
+    ($($var:ident),+ $(,)?) => {
+        $(let $var = $var.map(|e| e.inner.clone()).unwrap_or(polars::lazy::dsl::lit(0));)+
+    };
+}
+
+pub fn arange(low: &RbExpr, high: &RbExpr, step: i64) -> RbExpr {
+    dsl::arange(low.inner.clone(), high.inner.clone(), step).into()
+}
+
+pub fn arg_sort_by(by: RArray, descending: Vec<bool>) -> RbResult<RbExpr> {
+    let by = rb_exprs_to_exprs(by)?;
+    Ok(dsl::arg_sort_by(by, &descending).into())
+}
+
+pub fn arg_where(condition: &RbExpr) -> RbExpr {
+    dsl::arg_where(condition.inner.clone()).into()
+}
+
+pub fn as_struct(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    Ok(dsl::as_struct(&exprs).into())
+}
+
+pub fn coalesce(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    Ok(dsl::coalesce(&exprs).into())
+}
+
+pub fn col(name: String) -> RbExpr {
+    dsl::col(&name).into()
+}
+
+pub fn collect_all(lfs: RArray) -> RbResult<RArray> {
+    let lfs = lfs
+        .each()
+        .map(|v| v?.try_convert::<&RbLazyFrame>())
+        .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
+
+    Ok(RArray::from_iter(lfs.iter().map(|lf| {
+        let df = lf.ldf.clone().collect().unwrap();
+        RbDataFrame::new(df)
+    })))
+}
+
+pub fn cols(names: Vec<String>) -> RbExpr {
+    dsl::cols(names).into()
+}
+
+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);
+
+    for res in seq.each() {
+        let item = res?;
+        let lf = get_lf(item)?;
+        lfs.push(lf);
+    }
+
+    let lf = dsl::concat(
+        lfs,
+        UnionArgs {
+            rechunk,
+            parallel,
+            to_supertypes,
+        },
+    )
+    .map_err(RbPolarsErr::from)?;
+    Ok(lf.into())
+}
+
+#[allow(clippy::too_many_arguments)]
+pub fn duration(
+    days: Option<&RbExpr>,
+    seconds: Option<&RbExpr>,
+    nanoseconds: Option<&RbExpr>,
+    microseconds: Option<&RbExpr>,
+    milliseconds: Option<&RbExpr>,
+    minutes: Option<&RbExpr>,
+    hours: Option<&RbExpr>,
+    weeks: Option<&RbExpr>,
+) -> RbExpr {
+    set_unwrapped_or_0!(
+        days,
+        seconds,
+        nanoseconds,
+        microseconds,
+        milliseconds,
+        minutes,
+        hours,
+        weeks,
+    );
+    let args = DurationArgs {
+        days,
+        seconds,
+        nanoseconds,
+        microseconds,
+        milliseconds,
+        minutes,
+        hours,
+        weeks,
+    };
+    dsl::duration(args).into()
+}
+
+pub fn count() -> RbExpr {
+    dsl::count().into()
+}
+
+pub fn first() -> RbExpr {
+    dsl::first().into()
+}
+
+pub fn last() -> RbExpr {
+    dsl::last().into()
+}
+
+pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
+    dsl::dtype_cols(dtypes).into()
+}
+
+pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+
+    let func = move |a: Series, b: Series| binary_lambda(lambda, a, b);
+    Ok(polars::lazy::dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
+}
+
+pub fn cumfold(acc: &RbExpr, lambda: Value, exprs: RArray, include_init: bool) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+
+    let func = move |a: Series, b: Series| binary_lambda(lambda, a, b);
+    Ok(polars::lazy::dsl::cumfold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+}
+
+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())
+                } else {
+                    Ok(dsl::lit(val).into())
+                }
+            }
+            _ => {
+                let val = value.try_convert::<u64>()?;
+                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 {
+        Err(RbValueError::new_err(format!(
+            "could not convert value {:?} as a Literal",
+            value.to_string()
+        )))
+    }
+}
+
+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 {
+    polars::lazy::dsl::pearson_corr(a.inner.clone(), b.inner.clone(), ddof).into()
+}
+
+pub fn spearman_rank_corr(a: &RbExpr, b: &RbExpr, ddof: u8, propagate_nans: bool) -> RbExpr {
+    polars::lazy::dsl::spearman_rank_corr(a.inner.clone(), b.inner.clone(), ddof, propagate_nans)
+        .into()
+}
+
+pub fn cov(a: &RbExpr, b: &RbExpr) -> RbExpr {
+    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone()).into()
+}
+
+pub fn concat_str(s: RArray, sep: String) -> RbResult<RbExpr> {
+    let s = rb_exprs_to_exprs(s)?;
+    Ok(dsl::concat_str(s, &sep).into())
+}
+
+pub fn concat_lst(s: RArray) -> RbResult<RbExpr> {
+    let s = rb_exprs_to_exprs(s)?;
+    let expr = dsl::concat_list(s).map_err(RbPolarsErr::from)?;
+    Ok(expr.into())
+}
+
+pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
+    let dtypes = dtypes
+        .each()
+        .map(|v| v?.try_convert::<Wrap<DataType>>())
+        .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
+    let dtypes = vec_extract_wrapped(dtypes);
+    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_horizontal(exprs).into())
+}

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

@@ -0,0 +1,8 @@
+use magnus::{IntoValue, Value};
+use polars_core::prelude::IDX_DTYPE;
+
+use crate::conversion::Wrap;
+
+pub fn get_idx_type() -> Value {
+    Wrap(IDX_DTYPE).into_value()
+}

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

@@ -0,0 +1,5 @@
+pub mod eager;
+pub mod io;
+pub mod lazy;
+pub mod meta;
+pub mod whenthen;

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

@@ -0,0 +1,43 @@
+use polars::lazy::dsl;
+
+use crate::RbExpr;
+
+#[magnus::wrap(class = "Polars::RbWhen")]
+#[derive(Clone)]
+pub struct RbWhen {
+    pub inner: dsl::When,
+}
+
+impl From<dsl::When> for RbWhen {
+    fn from(inner: dsl::When) -> Self {
+        RbWhen { inner }
+    }
+}
+
+#[magnus::wrap(class = "Polars::RbWhenThen")]
+#[derive(Clone)]
+pub struct RbWhenThen {
+    pub inner: dsl::WhenThen,
+}
+
+impl From<dsl::WhenThen> for RbWhenThen {
+    fn from(inner: dsl::WhenThen) -> Self {
+        RbWhenThen { inner }
+    }
+}
+
+impl RbWhen {
+    pub fn then(&self, expr: &RbExpr) -> RbWhenThen {
+        self.inner.clone().then(expr.inner.clone()).into()
+    }
+}
+
+impl RbWhenThen {
+    pub fn overwise(&self, expr: &RbExpr) -> RbExpr {
+        self.inner.clone().otherwise(expr.inner.clone()).into()
+    }
+}
+
+pub fn when(predicate: &RbExpr) -> RbWhen {
+    dsl::when(predicate.inner.clone()).into()
+}

data/ext/polars/src/{lazy/dataframe.rs → lazyframe.rs}

@@ -1,38 +1,15 @@
 use magnus::{IntoValue, RArray, RHash, Value};
 use polars::io::RowCount;
-use polars::lazy::frame::{LazyFrame, LazyGroupBy};
+use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
 use std::cell::RefCell;
 use std::io::{BufWriter, Read};
 use std::path::PathBuf;
 
 use crate::conversion::*;
+use crate::expr::rb_exprs_to_exprs;
 use crate::file::get_file_like;
-use crate::lazy::utils::rb_exprs_to_exprs;
-use crate::{RbDataFrame, RbExpr, RbPolarsErr, RbResult, RbValueError};
-
-#[magnus::wrap(class = "Polars::RbLazyGroupBy")]
-pub struct RbLazyGroupBy {
-    lgb: RefCell<Option<LazyGroupBy>>,
-}
-
-impl RbLazyGroupBy {
-    pub fn agg(&self, aggs: RArray) -> RbResult<RbLazyFrame> {
-        let lgb = self.lgb.borrow_mut().take().unwrap();
-        let aggs = rb_exprs_to_exprs(aggs)?;
-        Ok(lgb.agg(aggs).into())
-    }
-
-    pub fn head(&self, n: usize) -> RbLazyFrame {
-        let lgb = self.lgb.take().unwrap();
-        lgb.head(Some(n)).into()
-    }
-
-    pub fn tail(&self, n: usize) -> RbLazyFrame {
-        let lgb = self.lgb.take().unwrap();
-        lgb.tail(Some(n)).into()
-    }
-}
+use crate::{RbDataFrame, RbExpr, RbLazyGroupBy, RbPolarsErr, RbResult, RbValueError};
 
 #[magnus::wrap(class = "Polars::RbLazyFrame")]
 #[derive(Clone)]
@@ -132,10 +109,10 @@ impl RbLazyFrame {
         let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
 
         let overwrite_dtype = overwrite_dtype.map(|overwrite_dtype| {
-            let fields = overwrite_dtype
+            overwrite_dtype
                 .into_iter()
-                .map(|(name, dtype)| Field::new(&name, dtype.0));
-            Schema::from(fields)
+                .map(|(name, dtype)| Field::new(&name, dtype.0))
+                .collect::<Schema>()
         });
         let r = LazyCsvReader::new(path)
             .with_infer_schema_length(infer_schema_length)
@@ -253,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,
@@ -261,6 +244,7 @@ impl RbLazyFrame {
                 descending: reverse,
                 nulls_last,
                 multithreaded: true,
+                maintain_order,
             },
         )
         .into()
@@ -271,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 {
@@ -346,22 +333,25 @@ impl RbLazyFrame {
 
     pub fn groupby_rolling(
         &self,
-        index_column: String,
+        index_column: &RbExpr,
         period: String,
         offset: String,
         closed: Wrap<ClosedWindow>,
         by: RArray,
+        check_sorted: bool,
     ) -> RbResult<RbLazyGroupBy> {
         let closed_window = closed.0;
         let ldf = self.ldf.clone();
         let by = rb_exprs_to_exprs(by)?;
         let lazy_gb = ldf.groupby_rolling(
+            index_column.inner.clone(),
             by,
             RollingGroupOptions {
-                index_column: index_column.into(),
+                index_column: "".into(),
                 period: Duration::parse(&period),
                 offset: Duration::parse(&offset),
                 closed_window,
+                check_sorted,
             },
         );
 
@@ -373,7 +363,7 @@ impl RbLazyFrame {
     #[allow(clippy::too_many_arguments)]
     pub fn groupby_dynamic(
         &self,
-        index_column: String,
+        index_column: &RbExpr,
         every: String,
         period: String,
         offset: String,
@@ -387,9 +377,9 @@ impl RbLazyFrame {
         let by = rb_exprs_to_exprs(by)?;
         let ldf = self.ldf.clone();
         let lazy_gb = ldf.groupby_dynamic(
+            index_column.inner.clone(),
             by,
             DynamicGroupOptions {
-                index_column: index_column.into(),
                 every: Duration::parse(&every),
                 period: Duration::parse(&period),
                 offset: Duration::parse(&offset),
@@ -397,6 +387,7 @@ impl RbLazyFrame {
                 include_boundaries,
                 closed_window,
                 start_by: start_by.0,
+                ..Default::default()
             },
         );
 

data/ext/polars/src/lazygroupby.rs

@@ -0,0 +1,29 @@
+use magnus::RArray;
+use polars::lazy::frame::LazyGroupBy;
+use std::cell::RefCell;
+
+use crate::expr::rb_exprs_to_exprs;
+use crate::{RbLazyFrame, RbResult};
+
+#[magnus::wrap(class = "Polars::RbLazyGroupBy")]
+pub struct RbLazyGroupBy {
+    pub lgb: RefCell<Option<LazyGroupBy>>,
+}
+
+impl RbLazyGroupBy {
+    pub fn agg(&self, aggs: RArray) -> RbResult<RbLazyFrame> {
+        let lgb = self.lgb.borrow_mut().take().unwrap();
+        let aggs = rb_exprs_to_exprs(aggs)?;
+        Ok(lgb.agg(aggs).into())
+    }
+
+    pub fn head(&self, n: usize) -> RbLazyFrame {
+        let lgb = self.lgb.take().unwrap();
+        lgb.head(Some(n)).into()
+    }
+
+    pub fn tail(&self, n: usize) -> RbLazyFrame {
+        let lgb = self.lgb.take().unwrap();
+        lgb.tail(Some(n)).into()
+    }
+}