polars-df 0.20.0 → 0.21.1

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

@@ -1,5 +1,5 @@
-use magnus::{prelude::*, r_hash::ForEach, RArray, RHash, Symbol, Value};
-use polars::frame::row::{rows_to_schema_supertypes, rows_to_supertypes, Row};
+use magnus::{RArray, RHash, Symbol, Value, prelude::*, r_hash::ForEach};
+use polars::frame::row::{Row, rows_to_schema_supertypes, rows_to_supertypes};
 use polars::prelude::*;
 
 use super::*;

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

@@ -1,9 +1,9 @@
-use magnus::{prelude::*, IntoValue, RArray, Value};
+use magnus::{IntoValue, RArray, Value, prelude::*};
 
 use super::*;
+use crate::RbResult;
 use crate::conversion::{ObjectValue, Wrap};
 use crate::interop::arrow::to_ruby::dataframe_to_stream;
-use crate::RbResult;
 
 impl RbDataFrame {
     pub fn row_tuple(&self, idx: i64) -> Value {

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

@@ -1,7 +1,7 @@
 use std::hash::BuildHasher;
 
 use either::Either;
-use magnus::{prelude::*, typed_data::Obj, IntoValue, RArray, Value};
+use magnus::{IntoValue, RArray, Value, prelude::*, typed_data::Obj};
 use polars::prelude::pivot::{pivot, pivot_stable};
 use polars::prelude::*;
 
@@ -416,17 +416,19 @@ impl RbDataFrame {
         columns: Option<Vec<String>>,
         separator: Option<String>,
         drop_first: bool,
+        drop_nulls: bool,
     ) -> RbResult<Self> {
         let df = match columns {
             Some(cols) => self.df.borrow().columns_to_dummies(
                 cols.iter().map(|x| x as &str).collect(),
                 separator.as_deref(),
                 drop_first,
+                drop_nulls,
             ),
             None => self
                 .df
                 .borrow()
-                .to_dummies(separator.as_deref(), drop_first),
+                .to_dummies(separator.as_deref(), drop_first, drop_nulls),
         }
         .map_err(RbPolarsErr::from)?;
         Ok(df.into())

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

@@ -1,6 +1,6 @@
-use magnus::{prelude::*, Value};
-use polars::io::avro::AvroCompression;
+use magnus::{Value, prelude::*};
 use polars::io::RowIndex;
+use polars::io::avro::AvroCompression;
 use polars::prelude::*;
 use std::io::BufWriter;
 use std::num::NonZeroUsize;

data/ext/polars/src/exceptions.rs

@@ -1,5 +1,5 @@
 use crate::rb_modules;
-use magnus::{exception, Error};
+use magnus::{Error, exception};
 use std::borrow::Cow;
 
 macro_rules! create_exception {
@@ -20,5 +20,6 @@ macro_rules! create_exception {
 create_exception!(RbTypeError, exception::type_error());
 create_exception!(RbValueError, exception::arg_error());
 create_exception!(RbOverflowError, exception::range_error());
+create_exception!(RbIndexError, exception::index_error());
 create_exception!(ComputeError, rb_modules::compute_error());
 create_exception!(InvalidOperationError, rb_modules::invalid_operation_error());

data/ext/polars/src/expr/array.rs

@@ -1,20 +1,41 @@
+use magnus::Value;
 use polars::prelude::*;
 
-use crate::RbExpr;
+use crate::{RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
-    pub fn array_max(&self) -> Self {
+    pub fn arr_len(&self) -> Self {
+        self.inner.clone().arr().len().into()
+    }
+
+    pub fn arr_max(&self) -> Self {
         self.inner.clone().arr().max().into()
     }
 
-    pub fn array_min(&self) -> Self {
+    pub fn arr_min(&self) -> Self {
         self.inner.clone().arr().min().into()
     }
 
-    pub fn array_sum(&self) -> Self {
+    pub fn arr_sum(&self) -> Self {
         self.inner.clone().arr().sum().into()
     }
 
+    pub fn arr_std(&self, ddof: u8) -> Self {
+        self.inner.clone().arr().std(ddof).into()
+    }
+
+    pub fn arr_var(&self, ddof: u8) -> Self {
+        self.inner.clone().arr().var(ddof).into()
+    }
+
+    pub fn arr_mean(&self) -> Self {
+        self.inner.clone().arr().mean().into()
+    }
+
+    pub fn arr_median(&self) -> Self {
+        self.inner.clone().arr().median().into()
+    }
+
     pub fn arr_unique(&self, maintain_order: bool) -> Self {
         if maintain_order {
             self.inner.clone().arr().unique_stable().into()
@@ -23,6 +44,10 @@ impl RbExpr {
         }
     }
 
+    pub fn arr_n_unique(&self) -> Self {
+        self.inner.clone().arr().n_unique().into()
+    }
+
     pub fn arr_to_list(&self) -> Self {
         self.inner.clone().arr().to_list().into()
     }
@@ -90,4 +115,48 @@ impl RbExpr {
             .count_matches(expr.inner.clone())
             .into()
     }
+
+    pub fn arr_to_struct(&self, name_gen: Option<Value>) -> Self {
+        if name_gen.is_some() {
+            todo!();
+        }
+        self.inner.clone().arr().to_struct(None).into()
+    }
+
+    pub fn arr_slice(
+        &self,
+        offset: &RbExpr,
+        length: Option<&RbExpr>,
+        as_array: bool,
+    ) -> RbResult<Self> {
+        let length = match length {
+            Some(i) => i.inner.clone(),
+            None => lit(i64::MAX),
+        };
+        Ok(self
+            .inner
+            .clone()
+            .arr()
+            .slice(offset.inner.clone(), length, as_array)
+            .map_err(RbPolarsErr::from)?
+            .into())
+    }
+
+    pub fn arr_tail(&self, n: &RbExpr, as_array: bool) -> RbResult<Self> {
+        Ok(self
+            .inner
+            .clone()
+            .arr()
+            .tail(n.inner.clone(), as_array)
+            .map_err(RbPolarsErr::from)?
+            .into())
+    }
+
+    pub fn arr_shift(&self, n: &RbExpr) -> Self {
+        self.inner.clone().arr().shift(n.inner.clone()).into()
+    }
+
+    pub fn arr_explode(&self) -> Self {
+        self.inner.clone().arr().explode().into()
+    }
 }

data/ext/polars/src/expr/binary.rs

@@ -1,4 +1,5 @@
-use crate::RbExpr;
+use crate::expr::datatype::RbDataTypeExpr;
+use crate::{RbExpr, RbResult};
 
 impl RbExpr {
     pub fn bin_contains(&self, lit: &RbExpr) -> Self {
@@ -40,4 +41,28 @@ impl RbExpr {
     pub fn bin_base64_encode(&self) -> Self {
         self.inner.clone().binary().base64_encode().into()
     }
+
+    pub fn bin_reinterpret(&self, dtype: &RbDataTypeExpr, kind: String) -> RbResult<Self> {
+        use crate::RbValueError;
+
+        let is_little_endian = match kind.to_lowercase().as_str() {
+            "little" => true,
+            "big" => false,
+            _ => {
+                return Err(RbValueError::new_err(format!(
+                    "Invalid endianness: {kind}. Valid values are \"little\" or \"big\"."
+                )));
+            }
+        };
+        Ok(self
+            .inner
+            .clone()
+            .binary()
+            .reinterpret(dtype.inner.clone(), is_little_endian)
+            .into())
+    }
+
+    pub fn bin_size_bytes(&self) -> Self {
+        self.inner.clone().binary().size_bytes().into()
+    }
 }

data/ext/polars/src/expr/bitwise.rs

@@ -0,0 +1,39 @@
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn bitwise_count_ones(&self) -> Self {
+        self.inner.clone().bitwise_count_ones().into()
+    }
+
+    pub fn bitwise_count_zeros(&self) -> Self {
+        self.inner.clone().bitwise_count_zeros().into()
+    }
+
+    pub fn bitwise_leading_ones(&self) -> Self {
+        self.inner.clone().bitwise_leading_ones().into()
+    }
+
+    pub fn bitwise_leading_zeros(&self) -> Self {
+        self.inner.clone().bitwise_leading_zeros().into()
+    }
+
+    pub fn bitwise_trailing_ones(&self) -> Self {
+        self.inner.clone().bitwise_trailing_ones().into()
+    }
+
+    pub fn bitwise_trailing_zeros(&self) -> Self {
+        self.inner.clone().bitwise_trailing_zeros().into()
+    }
+
+    pub fn bitwise_and(&self) -> Self {
+        self.inner.clone().bitwise_and().into()
+    }
+
+    pub fn bitwise_or(&self) -> Self {
+        self.inner.clone().bitwise_or().into()
+    }
+
+    pub fn bitwise_xor(&self) -> Self {
+        self.inner.clone().bitwise_xor().into()
+    }
+}

data/ext/polars/src/expr/categorical.rs

@@ -4,4 +4,24 @@ impl RbExpr {
     pub fn cat_get_categories(&self) -> Self {
         self.inner.clone().cat().get_categories().into()
     }
+
+    pub fn cat_len_bytes(&self) -> Self {
+        self.inner.clone().cat().len_bytes().into()
+    }
+
+    pub fn cat_len_chars(&self) -> Self {
+        self.inner.clone().cat().len_chars().into()
+    }
+
+    pub fn cat_starts_with(&self, prefix: String) -> Self {
+        self.inner.clone().cat().starts_with(prefix).into()
+    }
+
+    pub fn cat_ends_with(&self, suffix: String) -> Self {
+        self.inner.clone().cat().ends_with(suffix).into()
+    }
+
+    pub fn cat_slice(&self, offset: i64, length: Option<usize>) -> Self {
+        self.inner.clone().cat().slice(offset, length).into()
+    }
 }

data/ext/polars/src/expr/datatype.rs

@@ -0,0 +1,37 @@
+use magnus::{IntoValue, Value};
+use polars::prelude::{DataType, DataTypeExpr, Schema};
+
+use crate::prelude::Wrap;
+use crate::{RbExpr, RbPolarsErr, RbResult};
+
+#[magnus::wrap(class = "Polars::RbDataTypeExpr")]
+#[repr(transparent)]
+#[derive(Clone)]
+pub struct RbDataTypeExpr {
+    pub inner: DataTypeExpr,
+}
+
+impl From<DataTypeExpr> for RbDataTypeExpr {
+    fn from(expr: DataTypeExpr) -> Self {
+        RbDataTypeExpr { inner: expr }
+    }
+}
+
+impl RbDataTypeExpr {
+    pub fn from_dtype(datatype: Wrap<DataType>) -> Self {
+        DataTypeExpr::Literal(datatype.0).into()
+    }
+
+    pub fn of_expr(expr: &RbExpr) -> Self {
+        DataTypeExpr::OfExpr(Box::new(expr.inner.clone())).into()
+    }
+
+    pub fn collect_dtype(&self, schema: Wrap<Schema>) -> RbResult<Value> {
+        let dtype = self
+            .clone()
+            .inner
+            .into_datatype(&schema.0)
+            .map_err(RbPolarsErr::from)?;
+        Ok(Wrap(dtype).into_value())
+    }
+}

data/ext/polars/src/expr/datetime.rs

@@ -4,6 +4,20 @@ use crate::conversion::Wrap;
 use crate::{RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
+    pub fn dt_add_business_days(
+        &self,
+        n: &RbExpr,
+        week_mask: [bool; 7],
+        holidays: Vec<i32>,
+        roll: Wrap<Roll>,
+    ) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .add_business_days(n.inner.clone(), week_mask, holidays, roll.0)
+            .into()
+    }
+
     pub fn dt_to_string(&self, format: String) -> Self {
         self.inner.clone().dt().to_string(&format).into()
     }
@@ -90,6 +104,34 @@ impl RbExpr {
         self.inner.clone().dt().round(every.inner.clone()).into()
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub fn dt_replace(
+        &self,
+        year: &Self,
+        month: &Self,
+        day: &Self,
+        hour: &Self,
+        minute: &Self,
+        second: &Self,
+        microsecond: &Self,
+        ambiguous: &Self,
+    ) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .replace(
+                year.inner.clone(),
+                month.inner.clone(),
+                day.inner.clone(),
+                hour.inner.clone(),
+                minute.inner.clone(),
+                second.inner.clone(),
+                microsecond.inner.clone(),
+                ambiguous.inner.clone(),
+            )
+            .into()
+    }
+
     pub fn dt_combine(&self, time: &Self, time_unit: Wrap<TimeUnit>) -> Self {
         self.inner
             .clone()
@@ -98,10 +140,26 @@ impl RbExpr {
             .into()
     }
 
+    pub fn dt_millennium(&self) -> Self {
+        self.inner.clone().dt().millennium().into()
+    }
+
+    pub fn dt_century(&self) -> Self {
+        self.inner.clone().dt().century().into()
+    }
+
     pub fn dt_year(&self) -> Self {
         self.clone().inner.dt().year().into()
     }
 
+    pub fn dt_is_business_day(&self, week_mask: [bool; 7], holidays: Vec<i32>) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .is_business_day(week_mask, holidays)
+            .into()
+    }
+
     pub fn dt_is_leap_year(&self) -> Self {
         self.clone().inner.dt().is_leap_year().into()
     }

data/ext/polars/src/expr/general.rs

@@ -6,10 +6,11 @@ use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 use polars_core::series::IsSorted;
 
-use crate::conversion::{parse_fill_null_strategy, Wrap};
+use super::selector::RbSelector;
+use crate::conversion::{Wrap, parse_fill_null_strategy};
 use crate::map::lazy::map_single;
 use crate::rb_exprs_to_exprs;
-use crate::{RbExpr, RbResult};
+use crate::{RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
     pub fn add(&self, rhs: &Self) -> RbResult<Self> {
@@ -276,17 +277,8 @@ impl RbExpr {
             .into()
     }
 
-    pub fn arg_sort(&self, reverse: bool, nulls_last: bool) -> Self {
-        self.clone()
-            .inner
-            .arg_sort(SortOptions {
-                descending: reverse,
-                nulls_last,
-                multithreaded: true,
-                maintain_order: false,
-                limit: None,
-            })
-            .into()
+    pub fn arg_sort(&self, descending: bool, nulls_last: bool) -> Self {
+        self.inner.clone().arg_sort(descending, nulls_last).into()
     }
 
     pub fn top_k(&self, k: &Self) -> Self {
@@ -331,6 +323,10 @@ impl RbExpr {
         self.inner.clone().arg_min().into()
     }
 
+    pub fn index_of(&self, element: &Self) -> Self {
+        self.inner.clone().index_of(element.inner.clone()).into()
+    }
+
     pub fn search_sorted(
         &self,
         element: &Self,
@@ -437,6 +433,13 @@ impl RbExpr {
             .into()
     }
 
+    pub fn is_close(&self, other: &Self, abs_tol: f64, rel_tol: f64, nans_equal: bool) -> Self {
+        self.inner
+            .clone()
+            .is_close(other.inner.clone(), abs_tol, rel_tol, nans_equal)
+            .into()
+    }
+
     pub fn approx_n_unique(&self) -> Self {
         self.inner.clone().approx_n_unique().into()
     }
@@ -490,6 +493,10 @@ impl RbExpr {
         self.inner.clone().round(decimals, mode.0).into()
     }
 
+    pub fn round_sig_figs(&self, digits: i32) -> Self {
+        self.clone().inner.round_sig_figs(digits).into()
+    }
+
     pub fn floor(&self) -> Self {
         self.inner.clone().floor().into()
     }
@@ -525,6 +532,10 @@ impl RbExpr {
         self.inner.clone().tan().into()
     }
 
+    pub fn cot(&self) -> Self {
+        self.inner.clone().cot().into()
+    }
+
     pub fn arcsin(&self) -> Self {
         self.inner.clone().arcsin().into()
     }
@@ -561,6 +572,14 @@ impl RbExpr {
         self.inner.clone().arctanh().into()
     }
 
+    pub fn degrees(&self) -> Self {
+        self.inner.clone().degrees().into()
+    }
+
+    pub fn radians(&self) -> Self {
+        self.inner.clone().radians().into()
+    }
+
     pub fn sign(&self) -> Self {
         self.inner.clone().sign().into()
     }
@@ -574,18 +593,35 @@ impl RbExpr {
         Ok(self.inner.clone().over(partition_by).into())
     }
 
-    pub fn _and(&self, expr: &Self) -> Self {
-        self.inner.clone().and(expr.inner.clone()).into()
+    pub fn rolling(
+        &self,
+        index_column: String,
+        period: String,
+        offset: String,
+        closed: Wrap<ClosedWindow>,
+    ) -> RbResult<Self> {
+        let options = RollingGroupOptions {
+            index_column: index_column.into(),
+            period: Duration::try_parse(&period).map_err(RbPolarsErr::from)?,
+            offset: Duration::try_parse(&offset).map_err(RbPolarsErr::from)?,
+            closed_window: closed.0,
+        };
+
+        Ok(self.inner.clone().rolling(options).into())
     }
 
-    pub fn _xor(&self, expr: &Self) -> Self {
-        self.inner.clone().xor(expr.inner.clone()).into()
+    pub fn and_(&self, expr: &Self) -> Self {
+        self.inner.clone().and(expr.inner.clone()).into()
     }
 
-    pub fn _or(&self, expr: &Self) -> Self {
+    pub fn or_(&self, expr: &Self) -> Self {
         self.inner.clone().or(expr.inner.clone()).into()
     }
 
+    pub fn xor_(&self, expr: &Self) -> Self {
+        self.inner.clone().xor(expr.inner.clone()).into()
+    }
+
     pub fn is_in(&self, expr: &Self, nulls_equal: bool) -> Self {
         self.inner
             .clone()
@@ -601,6 +637,14 @@ impl RbExpr {
         self.inner.clone().pow(exponent.inner.clone()).into()
     }
 
+    pub fn sqrt(&self) -> Self {
+        self.inner.clone().sqrt().into()
+    }
+
+    pub fn cbrt(&self) -> Self {
+        self.inner.clone().cbrt().into()
+    }
+
     pub fn cum_sum(&self, reverse: bool) -> Self {
         self.inner.clone().cum_sum(reverse).into()
     }
@@ -655,10 +699,6 @@ impl RbExpr {
         self.inner.clone().mode().into()
     }
 
-    pub fn exclude(&self, columns: Vec<String>) -> Self {
-        self.inner.clone().exclude(columns).into()
-    }
-
     pub fn interpolate(&self, method: Wrap<InterpolationMethod>) -> Self {
         self.inner.clone().interpolate(method.0).into()
     }
@@ -768,6 +808,15 @@ impl RbExpr {
         self.inner.clone().ewm_mean(options).into()
     }
 
+    pub fn ewm_mean_by(&self, times: &RbExpr, half_life: String) -> RbResult<Self> {
+        let half_life = Duration::try_parse(&half_life).map_err(RbPolarsErr::from)?;
+        Ok(self
+            .inner
+            .clone()
+            .ewm_mean_by(times.inner.clone(), half_life)
+            .into())
+    }
+
     pub fn ewm_std(
         &self,
         alpha: f64,
@@ -823,6 +872,10 @@ impl RbExpr {
         self.inner.clone().log(base).into()
     }
 
+    pub fn log1p(&self) -> Self {
+        self.inner.clone().log1p().into()
+    }
+
     pub fn exp(&self) -> Self {
         self.inner.clone().exp().into()
     }
@@ -868,4 +921,35 @@ impl RbExpr {
             )
             .into()
     }
+
+    pub fn hist(
+        &self,
+        bins: Option<&RbExpr>,
+        bin_count: Option<usize>,
+        include_category: bool,
+        include_breakpoint: bool,
+    ) -> Self {
+        let bins = bins.map(|e| e.inner.clone());
+        self.inner
+            .clone()
+            .hist(bins, bin_count, include_category, include_breakpoint)
+            .into()
+    }
+
+    #[allow(clippy::wrong_self_convention)]
+    pub fn into_selector(&self) -> RbResult<RbSelector> {
+        Ok(self
+            .inner
+            .clone()
+            .into_selector()
+            .ok_or_else(
+                || polars_err!(InvalidOperation: "expr `{}` is not a selector", &self.inner),
+            )
+            .map_err(RbPolarsErr::from)?
+            .into())
+    }
+
+    pub fn new_selector(selector: &RbSelector) -> Self {
+        Expr::Selector(selector.inner.clone()).into()
+    }
 }