polars-df 0.8.0 → 0.10.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::{HiveOptions, 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())
@@ -95,17 +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])?;
+        let truncate_ragged_lines = bool::try_convert(arguments[20])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.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
@@ -113,6 +115,7 @@ impl RbLazyFrame {
                 .map(|(name, dtype)| Field::new(&name, dtype.0))
                 .collect::<Schema>()
         });
+
         let r = LazyCsvReader::new(path)
             .with_infer_schema_length(infer_schema_length)
             .with_separator(separator)
@@ -122,6 +125,7 @@ impl RbLazyFrame {
             .with_n_rows(n_rows)
             .with_cache(cache)
             .with_dtype_overwrite(overwrite_dtype.as_ref())
+            // TODO add with_schema
             .low_memory(low_memory)
             .with_comment_prefix(comment_prefix.as_deref())
             .with_quote_char(quote_char)
@@ -129,9 +133,11 @@ impl RbLazyFrame {
             .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);
+            .with_null_values(null_values)
+            // TODO add with_missing_is_null
+            .truncate_ragged_lines(truncate_ragged_lines);
 
         if let Some(_lambda) = with_schema_modify {
             todo!();
@@ -142,30 +148,53 @@ impl RbLazyFrame {
 
     #[allow(clippy::too_many_arguments)]
     pub fn new_from_parquet(
-        path: String,
+        path: Option<PathBuf>,
+        paths: Vec<PathBuf>,
         n_rows: Option<usize>,
         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,
+        hive_schema: Option<Wrap<Schema>>,
     ) -> RbResult<Self> {
-        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let parallel = parallel.0;
+        let hive_schema = hive_schema.map(|s| Arc::new(s.0));
+
+        let first_path = if let Some(path) = &path {
+            path
+        } else {
+            paths
+                .first()
+                .ok_or_else(|| RbValueError::new_err("expected a path argument".to_string()))?
+        };
+
+        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let hive_options = HiveOptions {
+            enabled: hive_partitioning,
+            schema: hive_schema,
+        };
+
         let args = ScanArgsParquet {
             n_rows,
             cache,
-            parallel: parallel.0,
+            parallel,
             rechunk,
-            row_count,
+            row_index,
             low_memory,
-            // TODO support cloud options
             cloud_options: None,
             use_statistics,
-            hive_partitioning,
+            hive_options,
         };
-        let lf = LazyFrame::scan_parquet(path, args).map_err(RbPolarsErr::from)?;
+
+        let lf = if path.is_some() {
+            LazyFrame::scan_parquet(first_path, args)
+        } else {
+            LazyFrame::scan_parquet_files(Arc::from(paths), args)
+        }
+        .map_err(RbPolarsErr::from)?;
         Ok(lf.into())
     }
 
@@ -174,16 +203,17 @@ 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,
-            memmap: memory_map,
+            row_index,
+            memory_map,
+            cloud_options: None,
         };
         let lf = LazyFrame::scan_ipc(path, args).map_err(RbPolarsErr::from)?;
         Ok(lf.into())
@@ -216,37 +246,42 @@ 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()
     }
 
     pub fn sort(
         &self,
         by_column: String,
-        reverse: bool,
+        descending: bool,
         nulls_last: bool,
         maintain_order: bool,
+        multithreaded: bool,
     ) -> Self {
         let ldf = self.ldf.clone();
         ldf.sort(
-            &by_column,
-            SortOptions {
-                descending: reverse,
+            [&by_column],
+            SortMultipleOptions {
+                descending: vec![descending],
                 nulls_last,
-                multithreaded: true,
+                multithreaded,
                 maintain_order,
             },
         )
@@ -255,15 +290,24 @@ impl RbLazyFrame {
 
     pub fn sort_by_exprs(
         &self,
-        by_column: RArray,
-        reverse: Vec<bool>,
+        by: RArray,
+        descending: Vec<bool>,
         nulls_last: bool,
         maintain_order: bool,
+        multithreaded: bool,
     ) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        let exprs = rb_exprs_to_exprs(by_column)?;
+        let exprs = rb_exprs_to_exprs(by)?;
         Ok(ldf
-            .sort_by_exprs(exprs, reverse, nulls_last, maintain_order)
+            .sort_by_exprs(
+                exprs,
+                SortMultipleOptions {
+                    descending,
+                    nulls_last,
+                    maintain_order,
+                    multithreaded,
+                },
+            )
             .into())
     }
 
@@ -304,6 +348,76 @@ 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(())
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    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)?;
@@ -321,6 +435,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)?;
@@ -334,7 +454,7 @@ impl RbLazyFrame {
         })
     }
 
-    pub fn group_by_rolling(
+    pub fn rolling(
         &self,
         index_column: &RbExpr,
         period: String,
@@ -346,7 +466,7 @@ impl RbLazyFrame {
         let closed_window = closed.0;
         let ldf = self.ldf.clone();
         let by = rb_exprs_to_exprs(by)?;
-        let lazy_gb = ldf.group_by_rolling(
+        let lazy_gb = ldf.rolling(
             index_column.inner.clone(),
             by,
             RollingGroupOptions {
@@ -459,6 +579,7 @@ impl RbLazyFrame {
         right_on: RArray,
         allow_parallel: bool,
         force_parallel: bool,
+        join_nulls: bool,
         how: Wrap<JoinType>,
         suffix: String,
     ) -> RbResult<Self> {
@@ -474,17 +595,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()
@@ -569,6 +701,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,
@@ -619,14 +756,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 {
@@ -668,4 +809,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())
+    }
 }