polars-df 0.8.0 → 0.9.0

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

@@ -17,6 +17,44 @@ macro_rules! set_unwrapped_or_0 {
     };
 }
 
+pub fn rolling_corr(
+    x: &RbExpr,
+    y: &RbExpr,
+    window_size: IdxSize,
+    min_periods: IdxSize,
+    ddof: u8,
+) -> RbExpr {
+    dsl::rolling_corr(
+        x.inner.clone(),
+        y.inner.clone(),
+        RollingCovOptions {
+            min_periods,
+            window_size,
+            ddof,
+        },
+    )
+    .into()
+}
+
+pub fn rolling_cov(
+    x: &RbExpr,
+    y: &RbExpr,
+    window_size: IdxSize,
+    min_periods: IdxSize,
+    ddof: u8,
+) -> RbExpr {
+    dsl::rolling_cov(
+        x.inner.clone(),
+        y.inner.clone(),
+        RollingCovOptions {
+            min_periods,
+            window_size,
+            ddof,
+        },
+    )
+    .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())
@@ -83,6 +121,47 @@ pub fn concat_lf(
     Ok(lf.into())
 }
 
+pub fn concat_list(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 concat_str(s: RArray, separator: String, ignore_nulls: bool) -> RbResult<RbExpr> {
+    let s = rb_exprs_to_exprs(s)?;
+    Ok(dsl::concat_str(s, &separator, ignore_nulls).into())
+}
+
+pub fn len() -> RbExpr {
+    dsl::len().into()
+}
+
+pub fn cov(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
+    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone(), ddof).into()
+}
+
+pub fn arctan2(y: &RbExpr, x: &RbExpr) -> RbExpr {
+    y.inner.clone().arctan2(x.inner.clone()).into()
+}
+
+pub fn arctan2d(y: &RbExpr, x: &RbExpr) -> RbExpr {
+    y.inner.clone().arctan2(x.inner.clone()).degrees().into()
+}
+
+pub fn cum_fold(
+    acc: &RbExpr,
+    lambda: Value,
+    exprs: RArray,
+    include_init: bool,
+) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let lambda = Opaque::from(lambda);
+
+    let func =
+        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
+    Ok(dsl::cum_fold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+}
+
 pub fn concat_lf_diagonal(
     lfs: RArray,
     rechunk: bool,
@@ -110,6 +189,19 @@ pub fn concat_lf_diagonal(
     Ok(lf.into())
 }
 
+pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
+    dsl::dtype_cols(dtypes).into()
+}
+
+pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
+    let dtypes = dtypes
+        .each()
+        .map(|v| Wrap::<DataType>::try_convert(v?))
+        .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
+    let dtypes = vec_extract_wrapped(dtypes);
+    Ok(crate::functions::lazy::dtype_cols(dtypes))
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn duration(
     weeks: Option<&RbExpr>,
@@ -146,38 +238,21 @@ pub fn duration(
     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 lambda = Opaque::from(lambda);
 
     let func =
         move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
-    Ok(polars::lazy::dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
+    Ok(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 lambda = Opaque::from(lambda);
-
-    let func =
-        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
-    Ok(polars::lazy::dsl::cum_fold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+pub fn last() -> RbExpr {
+    dsl::last().into()
 }
 
 pub fn lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
@@ -219,6 +294,10 @@ pub fn lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
     }
 }
 
+pub fn pearson_corr(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
+    dsl::pearson_corr(a.inner.clone(), b.inner.clone(), ddof).into()
+}
+
 pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbResult<RbExpr> {
     let mut value = value.inner.clone();
     let n = n.inner.clone();
@@ -228,7 +307,7 @@ pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbRe
     }
 
     if let Expr::Literal(lv) = &value {
-        let av = lv.to_anyvalue().unwrap();
+        let av = lv.to_any_value().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();
@@ -240,35 +319,11 @@ pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbRe
     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, ddof: u8) -> RbExpr {
-    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone(), ddof).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())
+    dsl::spearman_rank_corr(a.inner.clone(), b.inner.clone(), ddof, propagate_nans).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)?;
+pub fn sql_expr(sql: String) -> RbResult<RbExpr> {
+    let expr = polars::sql::sql_expr(&sql).map_err(RbPolarsErr::from)?;
     Ok(expr.into())
 }
-
-pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
-    let dtypes = dtypes
-        .each()
-        .map(|v| Wrap::<DataType>::try_convert(v?))
-        .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
-    let dtypes = vec_extract_wrapped(dtypes);
-    Ok(crate::functions::lazy::dtype_cols(dtypes))
-}

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

@@ -7,7 +7,7 @@ use polars_core::POOL;
 use crate::conversion::Wrap;
 use crate::{RbResult, RbValueError};
 
-pub fn get_idx_type() -> Value {
+pub fn get_index_type() -> Value {
     Wrap(IDX_DTYPE).into_value()
 }
 
@@ -36,3 +36,47 @@ pub fn get_float_fmt() -> RbResult<String> {
     };
     Ok(strfmt.to_string())
 }
+
+pub fn set_float_precision(precision: Option<usize>) -> RbResult<()> {
+    use polars_core::fmt::set_float_precision;
+    set_float_precision(precision);
+    Ok(())
+}
+
+pub fn get_float_precision() -> RbResult<Option<usize>> {
+    use polars_core::fmt::get_float_precision;
+    Ok(get_float_precision())
+}
+
+pub fn set_thousands_separator(sep: Option<char>) -> RbResult<()> {
+    use polars_core::fmt::set_thousands_separator;
+    set_thousands_separator(sep);
+    Ok(())
+}
+
+pub fn get_thousands_separator() -> RbResult<Option<String>> {
+    use polars_core::fmt::get_thousands_separator;
+    Ok(Some(get_thousands_separator()))
+}
+
+pub fn set_decimal_separator(sep: Option<char>) -> RbResult<()> {
+    use polars_core::fmt::set_decimal_separator;
+    set_decimal_separator(sep);
+    Ok(())
+}
+
+pub fn get_decimal_separator() -> RbResult<Option<char>> {
+    use polars_core::fmt::get_decimal_separator;
+    Ok(Some(get_decimal_separator()))
+}
+
+pub fn set_trim_decimal_zeros(trim: Option<bool>) -> RbResult<()> {
+    use polars_core::fmt::set_trim_decimal_zeros;
+    set_trim_decimal_zeros(trim);
+    Ok(())
+}
+
+pub fn get_trim_decimal_zeros() -> RbResult<Option<bool>> {
+    use polars_core::fmt::get_trim_decimal_zeros;
+    Ok(Some(get_trim_decimal_zeros()))
+}

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

@@ -1,3 +1,7 @@
+use crate::RbResult;
+use magnus::{RArray, Ruby, Value};
+use polars_core::StringCacheHolder;
+
 pub fn enable_string_cache() {
     polars_core::enable_string_cache()
 }
@@ -9,3 +13,13 @@ pub fn disable_string_cache() {
 pub fn using_string_cache() -> bool {
     polars_core::using_string_cache()
 }
+
+#[magnus::wrap(class = "Polars::RbStringCacheHolder")]
+pub struct RbStringCacheHolder {}
+
+impl RbStringCacheHolder {
+    pub fn hold() -> RbResult<Value> {
+        let _hold = StringCacheHolder::hold();
+        Ruby::get().unwrap().yield_splat(RArray::new())
+    }
+}

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())
@@ -95,7 +97,7 @@ 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
@@ -105,7 +107,7 @@ impl RbLazyFrame {
         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
@@ -129,7 +131,7 @@ 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);
 
@@ -147,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,
@@ -174,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)?;
@@ -216,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()
     }
 
@@ -304,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)?;
@@ -321,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)?;
@@ -334,7 +415,7 @@ impl RbLazyFrame {
         })
     }
 
-    pub fn group_by_rolling(
+    pub fn rolling(
         &self,
         index_column: &RbExpr,
         period: String,
@@ -459,6 +540,7 @@ impl RbLazyFrame {
         right_on: RArray,
         allow_parallel: bool,
         force_parallel: bool,
+        join_nulls: bool,
         how: Wrap<JoinType>,
         suffix: String,
     ) -> RbResult<Self> {
@@ -474,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()
@@ -569,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,
@@ -619,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 {
@@ -668,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())
+    }
 }