polars-df 0.21.0 → 0.22.0

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

@@ -4,10 +4,14 @@ mod general;
 mod io;
 mod serde;
 
+use magnus::{DataTypeFunctions, TypedData, gc};
 use polars::prelude::*;
 use std::cell::RefCell;
 
-#[magnus::wrap(class = "Polars::RbDataFrame")]
+use crate::series::mark_series;
+
+#[derive(TypedData)]
+#[magnus(class = "Polars::RbDataFrame", mark)]
 pub struct RbDataFrame {
     pub df: RefCell<DataFrame>,
 }
@@ -25,3 +29,23 @@ impl RbDataFrame {
         }
     }
 }
+
+impl DataTypeFunctions for RbDataFrame {
+    fn mark(&self, marker: &gc::Marker) {
+        // this is really, really not ideal, as objects will not be marked if unable to borrow
+        // currently, this should only happen for write_* methods,
+        // which should refuse to write Object datatype, and therefore be safe,
+        // since GC will not have a chance to run
+        if let Ok(df) = self.df.try_borrow() {
+            for column in df.get_columns() {
+                if let DataType::Object(_) = column.dtype() {
+                    match column {
+                        Column::Series(s) => mark_series(marker, s),
+                        Column::Partitioned(s) => mark_series(marker, s.partitions()),
+                        Column::Scalar(s) => mark_series(marker, &s.as_single_value_series()),
+                    }
+                }
+            }
+        }
+    }
+}

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

@@ -1,15 +1,30 @@
-use crate::exceptions::ComputeError;
+use std::io::{BufReader, BufWriter};
+
+use polars::prelude::*;
+
 use crate::file::get_file_like;
-use crate::{RbDataFrame, RbResult};
+use crate::utils::to_rb_err;
+use crate::{RbDataFrame, RbPolarsErr, RbResult};
 use magnus::Value;
-use std::io::BufWriter;
 
 impl RbDataFrame {
-    // TODO add to Ruby
-    pub fn serialize_json(&self, rb_f: Value) -> RbResult<()> {
+    pub fn serialize_binary(&self, rb_f: Value) -> RbResult<()> {
         let file = get_file_like(rb_f, true)?;
-        let writer = BufWriter::new(file);
-        serde_json::to_writer(writer, &self.df)
-            .map_err(|err| ComputeError::new_err(err.to_string()))
+        let mut writer = BufWriter::new(file);
+
+        Ok(self
+            .df
+            .borrow_mut()
+            .serialize_into_writer(&mut writer)
+            .map_err(RbPolarsErr::from)?)
+    }
+
+    pub fn deserialize_binary(rb_f: Value) -> RbResult<Self> {
+        let file = get_file_like(rb_f, false)?;
+        let mut file = BufReader::new(file);
+
+        DataFrame::deserialize_from_reader(&mut file)
+            .map(|v| v.into())
+            .map_err(to_rb_err)
     }
 }

data/ext/polars/src/exceptions.rs

@@ -1,5 +1,5 @@
 use crate::rb_modules;
-use magnus::{Error, exception};
+use magnus::{Error, Ruby};
 use std::borrow::Cow;
 
 macro_rules! create_exception {
@@ -17,8 +17,12 @@ macro_rules! create_exception {
     };
 }
 
-create_exception!(RbTypeError, exception::type_error());
-create_exception!(RbValueError, exception::arg_error());
-create_exception!(RbOverflowError, exception::range_error());
+create_exception!(RbTypeError, Ruby::get().unwrap().exception_type_error());
+create_exception!(RbValueError, Ruby::get().unwrap().exception_arg_error());
+create_exception!(
+    RbOverflowError,
+    Ruby::get().unwrap().exception_range_error()
+);
+create_exception!(RbIndexError, Ruby::get().unwrap().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

@@ -1,4 +1,8 @@
-use polars::prelude::DataTypeExpr;
+use magnus::{IntoValue, Ruby, Value};
+use polars::prelude::{DataType, DataTypeExpr, Schema};
+
+use crate::prelude::Wrap;
+use crate::{RbExpr, RbPolarsErr, RbResult};
 
 #[magnus::wrap(class = "Polars::RbDataTypeExpr")]
 #[repr(transparent)]
@@ -12,3 +16,22 @@ impl From<DataTypeExpr> for RbDataTypeExpr {
         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(ruby: &Ruby, rb_self: &Self, schema: Wrap<Schema>) -> RbResult<Value> {
+        let dtype = rb_self
+            .clone()
+            .inner
+            .into_datatype(&schema.0)
+            .map_err(RbPolarsErr::from)?;
+        Ok(Wrap(dtype).into_value_with(ruby))
+    }
+}

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()
     }
@@ -12,20 +26,6 @@ impl RbExpr {
         self.inner.clone().dt().offset_by(by.inner.clone()).into()
     }
 
-    pub fn dt_epoch_seconds(&self) -> Self {
-        self.inner
-            .clone()
-            .map(
-                |s| {
-                    s.take_materialized_series()
-                        .timestamp(TimeUnit::Milliseconds)
-                        .map(|ca| Some((ca / 1000).into_column()))
-                },
-                GetOutput::from_type(DataType::Int64),
-            )
-            .into()
-    }
-
     pub fn dt_with_time_unit(&self, tu: Wrap<TimeUnit>) -> Self {
         self.inner.clone().dt().with_time_unit(tu.0).into()
     }
@@ -90,6 +90,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 +126,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

@@ -323,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,
@@ -429,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()
     }
@@ -472,16 +483,17 @@ impl RbExpr {
     }
 
     pub fn rechunk(&self) -> Self {
-        self.inner
-            .clone()
-            .map(|s| Ok(Some(s.rechunk())), GetOutput::same_type())
-            .into()
+        self.inner.clone().rechunk().into()
     }
 
     pub fn round(&self, decimals: u32, mode: Wrap<RoundMode>) -> Self {
         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()
     }
@@ -517,6 +529,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()
     }
@@ -553,6 +569,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()
     }
@@ -566,18 +590,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()
@@ -593,6 +634,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()
     }
@@ -613,10 +662,6 @@ impl RbExpr {
         self.inner.clone().product().into()
     }
 
-    pub fn shrink_dtype(&self) -> Self {
-        self.inner.clone().shrink_dtype().into()
-    }
-
     pub fn map_batches(
         &self,
         lambda: Value,
@@ -756,6 +801,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,
@@ -807,8 +861,12 @@ impl RbExpr {
         self.inner.clone().all(drop_nulls).into()
     }
 
-    pub fn log(&self, base: f64) -> Self {
-        self.inner.clone().log(base).into()
+    pub fn log(&self, base: &RbExpr) -> Self {
+        self.inner.clone().log(base.inner.clone()).into()
+    }
+
+    pub fn log1p(&self) -> Self {
+        self.inner.clone().log1p().into()
     }
 
     pub fn exp(&self) -> Self {
@@ -857,6 +915,20 @@ 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