polars-df 0.3.1 → 0.5.0

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

@@ -0,0 +1,209 @@
+use magnus::{class, 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};
+
+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) -> 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 = polars::lazy::dsl::concat(lfs, rechunk, parallel).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())
+}
+
+// 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>() {
+            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 {
+        Ok(dsl::lit(value.try_convert::<f64>()?).into())
+    }
+}
+
+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 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))
+}
+
+pub fn sum_exprs(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    Ok(polars::lazy::dsl::sum_exprs(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,37 +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)]
@@ -118,7 +96,7 @@ impl RbLazyFrame {
         let skip_rows_after_header: usize = arguments[15].try_convert()?;
         let encoding: Wrap<CsvEncoding> = arguments[16].try_convert()?;
         let row_count: Option<(String, IdxSize)> = arguments[17].try_convert()?;
-        let parse_dates: bool = arguments[18].try_convert()?;
+        let try_parse_dates: bool = arguments[18].try_convert()?;
         let eol_char: String = arguments[19].try_convert()?;
         // end arguments
 
@@ -131,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)
@@ -153,7 +131,7 @@ impl RbLazyFrame {
             .with_skip_rows_after_header(skip_rows_after_header)
             .with_encoding(encoding.0)
             .with_row_count(row_count)
-            .with_parse_dates(parse_dates)
+            .with_try_parse_dates(try_parse_dates)
             .with_null_values(null_values);
 
         if let Some(_lambda) = with_schema_modify {
@@ -163,6 +141,7 @@ impl RbLazyFrame {
         Ok(r.finish().map_err(RbPolarsErr::from)?.into())
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn new_from_parquet(
         path: String,
         n_rows: Option<usize>,
@@ -171,6 +150,7 @@ impl RbLazyFrame {
         rechunk: bool,
         row_count: Option<(String, IdxSize)>,
         low_memory: bool,
+        use_statistics: bool,
     ) -> RbResult<Self> {
         let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
         let args = ScanArgsParquet {
@@ -182,6 +162,7 @@ impl RbLazyFrame {
             low_memory,
             // TODO support cloud options
             cloud_options: None,
+            use_statistics,
         };
         let lf = LazyFrame::scan_parquet(path, args).map_err(RbPolarsErr::from)?;
         Ok(lf.into())
@@ -284,6 +265,32 @@ impl RbLazyFrame {
         Ok(df.into())
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub fn sink_parquet(
+        &self,
+        path: PathBuf,
+        compression: String,
+        compression_level: Option<i32>,
+        statistics: bool,
+        row_group_size: Option<usize>,
+        data_pagesize_limit: Option<usize>,
+        maintain_order: bool,
+    ) -> RbResult<()> {
+        let compression = parse_parquet_compression(&compression, compression_level)?;
+
+        let options = ParquetWriteOptions {
+            compression,
+            statistics,
+            row_group_size,
+            data_pagesize_limit,
+            maintain_order,
+        };
+
+        let ldf = self.ldf.clone();
+        ldf.sink_parquet(path, options).map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
     pub fn fetch(&self, n_rows: usize) -> RbResult<RbDataFrame> {
         let ldf = self.ldf.clone();
         let df = ldf.fetch(n_rows).map_err(RbPolarsErr::from)?;
@@ -316,7 +323,7 @@ impl RbLazyFrame {
 
     pub fn groupby_rolling(
         &self,
-        index_column: String,
+        index_column: &RbExpr,
         period: String,
         offset: String,
         closed: Wrap<ClosedWindow>,
@@ -326,9 +333,10 @@ impl RbLazyFrame {
         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(),
                 period: Duration::parse(&period),
                 offset: Duration::parse(&offset),
                 closed_window,
@@ -343,7 +351,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,
@@ -357,9 +365,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,
                 every: Duration::parse(&every),
                 period: Duration::parse(&period),
                 offset: Duration::parse(&offset),
@@ -367,6 +375,7 @@ impl RbLazyFrame {
                 include_boundaries,
                 closed_window,
                 start_by: start_by.0,
+                ..Default::default()
             },
         );
 
@@ -415,10 +424,10 @@ impl RbLazyFrame {
             .force_parallel(force_parallel)
             .how(JoinType::AsOf(AsOfOptions {
                 strategy: strategy.0,
-                left_by,
-                right_by,
+                left_by: left_by.map(strings_to_smartstrings),
+                right_by: right_by.map(strings_to_smartstrings),
                 tolerance: tolerance.map(|t| t.0.into_static().unwrap()),
-                tolerance_str,
+                tolerance_str: tolerance_str.map(|s| s.into()),
             }))
             .suffix(suffix)
             .finish()
@@ -570,12 +579,14 @@ impl RbLazyFrame {
         value_vars: Vec<String>,
         value_name: Option<String>,
         variable_name: Option<String>,
+        streamable: bool,
     ) -> Self {
         let args = MeltArgs {
-            id_vars,
-            value_vars,
-            value_name,
-            variable_name,
+            id_vars: strings_to_smartstrings(id_vars),
+            value_vars: strings_to_smartstrings(value_vars),
+            value_name: value_name.map(|s| s.into()),
+            variable_name: variable_name.map(|s| s.into()),
+            streamable,
         };
 
         let ldf = self.ldf.clone();
@@ -596,8 +607,10 @@ impl RbLazyFrame {
         self.ldf.clone().into()
     }
 
-    pub fn columns(&self) -> RbResult<Vec<String>> {
-        Ok(self.get_schema()?.iter_names().cloned().collect())
+    pub fn columns(&self) -> RbResult<RArray> {
+        let schema = self.get_schema()?;
+        let iter = schema.iter_names().map(|s| s.as_str());
+        Ok(RArray::from_iter(iter))
     }
 
     pub fn dtypes(&self) -> RbResult<RArray> {
@@ -614,7 +627,7 @@ impl RbLazyFrame {
             // TODO remove unwrap
             schema_dict
                 .aset::<String, Value>(
-                    fld.name().clone(),
+                    fld.name().to_string(),
                     Wrap(fld.data_type().clone()).into_value(),
                 )
                 .unwrap();

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()
+    }
+}