polars-df 0.7.0 → 0.9.0

data/ext/polars/src/{lazyframe.rs → lazyframe/mod.rs}

@@ -1,9 +1,10 @@
 use magnus::{IntoValue, RArray, RHash, TryConvert, Value};
-use polars::io::RowCount;
+use polars::io::RowIndex;
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
 use std::cell::RefCell;
 use std::io::{BufWriter, Read};
+use std::num::NonZeroUsize;
 use std::path::PathBuf;
 
 use crate::conversion::*;
@@ -55,13 +56,14 @@ impl RbLazyFrame {
     pub fn new_from_ndjson(
         path: String,
         infer_schema_length: Option<usize>,
-        batch_size: Option<usize>,
+        batch_size: Option<Wrap<NonZeroUsize>>,
         n_rows: Option<usize>,
         low_memory: bool,
         rechunk: bool,
-        row_count: Option<(String, IdxSize)>,
+        row_index: Option<(String, IdxSize)>,
     ) -> RbResult<Self> {
-        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let batch_size = batch_size.map(|v| v.0);
+        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
 
         let lf = LazyJsonLineReader::new(path)
             .with_infer_schema_length(infer_schema_length)
@@ -69,7 +71,7 @@ impl RbLazyFrame {
             .with_n_rows(n_rows)
             .low_memory(low_memory)
             .with_rechunk(rechunk)
-            .with_row_count(row_count)
+            .with_row_index(row_index)
             .finish()
             .map_err(RbPolarsErr::from)?;
         Ok(lf.into())
@@ -87,7 +89,7 @@ impl RbLazyFrame {
         let cache = bool::try_convert(arguments[6])?;
         let overwrite_dtype = Option::<Vec<(String, Wrap<DataType>)>>::try_convert(arguments[7])?;
         let low_memory = bool::try_convert(arguments[8])?;
-        let comment_char = Option::<String>::try_convert(arguments[9])?;
+        let comment_prefix = Option::<String>::try_convert(arguments[9])?;
         let quote_char = Option::<String>::try_convert(arguments[10])?;
         let null_values = Option::<Wrap<NullValues>>::try_convert(arguments[11])?;
         let infer_schema_length = Option::<usize>::try_convert(arguments[12])?;
@@ -95,18 +97,17 @@ impl RbLazyFrame {
         let rechunk = bool::try_convert(arguments[14])?;
         let skip_rows_after_header = usize::try_convert(arguments[15])?;
         let encoding = Wrap::<CsvEncoding>::try_convert(arguments[16])?;
-        let row_count = Option::<(String, IdxSize)>::try_convert(arguments[17])?;
+        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[17])?;
         let try_parse_dates = bool::try_convert(arguments[18])?;
         let eol_char = String::try_convert(arguments[19])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
-        let comment_char = comment_char.map(|s| s.as_bytes()[0]);
         let quote_char = quote_char.map(|s| s.as_bytes()[0]);
         let separator = separator.as_bytes()[0];
         let eol_char = eol_char.as_bytes()[0];
 
-        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
 
         let overwrite_dtype = overwrite_dtype.map(|overwrite_dtype| {
             overwrite_dtype
@@ -124,13 +125,13 @@ impl RbLazyFrame {
             .with_cache(cache)
             .with_dtype_overwrite(overwrite_dtype.as_ref())
             .low_memory(low_memory)
-            .with_comment_char(comment_char)
+            .with_comment_prefix(comment_prefix.as_deref())
             .with_quote_char(quote_char)
             .with_end_of_line_char(eol_char)
             .with_rechunk(rechunk)
             .with_skip_rows_after_header(skip_rows_after_header)
             .with_encoding(encoding.0)
-            .with_row_count(row_count)
+            .with_row_index(row_index)
             .with_try_parse_dates(try_parse_dates)
             .with_null_values(null_values);
 
@@ -148,18 +149,18 @@ impl RbLazyFrame {
         cache: bool,
         parallel: Wrap<ParallelStrategy>,
         rechunk: bool,
-        row_count: Option<(String, IdxSize)>,
+        row_index: Option<(String, IdxSize)>,
         low_memory: bool,
         use_statistics: bool,
         hive_partitioning: bool,
     ) -> RbResult<Self> {
-        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
         let args = ScanArgsParquet {
             n_rows,
             cache,
             parallel: parallel.0,
             rechunk,
-            row_count,
+            row_index,
             low_memory,
             // TODO support cloud options
             cloud_options: None,
@@ -175,15 +176,15 @@ impl RbLazyFrame {
         n_rows: Option<usize>,
         cache: bool,
         rechunk: bool,
-        row_count: Option<(String, IdxSize)>,
+        row_index: Option<(String, IdxSize)>,
         memory_map: bool,
     ) -> RbResult<Self> {
-        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
         let args = ScanArgsIpc {
             n_rows,
             cache,
             rechunk,
-            row_count,
+            row_index,
             memmap: memory_map,
         };
         let lf = LazyFrame::scan_ipc(path, args).map_err(RbPolarsErr::from)?;
@@ -217,20 +218,24 @@ impl RbLazyFrame {
         projection_pushdown: bool,
         simplify_expr: bool,
         slice_pushdown: bool,
-        cse: bool,
+        comm_subplan_elim: bool,
+        comm_subexpr_elim: bool,
         allow_streaming: bool,
         _eager: bool,
     ) -> RbLazyFrame {
         let ldf = self.ldf.clone();
-        let ldf = ldf
+        let mut ldf = ldf
             .with_type_coercion(type_coercion)
             .with_predicate_pushdown(predicate_pushdown)
             .with_simplify_expr(simplify_expr)
             .with_slice_pushdown(slice_pushdown)
-            .with_comm_subplan_elim(cse)
             .with_streaming(allow_streaming)
             ._with_eager(_eager)
             .with_projection_pushdown(projection_pushdown);
+
+        ldf = ldf.with_comm_subplan_elim(comm_subplan_elim);
+        ldf = ldf.with_comm_subexpr_elim(comm_subexpr_elim);
+
         ldf.into()
     }
 
@@ -305,6 +310,75 @@ impl RbLazyFrame {
         Ok(())
     }
 
+    pub fn sink_ipc(
+        &self,
+        path: PathBuf,
+        compression: Option<Wrap<IpcCompression>>,
+        maintain_order: bool,
+    ) -> RbResult<()> {
+        let options = IpcWriterOptions {
+            compression: compression.map(|c| c.0),
+            maintain_order,
+        };
+
+        let ldf = self.ldf.clone();
+        ldf.sink_ipc(path, options).map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn sink_csv(
+        &self,
+        path: PathBuf,
+        include_bom: bool,
+        include_header: bool,
+        separator: u8,
+        line_terminator: String,
+        quote_char: u8,
+        batch_size: Wrap<NonZeroUsize>,
+        datetime_format: Option<String>,
+        date_format: Option<String>,
+        time_format: Option<String>,
+        float_precision: Option<usize>,
+        null_value: Option<String>,
+        quote_style: Option<Wrap<QuoteStyle>>,
+        maintain_order: bool,
+    ) -> RbResult<()> {
+        let quote_style = quote_style.map_or(QuoteStyle::default(), |wrap| wrap.0);
+        let null_value = null_value.unwrap_or(SerializeOptions::default().null);
+
+        let serialize_options = SerializeOptions {
+            date_format,
+            time_format,
+            datetime_format,
+            float_precision,
+            separator,
+            quote_char,
+            null: null_value,
+            line_terminator,
+            quote_style,
+        };
+
+        let options = CsvWriterOptions {
+            include_bom,
+            include_header,
+            maintain_order,
+            batch_size: batch_size.0,
+            serialize_options,
+        };
+
+        let ldf = self.ldf.clone();
+        ldf.sink_csv(path, options).map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn sink_json(&self, path: PathBuf, maintain_order: bool) -> RbResult<()> {
+        let options = JsonWriterOptions { maintain_order };
+
+        let ldf = self.ldf.clone();
+        ldf.sink_json(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)?;
@@ -322,6 +396,12 @@ impl RbLazyFrame {
         Ok(ldf.select(exprs).into())
     }
 
+    pub fn select_seq(&self, exprs: RArray) -> RbResult<Self> {
+        let ldf = self.ldf.clone();
+        let exprs = rb_exprs_to_exprs(exprs)?;
+        Ok(ldf.select_seq(exprs).into())
+    }
+
     pub fn group_by(&self, by: RArray, maintain_order: bool) -> RbResult<RbLazyGroupBy> {
         let ldf = self.ldf.clone();
         let by = rb_exprs_to_exprs(by)?;
@@ -335,7 +415,7 @@ impl RbLazyFrame {
         })
     }
 
-    pub fn group_by_rolling(
+    pub fn rolling(
         &self,
         index_column: &RbExpr,
         period: String,
@@ -460,6 +540,7 @@ impl RbLazyFrame {
         right_on: RArray,
         allow_parallel: bool,
         force_parallel: bool,
+        join_nulls: bool,
         how: Wrap<JoinType>,
         suffix: String,
     ) -> RbResult<Self> {
@@ -475,17 +556,28 @@ impl RbLazyFrame {
             .right_on(right_on)
             .allow_parallel(allow_parallel)
             .force_parallel(force_parallel)
+            .join_nulls(join_nulls)
             .how(how.0)
             .suffix(suffix)
             .finish()
             .into())
     }
 
+    pub fn with_column(&self, expr: &RbExpr) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.with_column(expr.inner.clone()).into()
+    }
+
     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 with_columns_seq(&self, exprs: RArray) -> RbResult<Self> {
+        let ldf = self.ldf.clone();
+        Ok(ldf.with_columns_seq(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()
@@ -510,48 +602,58 @@ impl RbLazyFrame {
         ldf.fill_nan(fill_value.inner.clone()).into()
     }
 
-    pub fn min(&self) -> Self {
+    pub fn min(&self) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        ldf.min().into()
+        let out = ldf.min().map_err(RbPolarsErr::from)?;
+        Ok(out.into())
     }
 
-    pub fn max(&self) -> Self {
+    pub fn max(&self) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        ldf.max().into()
+        let out = ldf.max().map_err(RbPolarsErr::from)?;
+        Ok(out.into())
     }
 
-    pub fn sum(&self) -> Self {
+    pub fn sum(&self) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        ldf.sum().into()
+        let out = ldf.sum().map_err(RbPolarsErr::from)?;
+        Ok(out.into())
     }
 
-    pub fn mean(&self) -> Self {
+    pub fn mean(&self) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        ldf.mean().into()
+        let out = ldf.mean().map_err(RbPolarsErr::from)?;
+        Ok(out.into())
     }
 
-    pub fn std(&self, ddof: u8) -> Self {
+    pub fn std(&self, ddof: u8) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        ldf.std(ddof).into()
+        let out = ldf.std(ddof).map_err(RbPolarsErr::from)?;
+        Ok(out.into())
     }
 
-    pub fn var(&self, ddof: u8) -> Self {
+    pub fn var(&self, ddof: u8) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        ldf.var(ddof).into()
+        let out = ldf.var(ddof).map_err(RbPolarsErr::from)?;
+        Ok(out.into())
     }
 
-    pub fn median(&self) -> Self {
+    pub fn median(&self) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        ldf.median().into()
+        let out = ldf.median().map_err(RbPolarsErr::from)?;
+        Ok(out.into())
     }
 
     pub fn quantile(
         &self,
         quantile: &RbExpr,
         interpolation: Wrap<QuantileInterpolOptions>,
-    ) -> Self {
+    ) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        ldf.quantile(quantile.inner.clone(), interpolation.0).into()
+        let out = ldf
+            .quantile(quantile.inner.clone(), interpolation.0)
+            .map_err(RbPolarsErr::from)?;
+        Ok(out.into())
     }
 
     pub fn explode(&self, column: RArray) -> RbResult<Self> {
@@ -560,6 +662,11 @@ impl RbLazyFrame {
         Ok(ldf.explode(column).into())
     }
 
+    pub fn null_count(&self) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.null_count().into()
+    }
+
     pub fn unique(
         &self,
         maintain_order: bool,
@@ -610,14 +717,18 @@ impl RbLazyFrame {
         ldf.melt(args).into()
     }
 
-    pub fn with_row_count(&self, name: String, offset: Option<IdxSize>) -> Self {
+    pub fn with_row_index(&self, name: String, offset: Option<IdxSize>) -> Self {
         let ldf = self.ldf.clone();
-        ldf.with_row_count(&name, offset).into()
+        ldf.with_row_index(&name, offset).into()
     }
 
-    pub fn drop_columns(&self, cols: Vec<String>) -> Self {
+    pub fn drop(&self, cols: Vec<String>) -> Self {
         let ldf = self.ldf.clone();
-        ldf.drop_columns(cols).into()
+        ldf.drop(cols).into()
+    }
+
+    pub fn cast_all(&self, dtype: Wrap<DataType>, strict: bool) -> Self {
+        self.ldf.clone().cast_all(dtype.0, strict).into()
     }
 
     pub fn clone(&self) -> Self {
@@ -659,4 +770,18 @@ impl RbLazyFrame {
     pub fn width(&self) -> RbResult<usize> {
         Ok(self.get_schema()?.len())
     }
+
+    pub fn count(&self) -> Self {
+        let ldf = self.ldf.clone();
+        ldf.count().into()
+    }
+
+    pub fn merge_sorted(&self, other: &Self, key: String) -> RbResult<Self> {
+        let out = self
+            .ldf
+            .clone()
+            .merge_sorted(other.ldf.clone(), &key)
+            .map_err(RbPolarsErr::from)?;
+        Ok(out.into())
+    }
 }