polars-df 0.1.0 → 0.1.1

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

@@ -1,10 +1,13 @@
-use magnus::RArray;
+use magnus::{RArray, RHash, Value};
 use polars::lazy::frame::{LazyFrame, LazyGroupBy};
+use polars::prelude::*;
 use std::cell::RefCell;
+use std::io::BufWriter;
 
-use crate::conversion::wrap_join_type;
+use crate::conversion::*;
+use crate::file::get_file_like;
 use crate::lazy::utils::rb_exprs_to_exprs;
-use crate::{RbDataFrame, RbExpr, RbPolarsErr, RbResult};
+use crate::{RbDataFrame, RbExpr, RbPolarsErr, RbResult, RbValueError};
 
 #[magnus::wrap(class = "Polars::RbLazyGroupBy")]
 pub struct RbLazyGroupBy {
@@ -17,6 +20,16 @@ impl RbLazyGroupBy {
         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()
+    }
 }
 
 #[magnus::wrap(class = "Polars::RbLazyFrame")]
@@ -25,6 +38,13 @@ pub struct RbLazyFrame {
     pub ldf: LazyFrame,
 }
 
+impl RbLazyFrame {
+    fn get_schema(&self) -> RbResult<SchemaRef> {
+        let schema = self.ldf.schema().map_err(RbPolarsErr::from)?;
+        Ok(schema)
+    }
+}
+
 impl From<LazyFrame> for RbLazyFrame {
     fn from(ldf: LazyFrame) -> Self {
         RbLazyFrame { ldf }
@@ -32,6 +52,25 @@ impl From<LazyFrame> for RbLazyFrame {
 }
 
 impl RbLazyFrame {
+    pub fn write_json(&self, rb_f: Value) -> RbResult<()> {
+        let file = BufWriter::new(get_file_like(rb_f, true)?);
+        serde_json::to_writer(file, &self.ldf.logical_plan)
+            .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
+        Ok(())
+    }
+
+    pub fn describe_plan(&self) -> String {
+        self.ldf.describe_plan()
+    }
+
+    pub fn describe_optimized_plan(&self) -> RbResult<String> {
+        let result = self
+            .ldf
+            .describe_optimized_plan()
+            .map_err(RbPolarsErr::from)?;
+        Ok(result)
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn optimization_toggle(
         &self,
@@ -55,18 +94,52 @@ impl RbLazyFrame {
         ldf.into()
     }
 
+    pub fn sort(&self, by_column: String, reverse: bool, nulls_last: bool) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.sort(
+            &by_column,
+            SortOptions {
+                descending: reverse,
+                nulls_last,
+            },
+        )
+        .into()
+    }
+
+    pub fn sort_by_exprs(
+        &self,
+        by_column: RArray,
+        reverse: Vec<bool>,
+        nulls_last: 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())
+    }
+
+    pub fn cache(&self) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.cache().into()
+    }
+
     pub fn collect(&self) -> RbResult<RbDataFrame> {
         let ldf = self.ldf.clone();
         let df = ldf.collect().map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }
 
-    pub fn filter(&self, predicate: &RbExpr) -> RbLazyFrame {
+    pub fn fetch(&self, n_rows: usize) -> RbResult<RbDataFrame> {
+        let ldf = self.ldf.clone();
+        let df = ldf.fetch(n_rows).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn filter(&self, predicate: &RbExpr) -> Self {
         let ldf = self.ldf.clone();
         ldf.filter(predicate.inner.clone()).into()
     }
 
-    pub fn select(&self, exprs: RArray) -> RbResult<RbLazyFrame> {
+    pub fn select(&self, exprs: RArray) -> RbResult<Self> {
         let ldf = self.ldf.clone();
         let exprs = rb_exprs_to_exprs(exprs)?;
         Ok(ldf.select(exprs).into())
@@ -85,6 +158,65 @@ impl RbLazyFrame {
         })
     }
 
+    pub fn groupby_rolling(
+        &self,
+        index_column: String,
+        period: String,
+        offset: String,
+        closed: Wrap<ClosedWindow>,
+        by: RArray,
+    ) -> 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(
+            by,
+            RollingGroupOptions {
+                index_column,
+                period: Duration::parse(&period),
+                offset: Duration::parse(&offset),
+                closed_window,
+            },
+        );
+
+        Ok(RbLazyGroupBy {
+            lgb: RefCell::new(Some(lazy_gb)),
+        })
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn groupby_dynamic(
+        &self,
+        index_column: String,
+        every: String,
+        period: String,
+        offset: String,
+        truncate: bool,
+        include_boundaries: bool,
+        closed: Wrap<ClosedWindow>,
+        by: RArray,
+    ) -> RbResult<RbLazyGroupBy> {
+        let closed_window = closed.0;
+        let by = rb_exprs_to_exprs(by)?;
+        let ldf = self.ldf.clone();
+        let lazy_gb = ldf.groupby_dynamic(
+            by,
+            DynamicGroupOptions {
+                index_column,
+                every: Duration::parse(&every),
+                period: Duration::parse(&period),
+                offset: Duration::parse(&offset),
+                truncate,
+                include_boundaries,
+                closed_window,
+            },
+        );
+
+        Ok(RbLazyGroupBy {
+            lgb: RefCell::new(Some(lazy_gb)),
+        })
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn join(
         &self,
@@ -93,11 +225,9 @@ impl RbLazyFrame {
         right_on: RArray,
         allow_parallel: bool,
         force_parallel: bool,
-        how: String,
+        how: Wrap<JoinType>,
         suffix: String,
     ) -> RbResult<Self> {
-        let how = wrap_join_type(&how)?;
-
         let ldf = self.ldf.clone();
         let other = other.ldf.clone();
         let left_on = rb_exprs_to_exprs(left_on)?;
@@ -110,14 +240,178 @@ impl RbLazyFrame {
             .right_on(right_on)
             .allow_parallel(allow_parallel)
             .force_parallel(force_parallel)
-            .how(how)
+            .how(how.0)
             .suffix(suffix)
             .finish()
             .into())
     }
 
-    pub fn with_columns(&self, exprs: RArray) -> RbResult<RbLazyFrame> {
+    pub fn with_columns(&self, exprs: RArray) -> RbResult<Self> {
         let ldf = self.ldf.clone();
         Ok(ldf.with_columns(rb_exprs_to_exprs(exprs)?).into())
     }
+
+    pub fn rename(&self, existing: Vec<String>, new: Vec<String>) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.rename(existing, new).into()
+    }
+
+    pub fn reverse(&self) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.reverse().into()
+    }
+
+    pub fn shift(&self, periods: i64) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.shift(periods).into()
+    }
+
+    pub fn shift_and_fill(&self, periods: i64, fill_value: &RbExpr) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.shift_and_fill(periods, fill_value.inner.clone()).into()
+    }
+
+    pub fn fill_nan(&self, fill_value: &RbExpr) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.fill_nan(fill_value.inner.clone()).into()
+    }
+
+    pub fn min(&self) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.min().into()
+    }
+
+    pub fn max(&self) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.max().into()
+    }
+
+    pub fn sum(&self) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.sum().into()
+    }
+
+    pub fn mean(&self) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.mean().into()
+    }
+
+    pub fn std(&self, ddof: u8) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.std(ddof).into()
+    }
+
+    pub fn var(&self, ddof: u8) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.var(ddof).into()
+    }
+
+    pub fn median(&self) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.median().into()
+    }
+
+    pub fn quantile(&self, quantile: f64, interpolation: Wrap<QuantileInterpolOptions>) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.quantile(quantile, interpolation.0).into()
+    }
+
+    pub fn explode(&self, column: RArray) -> RbResult<Self> {
+        let ldf = self.ldf.clone();
+        let column = rb_exprs_to_exprs(column)?;
+        Ok(ldf.explode(column).into())
+    }
+
+    pub fn unique(
+        &self,
+        maintain_order: bool,
+        subset: Option<Vec<String>>,
+        keep: Wrap<UniqueKeepStrategy>,
+    ) -> RbResult<Self> {
+        let ldf = self.ldf.clone();
+        Ok(match maintain_order {
+            true => ldf.unique_stable(subset, keep.0),
+            false => ldf.unique(subset, keep.0),
+        }
+        .into())
+    }
+
+    pub fn drop_nulls(&self, subset: Option<Vec<String>>) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.drop_nulls(subset.map(|v| v.into_iter().map(|s| col(&s)).collect()))
+            .into()
+    }
+
+    pub fn slice(&self, offset: i64, len: Option<IdxSize>) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.slice(offset, len.unwrap_or(IdxSize::MAX)).into()
+    }
+
+    pub fn tail(&self, n: IdxSize) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.tail(n).into()
+    }
+
+    pub fn melt(
+        &self,
+        id_vars: Vec<String>,
+        value_vars: Vec<String>,
+        value_name: Option<String>,
+        variable_name: Option<String>,
+    ) -> Self {
+        let args = MeltArgs {
+            id_vars,
+            value_vars,
+            value_name,
+            variable_name,
+        };
+
+        let ldf = self.ldf.clone();
+        ldf.melt(args).into()
+    }
+
+    pub fn with_row_count(&self, name: String, offset: Option<IdxSize>) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.with_row_count(&name, offset).into()
+    }
+
+    pub fn drop_columns(&self, cols: Vec<String>) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.drop_columns(cols).into()
+    }
+
+    pub fn clone(&self) -> Self {
+        self.ldf.clone().into()
+    }
+
+    pub fn columns(&self) -> RbResult<Vec<String>> {
+        Ok(self.get_schema()?.iter_names().cloned().collect())
+    }
+
+    pub fn dtypes(&self) -> RbResult<Vec<String>> {
+        let schema = self.get_schema()?;
+        let iter = schema.iter_dtypes().map(|dt| dt.to_string());
+        Ok(iter.collect())
+    }
+
+    pub fn schema(&self) -> RbResult<RHash> {
+        let schema = self.get_schema()?;
+        let schema_dict = RHash::new();
+
+        schema.iter_fields().for_each(|fld| {
+            // TODO remove unwrap
+            schema_dict
+                .aset(fld.name().clone(), fld.data_type().to_string())
+                .unwrap();
+        });
+        Ok(schema_dict)
+    }
+
+    pub fn unnest(&self, cols: Vec<String>) -> Self {
+        self.ldf.clone().unnest(cols).into()
+    }
+
+    pub fn width(&self) -> RbResult<usize> {
+        Ok(self.get_schema()?.len())
+    }
 }